Duplicate from biohub/ESMFold2

Co-authored-by: Fausto Milletari <fmilletari-czi@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+images/folding_evals.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,228 @@
+---
+license:
+- mit
+- other
+license_link: https://github.com/Biohub/esm/blob/main/THIRD_PARTY_NOTICE.md
+language: en
+tags:
+- biology
+- esm
+- protein
+- protein-structure-prediction
+- structure-prediction
+- protein-design
+- 3d-structure
+- confidence-estimation
+- molecular-dynamics
+- transformers
+---
+
+# ESMFold2
+
+## Model Details
+
+ESMFold2 is a state-of-the-art model for protein structure prediction and design that defines a new frontier for speed and accuracy. The model predicts high-resolution, all-atom 3D protein structures directly from amino acid sequences, with optional multiple sequence alignment (MSA) input for enhanced accuracy on challenging targets. The model outputs comprehensive structural information including all-atom coordinates (backbone and side chains), confidence metrics (pLDDT, pAE, pTM, iPTM), and optional distogram predictions for detailed analysis of predicted structures. Unlike ESMFold, ESMFold2 is able to predict structures for all biomolecules, including small molecules, DNA, RNA, and modified amino acids.
+
+ESMFold2 is capable of either single-sequence or MSA conditioned structure prediction for improved accuracy on difficult targets. The ESMFold2-Fast variant is an inference optimized single-sequence structure prediction model and is not MSA conditioned.
+
+To run this model with the Biohub Platform API, visit the [Biohub Platform](https://biohub.ai/).
+
+Read more about ESMFold2 in our paper [here](https://biohub.ai/papers/esm_protein.pdf).
+
+## Model Variants
+
+| Model | MSA Conditioning | Description | Data Cutoff |
+| :---- | :---- | :---- | :---- |
+| [ESMFold2](https://huggingface.co/biohub/ESMFold2) | Yes | Large model, capable of either single-sequence or MSA conditioned structure prediction for improved accuracy on difficult targets | Sept 2021 |
+| [ESMFold2-Fast](https://huggingface.co/biohub/ESMFold2-Fast) | No | Inference optimized single-sequence structure prediction model | Sept 2021 |
+
+## Performance Metrics
+
+ESMfold2 was evaluated against state-of-the-art single-sequence and MSA-based structure prediction models on the FoldBench benchmark. ESMFold2 meets or exceeds performance by AlphaFold3 on antibody-antigen complex prediction, protein-protein complex prediction and [Runs N' Poses](https://www.biorxiv.org/content/10.1101/2025.02.03.636309v1) benchmarks. Inference-time compute can dramatically improve performance of ESMFold2, especially across antibody-antigen complexes.
+
+![][image1]
+
+Refer to the [paper](https://biohub.ai/papers/esm_protein.pdf) for details on additional performance metrics.
+
+### Usage
+
+Please install `esm` from GitHub (a PyPI release is coming soon):
+
+```
+pip install esm@git+https://github.com/Biohub/esm.git@main
+```
+
+You can fold your first protein with:
+
+```py
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+
+# Ubiquitin (PDB 1UBQ)
+sequence = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"
+
+# optionally use "biohub/ESMFold2"
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2-Fast").cuda().eval()
+output = model.infer_protein(sequence, num_loops=3, num_sampling_steps=50)
+
+print(f"pLDDT mean: {float(output['plddt'].mean()):.3f}, pTM: {float(output['ptm'].mean()):.3f}")
+```
+
+You can also fold complex biomolecules — proteins, DNA/RNA (with modified residues), and small-molecule ligands all at once. Here's an example folding the HhaI DNA methyltransferase + its cognate DNA (with a trapped 5-fluoro-2′-deoxycytidine, CCD `C36`) + the SAH cofactor (PDB [1MHT](https://www.rcsb.org/structure/1MHT)):
+
+```py
+from esm.models.esmfold2 import (
+    DNAInput,
+    ESMFold2InputBuilder,
+    LigandInput,
+    Modification,
+    ProteinInput,
+    StructurePredictionInput,
+)
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+
+HHAI_SEQ = (
+    "MIEIKDKQLTGLRFIDLFAGLGGFRLALESCGAECVYSNEWDKYAQEVYEMNFGEKPEGDITQVNEKTIPDH"
+    "DILCAGFPCQAFSISGKQKGFEDSRGTLFFDIARIVREKKPKVVFMENVKNFASHDNGNTLEVVKNTMNELD"
+    "YSFHAKVLNALDYGIPQKRERIYMICFRNDLNIQNFQFPKPFELNTFVKDLLLPDSEVEHLVIDRKDLVMTN"
+    "QEIEQTTPKTVRLGIVGKGGQGERIYSTRGIAITLSAYGGGIFAKTGGYLVNGKTRKLHPRECARVMGYPDS"
+    "YKVHPSTSQAYKQFGNSVVINVLQYIAYNIGSSLNFKPY"
+)
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
+
+spi = StructurePredictionInput(
+    sequences=[
+        ProteinInput(id="A", sequence=HHAI_SEQ),
+        DNAInput(
+            id="B",
+            sequence="GATAGCGCTATC",
+            modifications=[Modification(position=5, ccd="C36")],
+        ),
+        DNAInput(
+            id="C",
+            sequence="TGATAGCGCTATC",
+            modifications=[Modification(position=6, ccd="C36")],
+        ),
+        LigandInput(id="L", ccd=["SAH"]),
+    ]
+)
+
+result = ESMFold2InputBuilder().fold(
+    model, spi, num_loops=3, num_sampling_steps=50, num_diffusion_samples=1, seed=0
+)
+
+print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+
+with open("1mht_pred.cif", "w") as f:
+    f.write(result.complex.to_mmcif())
+```
+
+For the Biohub API, first generate an [API key](https://biohub.ai/developer-console/api-keys) and add it to your Biohub account. The code below assumes the environment variable `ESM_API_KEY=$YOUR_API_KEY`.
+
+```py
+import os
+
+from esm.models.esmfold2 import (
+    DNAInput,
+    LigandInput,
+    Modification,
+    ProteinInput,
+    StructurePredictionInput,
+)
+from esm.sdk import esmfold2_client
+from esm.sdk.api import FoldingConfig
+
+HHAI_SEQ = (
+    "MIEIKDKQLTGLRFIDLFAGLGGFRLALESCGAECVYSNEWDKYAQEVYEMNFGEKPEGDITQVNEKTIPDH"
+    "DILCAGFPCQAFSISGKQKGFEDSRGTLFFDIARIVREKKPKVVFMENVKNFASHDNGNTLEVVKNTMNELD"
+    "YSFHAKVLNALDYGIPQKRERIYMICFRNDLNIQNFQFPKPFELNTFVKDLLLPDSEVEHLVIDRKDLVMTN"
+    "QEIEQTTPKTVRLGIVGKGGQGERIYSTRGIAITLSAYGGGIFAKTGGYLVNGKTRKLHPRECARVMGYPDS"
+    "YKVHPSTSQAYKQFGNSVVINVLQYIAYNIGSSLNFKPY"
+)
+
+client = esmfold2_client(model="esmfold2-fast-2026-05", token=os.environ["ESM_API_KEY"])
+
+spi = StructurePredictionInput(
+    sequences=[
+        ProteinInput(id="A", sequence=HHAI_SEQ),
+        DNAInput(
+            id="B",
+            sequence="GATAGCGCTATC",
+            modifications=[Modification(position=5, ccd="C36")],
+        ),
+        DNAInput(
+            id="C",
+            sequence="TGATAGCGCTATC",
+            modifications=[Modification(position=6, ccd="C36")],
+        ),
+        LigandInput(id="L", ccd=["SAH"]),
+    ]
+)
+
+result = client.fold_all_atom(spi, config=FoldingConfig(num_loops=3, num_sampling_steps=50))
+
+print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+```
+
+## Training Data
+
+ESMfold2 was trained on sequences from the Protein Data Bank (PDB) and the AlphaFold DB (AFDB).
+
+## Frontier Safety
+
+Biohub has established a safety team to assess the benefits and potential risks of our models and tools prior to release, and develop mitigations where necessary. Risk assessment was conducted for ESMFold2 prior to release. Further details are available in our corresponding paper appendix.
+
+Informed by our risk assessments, we are releasing the source code and model weights for ESMFold2.
+
+[Biohub.ai](http://Biohub.ai) Platform: We implement guardrails that detect and restrict the use of keywords and sequences corresponding to controlled pathogens and toxins on our freely accessible platform. For further details regarding these guardrails, please refer to our Biohub platform Resources page.
+
+## Biases and Limitations
+
+- **Dataset biases**: The model may reflect biases present in the training data (PDB, AFDB), including over-representation of certain protein families, experimental conditions, or structural classes. Performance may vary for underrepresented protein types.
+- **Dataset limitations:** PDB historically lacks comprehensive data on protein conformations, post-translational modifications, disordered regions, etc. Like all other structure prediction models trained on the PDB, performance may degrade on other biomolecules.
+- **Computational demand:** Highest accuracy structure predictions require scaling inference time compute. Predictions made with reduced inference parameters may lead to suboptimal performance.
+- **Experimental validation required**: All predictions should be considered hypotheses requiring experimental validation. The model cannot replace experimental structure determination methods (X-ray crystallography, cryo-EM, NMR) for definitive structural characterization.
+
+### Out-of-Scope or Unauthorized Use Cases
+
+Do not use the model for the following purposes:
+
+- Any use that is prohibited by the [Acceptable Use Policy](https://biohub.org/acceptable-use-policy/).
+
+### Caveats and Recommendations
+
+- Always review and validate outputs generated by the model.
+- Treat model outputs as machine-generated hypotheses that require further experimental validation, not as established biological facts.
+- We are committed to advancing the responsible development and use of artificial intelligence.
+
+Should you have any security or privacy issues or questions related to this model, please reach out to our team at [support@biohub.org](mailto:support@biohub.org).
+
+### Citation
+
+```
+@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}
+}
+```
+
+## Acknowledgements
+
+Many people on the Biohub AI Research team and prior EvolutionaryScale team contributed to the development of this model. It would not have been possible without them.
+
+[image1]: images/folding_evals.png

ccd.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5
+size 417306584

config.json

@@ -0,0 +1,105 @@
+{
+  "architectures": [
+    "ESMFold2Model"
+  ],
+  "confidence_head": {
+    "distogram_bins": 39,
+    "enabled": true,
+    "folding_trunk": {
+      "dropout": 0.25,
+      "n_heads": 8,
+      "n_layers": 4
+    },
+    "max_dist": 50.75,
+    "min_dist": 3.25,
+    "num_pae_bins": 64,
+    "num_pde_bins": 64,
+    "num_plddt_bins": 50
+  },
+  "d_pair": 256,
+  "d_single": 384,
+  "disable_msa_features": false,
+  "dtype": "float32",
+  "esmc_id": "biohub/ESMC-6B",
+  "folding_trunk": {
+    "dropout": 0.25,
+    "n_heads": 8,
+    "n_layers": 48
+  },
+  "force_lm_dropout_during_inference": false,
+  "inputs": {
+    "atom_encoder": {
+      "d_atom": 128,
+      "d_token": 768,
+      "expansion_ratio": 2,
+      "n_blocks": 3,
+      "n_heads": 4,
+      "n_spatial_rope_pairs_per_axis": 2,
+      "n_uid_rope_pairs": 10,
+      "spatial_rope_base_frequency": 20.0,
+      "swa_window_size": 128,
+      "uid_rope_base_frequency": 10000.0
+    },
+    "d_inputs": 451
+  },
+  "lm_d_model": 2560,
+  "lm_dropout": 0.0,
+  "lm_encoder": {
+    "enabled": true,
+    "lm_dropout": 0.25,
+    "n_layers": 4,
+    "per_loop_lm_dropout": true
+  },
+  "lm_num_layers": 80,
+  "model_type": "esmfold2",
+  "msa_encoder": {
+    "d_hidden": 32,
+    "d_msa": 128,
+    "enabled": true,
+    "msa_head_width": 16,
+    "n_heads_msa": 8,
+    "n_layers": 4
+  },
+  "msa_encoder_overwrite": true,
+  "n_relative_chain_bins": 2,
+  "n_relative_residx_bins": 32,
+  "num_diffusion_samples": 32,
+  "num_loops": 3,
+  "parcae": {
+    "coda_n_layers": 2,
+    "enabled": true,
+    "max_steps": 6,
+    "min_steps": 1,
+    "poisson_mean": 3.0
+  },
+  "structure_head": {
+    "diffusion_module": {
+      "atom_num_blocks": 3,
+      "atom_num_heads": 4,
+      "c_atom": 128,
+      "c_s_inputs": 451,
+      "c_token": 768,
+      "c_z": 256,
+      "fourier_dim": 256,
+      "relpos_r_max": 32,
+      "relpos_s_max": 2,
+      "sigma_data": 16.0,
+      "token_num_blocks": 12,
+      "token_num_heads": 16,
+      "transition_multiplier": 2
+    },
+    "distogram_bins": 64,
+    "gamma_0": 0.8,
+    "gamma_min": 1.0,
+    "inference_num_steps": 14,
+    "inference_p": 7.0,
+    "inference_s_max": 160.0,
+    "inference_s_min": 0.0004,
+    "noise_scale": 1.003,
+    "step_scale": 1.5,
+    "train_noise_log_mean": -1.2,
+    "train_noise_log_std": 1.5
+  },
+  "transformers_version": "4.57.6",
+  "type": "release"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e
+size 939505228