Protein Language Models: Extract ESM-2 Embeddings with Python

A protein sequence is a string of amino acid letters. Most machine-learning tools, however, need numbers. A protein language model provides one way to turn that sequence into a numerical representation that can be reused for search, clustering, or a supervised prediction task.

This tutorial uses the small ESM-2 checkpoint to extract one 320-dimensional vector per sequence on a CPU. You will save the vectors as JSON, check their shape, and understand why an embedding similarity score does not establish biological function. Basic Python knowledge is enough; no model training or laboratory work is involved.

1. What does a protein language model learn?

In a masked protein language model, some amino acid tokens are hidden during training. The model learns to predict them from the surrounding sequence. The resulting hidden states represent each residue in its sequence context. ESM uses this masked-language-modeling approach; see the official Transformers ESM documentation.

An embedding is a list of learned numerical features. Individual coordinates usually have no simple interpretation such as temperature, charge, or enzyme activity. A later task can use the features, but the embedding alone is not a measured property.

Keep these three outputs separate:

ESM-2 supplies the representations used here. ESMFold adds a structure-prediction architecture. ESM3 is a separate multimodal model involving sequence, structure, and function. The ESM repository and ESM3 paper describe those distinctions. Loading ESM-2 in this tutorial does not run ESMFold or ESM3.

2. Install the small ESM-2 example

Use a recent Python version supported by your PyTorch installation. Create an isolated environment:

mkdir protein-lm-demo
cd protein-lm-demo
python3 -m venv .venv
. .venv/bin/activate
python -m pip install torch 'transformers==4.57.6'

The checkpoint is facebook/esm2_t6_8M_UR50D: six layers, roughly eight million parameters, and 320 hidden dimensions. Those dimensions are also recorded in its configuration. It is a manageable starting point for short sequences; it is not presented here as the best model for every task.

The first run downloads public model files and requires internet access. Later runs can reuse the Hugging Face cache. The code pins the model revision and runs on CPU, so a GPU and an API key are unnecessary for this example. Package downloads and model files still require disk space.

3. Convert sequences into embeddings

Save this as protein_embeddings.py:

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, EsmModel


MODEL_ID = "facebook/esm2_t6_8M_UR50D"
REVISION = "c731040fcd8d73dceaa04b0a8e6329b345b0f5df"
AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWY")
# Artificial strings for a software exercise; no biological function is claimed.
SEQUENCES = ["ACDEFGHIKLMNPQRSTVWY", "MKTAYIAKQRQISFVKSHFSRQ"]


def embed(sequences, tokenizer, model):
    for sequence in sequences:
        if not 1 <= len(sequence) <= 1022:
            raise ValueError("Use sequences of 1-1022 residues; no silent truncation")
        if set(sequence) - AMINO_ACIDS:
            raise ValueError("This demo accepts only the 20 standard uppercase amino acids")
    if not sequences:
        raise ValueError("Provide at least one sequence")

    batch = tokenizer(
        sequences, padding=True, truncation=False,
        return_special_tokens_mask=True, return_tensors="pt",
    )
    special = batch.pop("special_tokens_mask").bool()
    residues = batch["attention_mask"].bool() & ~special
    lengths = residues.sum(dim=1)
    expected = torch.tensor([len(s) for s in sequences])
    if not torch.equal(lengths, expected):
        raise ValueError("Token counts do not match residue counts")

    with torch.inference_mode():
        hidden = model(**batch).last_hidden_state
        weights = residues.unsqueeze(-1).to(hidden.dtype)
        pooled = (hidden * weights).sum(dim=1) / lengths.unsqueeze(-1)
    if not torch.isfinite(pooled).all():
        raise ValueError("Non-finite embedding values")
    return pooled, hidden.shape, lengths.tolist()


def main():
    parser = argparse.ArgumentParser(description="Extract ESM-2 protein embeddings on CPU")
    parser.add_argument("--output", type=Path, default=Path("embeddings.json"))
    args = parser.parse_args()
    if args.output.exists():
        parser.error("Output exists; choose another --output path")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=REVISION)
    model = EsmModel.from_pretrained(
        MODEL_ID, revision=REVISION, add_pooling_layer=False,
    )
    model.eval()
    vectors, token_shape, lengths = embed(SEQUENCES, tokenizer, model)
    similarity = F.cosine_similarity(vectors[0:1], vectors[1:2]).item()
    result = {
        "model": MODEL_ID,
        "revision": REVISION,
        "pooling": "last-layer mean over residue tokens only",
        "sequences": SEQUENCES,
        "embeddings": vectors.tolist(),
    }
    with args.output.open("x", encoding="utf-8") as output:
        json.dump(result, output, indent=2, allow_nan=False)
        output.write("\n")
    print("Residue counts:", lengths)
    print("Token tensor:", tuple(token_shape))
    print("Protein tensor:", tuple(vectors.shape))
    print(f"Cosine similarity (not a probability): {similarity:.6f}")
    print(f"Saved {args.output}")


if __name__ == "__main__":
    main()

The two input strings are artificial software examples. They are not a benchmark or a claim about real protein activity. Replace SEQUENCES with your own amino acid strings when the example works.

The script deliberately accepts only the 20 standard uppercase amino acid letters. Some real datasets contain ambiguous or nonstandard residues; decide how to handle those before extending the input policy. Do not paste FASTA headers, nucleotide sequences, alignment gaps, or stop symbols into this list. A DNA string containing only A, C, G, and T would pass the character check, so you must also verify the source data type.

The 1,022-residue limit is a conservative limit for this demonstration, not a universal limit for protein models. Longer sequences are rejected instead of silently shortened. For large datasets, process short batches of similar lengths; the longest sequence determines how much padding a batch needs.

Why the pooling mask matters

The tokenizer adds special tokens and pads shorter sequences. attention_mask identifies real input positions, but it also includes the special tokens. The code combines it with special_tokens_mask so only amino acid positions contribute to the mean.

The last hidden state has shape (batch, token_positions, hidden_dimensions). After averaging residue positions, the result has shape (batch, hidden_dimensions). add_pooling_layer=False avoids an unused pooler; the script performs its own explicit mean pooling.

Mean pooling is a baseline. It compresses a whole sequence into one vector and can hide local differences. Keep residue-level representations if your task needs a prediction at each position. The ESM API reference documents the hidden-state output used here.

4. Run and verify the result

python protein_embeddings.py --help
python protein_embeddings.py --output embeddings.json
python -m json.tool embeddings.json > /dev/null
python -m pip freeze > requirements-tested.txt

The shape-related output should be:

Residue counts: [20, 22]
Token tensor: (2, 24, 320)
Protein tensor: (2, 320)

There are 24 token positions because the longer string has 22 residues plus two special tokens. The shorter string is padded to match. The script also prints a cosine similarity and saves the vectors, input sequences, checkpoint revision, and pooling method.

Validate the saved file independently:

python - <<'PYTHON'
import json
import math

with open("embeddings.json", encoding="utf-8") as source:
    result = json.load(source)
vectors = result["embeddings"]
assert len(vectors) == len(result["sequences"]) == 2
assert all(len(vector) == 320 for vector in vectors)
assert all(math.isfinite(x) for vector in vectors for x in vector)
print("Validated: 2 protein embeddings, 320 finite values each")
PYTHON

The example was run on CPU with Python 3.14.7, PyTorch 2.14.0, Transformers 4.57.6, and the pinned checkpoint above. The shapes matched, and the two artificial inputs produced a cosine similarity of 0.937591. That high number is a useful reminder that this score is not evidence of shared function; these inputs have no functional labels. Small numerical differences across environments are possible.

An existing output file is preserved: choose another path with --output for a new run. Keep requirements-tested.txt with the JSON if you want to record the installed package versions.

5. Use the vectors without overreading them

Cosine similarity compares vector directions. It is not a probability, an alignment identity percentage, or proof that two proteins share a function. There is no universal threshold at which this tutorial’s mean-pooled embeddings establish a biological relationship.

A 2026 Nature Methods study on uncertainty in protein representations examines how embeddings vary across models and sequence types. Its practical relevance here is simple: evaluate representations on the data you actually plan to use.

For a first classification experiment, I would freeze ESM-2, extract vectors for a labeled dataset, and fit a simple classifier. Compare it with a basic baseline, such as amino acid composition, before spending time on fine-tuning.

If the intended use is prediction for unfamiliar protein families, a random train/test split can be misleading when related sequences occur on both sides. Use a similarity-aware split aligned with that intended use, and fit any scaling or feature selection only on the training portion. The DataSAIL paper explains how similarity across splits can inflate evaluation results. The two strings in this article only check the software path; they cannot measure model quality.

Troubleshooting

A useful next step

You now have a repeatable path from a protein sequence to a saved vector, with explicit token handling and basic output checks. Try a small labeled dataset next and assess a frozen-embedding baseline on a split that matches your intended use. Treat successful extraction and useful biological prediction as separate milestones.