蛋白质语言模型入门：用 Python 提取 ESM-2 向量

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

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()

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

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

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