README.md

---
library_name: transformers
license: mit
tags:
  - biology
  - esm
  - protein
  - protein-language-model
  - masked-language-modeling
---

# ESM++ 6B

[ESM++](https://github.com/Synthyra/FastPLMs) is a Hugging Face compatible implementation of [Biohub ESMC](https://biohub.ai/esm/protein) ([license](https://github.com/Biohub/esm/blob/main/LICENSE.md)).
This checkpoint corresponds to the 6 billion parameter ESMC model released as [`biohub/ESMC-6B`](https://huggingface.co/biohub/ESMC-6B).

This repository includes the Biohub ESM MIT license in `LICENSE`.

The 6B model has 80 transformer layers, hidden size 2560, and 40 attention heads. It is large enough that `dtype=torch.bfloat16` or `torch.float16` plus `device_map="auto"` is usually the practical loading path.

## Attention Backends

`sdpa` is the default backend. Set `config.attn_backend` before loading if you want a different attention implementation.

| Backend | Key | Notes |
| :--- | :--- | :--- |
| PyTorch SDPA | `"sdpa"` | Default. Exact numerics and stable on all hardware. |
| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs when `kernels` is installed. Outputs are not bitwise identical to SDPA. |
| Flex Attention | `"flex"` | Skips padding tokens via block masks. First use compiles a Triton kernel. |
| Auto | `"auto"` | Picks the best available backend: `kernels_flash`, then `flex`, then `sdpa`. |

```python
import torch
from transformers import AutoConfig, AutoModelForMaskedLM

config = AutoConfig.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    trust_remote_code=True,
)
config.attn_backend = "auto"

model = AutoModelForMaskedLM.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    config=config,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)
```

## Masked Language Modeling

```python
import torch
from transformers import AutoModelForMaskedLM

model = AutoModelForMaskedLM.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = model.tokenizer

sequences = ["MPRTEIN", "MSEQWENCE"]
inputs = tokenizer(sequences, padding=True, return_tensors="pt")
inputs = inputs.to(model.device)

with torch.no_grad():
    output = model(**inputs)

print(output.logits.shape)
print(output.last_hidden_state.shape)
```

Pass `output_hidden_states=True` if you need all intermediate hidden states.

## Experimental Test-Time Training

TTT is disabled by default. Normal ESM++ inference, embeddings, logits, and
`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`.
The current implementation is experimental and trains only local LoRA adapters
on the ESMC backbone with masked language modeling on the test protein. It can
help some difficult proteins, but it adds test-time compute and can degrade
already confident predictions. The 6B checkpoint is large, so start with small
`steps`, `ags`, and `batch_size` values.

```python
metrics = model.ttt(
    seq="MSTNPKPQRKTKRNT",
    ttt_config={"steps": 1, "ags": 1, "batch_size": 1},
)
model.ttt_reset()
print(metrics["losses"])
```

## Binder Design Regularizer

The FastPLMs binder design tutorial uses `Synthyra/ESMplusplus_6B` as the
ESMC-style masked-LM regularizer while FastPLMs ESMFold2 experimental models
provide differentiable folding losses and final critics. The script lives at
`cookbook/tutorials/binder_design_fastplms.py` and supports local CUDA Docker
runs plus Modal deployment.

Run the verified EGFR 128 amino acid de novo minibinder example:

```bash
cd /home/ubuntu/FastPLMs

sudo -n docker run --gpus all --rm \
  -v /home/ubuntu/FastPLMs:/app \
  -v /home/ubuntu/FastPLMs:/workspace \
  -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \
  -w /workspace fastplms-esmfold2 \
  python /app/cookbook/tutorials/binder_design_fastplms.py \
    --backend local \
    --target-name egfr \
    --binder-sequence '################################################################################################################################' \
    --not-antibody \
    --steps 150 \
    --batch-size 1 \
    --seed 103 \
    --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli
```

The run writes `trajectory.jsonl`, `best_sequences.fasta`, `results.parquet`,
`selection.parquet`, and per-critic PDB/CIF/logit files. The verified candidate
had hero mean iPTM `0.913870`, hero min iPTM `0.904600`, and all four ESMFold2
hero critics above `0.9`.

Binder sequence:

```text
SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ
```

See [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
for the full strategy, Modal backend, official pI and selection scoring,
per-critic metrics, and caveats.

## Embed Datasets

All FastPLMs sequence models include `embed_dataset`, which handles batching, length sorting, pooling, FASTA parsing, optional resume from existing outputs, and `.pth` or SQLite storage.

```python
import torch
from transformers import AutoModelForMaskedLM

model = AutoModelForMaskedLM.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)

embedding_dict = model.embed_dataset(
    sequences=[
        "MALWMRLLPLLALLALWGPDPAAA",
        "MSEQWENCE",
        "MPRTEIN",
    ],
    batch_size=1,
    max_len=1024,
    full_embeddings=False,
    embed_dtype=torch.float32,
    pooling_types=["mean", "cls"],
    num_workers=0,
    save=True,
    save_path="esmplusplus_6b_embeddings.pth",
)

print(embedding_dict["MPRTEIN"].shape)
```

For residue-level embeddings, set `full_embeddings=True`:

```python
residue_embeddings = model.embed_dataset(
    sequences=["MALWMRLLPLLALLALWGPDPAAA"],
    batch_size=1,
    max_len=1024,
    full_embeddings=True,
    embed_dtype=torch.float32,
    save=False,
)
```

For very large datasets, write embeddings directly to SQLite:

```python
model.embed_dataset(
    fasta_path="proteins.fasta",
    batch_size=1,
    max_len=1024,
    pooling_types=["mean"],
    sql=True,
    sql_db_path="esmplusplus_6b_embeddings.db",
    save=False,
)
```

`embed_dataset` returns a dictionary when `sql=False`. With `sql=True`, embeddings are written to the database and loaded as needed.

## Classification Heads

ESM++ supports sequence-level and token-level classification through the standard Transformers auto classes.

```python
import torch
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    num_labels=2,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)

tokenized = model.tokenizer(
    ["MPRTEIN", "MSEQWENCE"],
    padding=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    logits = model(**tokenized).logits

print(logits.shape)
```

## LoRA Fine-Tuning

```python
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "Synthyra/ESMplusplus_6B",
    num_labels=2,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.01,
    bias="none",
    target_modules=[
        "layernorm_qkv.1",
        "out_proj",
        "query",
        "key",
        "value",
        "dense",
    ],
)

model = get_peft_model(model, lora_config)
```

## Attention Maps

Optimized attention backends do not return attention maps directly. ESM++ can compute them manually with `output_attentions=True`, but this is much slower and memory-heavy for the 6B model.

```python
with torch.no_grad():
    output = model(**inputs, output_attentions=True)

attentions = output.attentions
print(len(attentions))
print(attentions[0].shape)
```

## Load Biohub Source Weights

You can also load the Biohub source weights directly through FastPLMs:

```python
from fastplms.esm_plusplus.modeling_esm_plusplus import ESMplusplusForMaskedLM

model = ESMplusplusForMaskedLM.from_pretrained_esm("esmc-6b")
```

The source repository is [`biohub/ESMC-6B`](https://huggingface.co/biohub/ESMC-6B).
The Biohub ESM license is available at https://github.com/Biohub/esm/blob/main/LICENSE.md.

## Citation

```bibtex
@misc{FastPLMs,
  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
  title={FastPLMs: Fast, efficient, protein language model inference from Hugging Face AutoModel.},
  year={2024},
  url={https://huggingface.co/Synthyra/ESMplusplus_6B},
  DOI={10.57967/hf/3726},
  publisher={Hugging Face}
}
```

```bibtex
@misc{candido2026language,
  title  = {Language Modeling Materializes a World Model of Protein Biology},
  author = {Candido, Salvatore and Hayes, Thomas and Derry, Alexander and Rao, Roshan
            and Lin, Zeming and Verkuil, Robert and Wu, Bryan and Lee, Jin Sub
            and Bruguera, Elise S. and Keval, Jehan A. and Kopylov, Mykhailo
            and Pak, John E. and Wu, Wesley and Thomas, Neil and Mataraso, Samson
            and Hsu, Alvin and Trotman-Grant, Ashton C. and Fatras, Kilian
            and dos Santos Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil
            and Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth
            and Sitwala, Hrishita and Yu, Yue and Wiggert, Marius
            and Carlin, Dylan Alexander and Goering, Anthony W. and Blazejewski, Tomasz
            and Sandora, McCullen and Hla, Michael and Jia, Tina Z.
            and Kloker, Leon H. and Sofroniew, Nicholas J. and Uehara, Masatoshi
            and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S.
            and Sercu, Tom and Rives, Alexander},
  year   = {2026},
  url    = {https://biohub.ai/papers/esm_protein.pdf},
  note   = {Preprint}
}
```