fair-esm2

fair-esm2 — ESM-2 (Meta AI)

ESM-2 code and weights are MIT (Meta AI, github.com/facebookresearch/esm).

Package disambiguation. pip install fair-esm gives you import esm with esm.pretrained.* (ESM-1/2). Biohub's github.com/Biohub/esm fork (MIT) gives you from esm.models.esmfold2 import ESMFold2InputBuilder — see the esmfold2 skill. Both share the esm namespace but are different libraries. This skill covers fair-esm (the Meta package).

Prerequisites

How to run

Embeddings

import torch, esm

model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
model = model.eval().cuda()
bc = alphabet.get_batch_converter()

_, _, toks = bc([("ubq", "MQIFVKTLTGKTITLEVEPSDTIENVK")])
with torch.no_grad():
    out = model(toks.cuda(), repr_layers=[33])
emb = out["representations"][33]      # (1, L+2, 1280) — includes BOS/EOS
seq_emb = emb[0, 1:-1].mean(0)        # per-sequence mean

Masked-LM scoring

with torch.no_grad():
    out = model(toks.cuda(), repr_layers=[33])
logits = out["logits"][0, 1:-1]       # (L, |vocab|)
# WT marginal log-likelihood; for mutation scoring, mask the position and
# compare logit[mut] − logit[wt].

Contact prediction

with torch.no_grad():
    out = model(toks.cuda(), repr_layers=[33], return_contacts=True)
contacts = out["contacts"][0]         # (L, L)

Models

Output format

out["representations"][layer] is (B, L+2, D); slice [ :, 1:-1, : ] to drop BOS/EOS. out["contacts"] (when return_contacts=True) is (B, L, L).

Remote compute

Needs ≥16 GB VRAM (650M model) and either pre-cached .pt checkpoints or egress to dl.fbaipublicfiles.com. Use a selected and probed ssh:<alias> context and load remote-compute-ssh. Confirm the fair-esm environment and torch-hub cache, then submit a self-contained runner with run_in_context:

{
  "context_id": "ssh:gpu-box",
  "title": "ESM-2 embeddings for 200 sequences",
  "command": "source ~/miniforge3/etc/profile.d/conda.sh && conda activate fair-esm && TORCH_HOME=/srv/torch-cache python embed_esm2.py --input seqs.fasta --output /home/me/wisp-results/esm2/embeddings.pt",
  "timeout_secs": 1800,
  "input_paths": ["runs/embed_esm2.py", "data/seqs.fasta"],
  "output_specs": [
    {
      "glob": "ssh://gpu-box/home/me/wisp-results/esm2/embeddings.pt",
      "kind": "pytorch",
      "residency": "remote"
    }
  ]
}

Replace context, environment, cache, and output paths with discovered values. For a large input already on the server, use its absolute path instead of staging it. Call monitor_run once to wait, get_run once for a snapshot, or cancel_run to stop.

Troubleshooting

Next: feed embeddings to a classifier. For structure prediction, use esmfold2.