File size: 7,617 Bytes
b44701d 407875b b44701d 407875b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 | ---
library_name: transformers
tags:
- biology
- protein-structure
- esmfold2
- multimodal-protein-model
---
# FastPLMs ESMFold2
FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for
Biohub's ESMFold2, ESMFold2-Fast, and experimental ESMFold2 structure
predictors. It vendors the released Biohub ESMFold2 model code, input builder,
MSA helpers, and structure export utilities, while loading the PLM backbone
through FastPLMs ESM++.
## Load With AutoModel
```python
import torch
from transformers import AutoModel
model = AutoModel.from_pretrained(
"Synthyra/ESMFold2-Fast",
trust_remote_code=True,
dtype=torch.float32,
).eval().cuda()
```
Use `Synthyra/ESMFold2` for the full model, `Synthyra/ESMFold2-Fast` for the
faster release variant, and the `Synthyra/ESMFold2-Experimental*` checkpoints
for differentiable binder design and experimental critic ensembles.
The folding trunk runs in fp32; the 6B FastPLMs ESM++ backbone is loaded in
bf16 by default via `esmc_precision="bf16"` and uses the flex attention backend
by default inside ESMFold2.
## Fold One Protein
```python
sequence = "MKTLLILAVVAAALA"
result = model.fold_protein(
sequence,
num_loops=3,
num_sampling_steps=50,
num_diffusion_samples=1,
seed=0,
)
print(float(result.plddt.mean()))
print(float(result.ptm))
```
## Experimental Test-Time Training
TTT is disabled by default. Standard `fold_protein(...)`, `fold(...)`, raw tensor
inference, and `state_dict()` keys are unchanged unless you explicitly pass
`ttt=True` or call `fold_protein_ttt(...)`.
The ESMFold2 TTT path is experimental and protein-only in v1. It trains local
LoRA adapters only on `_esmc` with a masked language modeling objective. The
folding trunk, confidence head, diffusion head, and structure input pipeline are
frozen. TTT can improve difficult low-confidence folds, but it adds substantial
test-time compute and can degrade already confident predictions.
```python
result = model.fold_protein(
"MSTNPKPQRKTKRNT",
num_loops=1,
num_sampling_steps=10,
num_diffusion_samples=1,
seed=0,
ttt=True,
ttt_config={
"steps": 1,
"ags": 1,
"batch_size": 1,
"lora_rank": 8,
"lora_alpha": 32.0,
},
)
print(result.ttt_metrics["losses"])
print(result.ttt_metrics["step_plddts"])
print(result.ttt_metrics["best_step"])
```
`load_esmc=True` is required for TTT because the ESM++ MLM head is loaded lazily
from `config.esmc_id`. If that pretrained MLM head cannot be loaded, TTT raises
an assertion instead of silently using a random head.
## Save mmCIF or PDB
```python
model.save_as_cif(result, "prediction.cif")
model.save_as_pdb(result, "prediction.pdb")
cif_text = model.result_to_cif(result)
pdb_text = model.result_to_pdb(result)
```
`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
## Fold Complexes
```python
types = model.input_types
complex_input = types.StructurePredictionInput(
sequences=[
types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
types.DNAInput(id="B", sequence="GATAGC"),
types.LigandInput(id="L", ccd=["SAH"]),
]
)
result = model.fold(
complex_input,
num_loops=3,
num_sampling_steps=50,
num_diffusion_samples=1,
seed=0,
)
model.save_as_cif(result, "complex_prediction.cif")
```
## Binder Design With FastPLMs ESMFold2
FastPLMs includes a FastPLMs-only port of the Biohub ESMFold2 binder design
tutorial at `cookbook/tutorials/binder_design_fastplms.py`. The workflow uses
ESMFold2 experimental checkpoints for differentiable folding losses, ESM++ for
sequence regularization, and ESMFold2 hero critics for final confidence scoring.
The optimizer follows the official strategy:
1. Optimize mutable `#` residues as continuous amino acid logits.
2. Suppress cysteine design by masking cysteine logits and gradients.
3. Backpropagate through ESMFold2 `res_type_soft` using intra-contact,
inter-contact, and globularity losses from the distogram.
4. Add an ESM++ masked-LM pseudoperplexity regularizer on mutable binder
residues.
5. Keep the late-trajectory sequence with the best iPTM.
6. Fold the selected sequence with the final critic ensemble and write
`results.parquet`, `selection.parquet`, `trajectory.jsonl`,
`best_sequences.fasta`, and per-critic PDB/CIF/logit files.
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
```
Verified result:
| Metric | Value |
| :--- | :--- |
| Binder length | `128` |
| Seed | `103` |
| Steps | `150` |
| Hero mean iPTM | `0.913870` |
| Hero min iPTM | `0.904600` |
| All four hero critics above 0.9 | `True` |
Binder sequence:
```text
SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ
```
See the full guide in [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
for Modal execution, official pI and selection scoring, per-critic metrics, and
the tested cheaper step-count boundary.
## Use MSAs
```python
types = model.input_types
msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
input_with_msa = types.StructurePredictionInput(
sequences=[
types.ProteinInput(id="A", sequence=msa.query, msa=msa),
]
)
result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
```
## Raw Tensor Inference
```python
features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
with torch.inference_mode():
output = model(
**features,
num_loops=3,
num_sampling_steps=50,
num_diffusion_samples=1,
)
decoded = model.input_builder.decode(output, features, chain_infos)
```
Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESM++ backbone:
```python
model = AutoModel.from_pretrained(
"Synthyra/ESMFold2-Fast",
trust_remote_code=True,
load_esmc=False,
).cuda().eval()
```
For FP8 LM inference, install `transformer_engine.pytorch` in a CUDA
environment with FP8-capable hardware and load the shared FastPLMs ESM++
backbone with:
```python
model = AutoModel.from_pretrained(
"Synthyra/ESMFold2-Fast",
trust_remote_code=True,
esmc_precision="fp8",
).cuda().eval()
```
FP8 is inference-only for the ESMFold2 LM backbone. TTT remains a bf16/fp32
path.
|