ferritin_plms/esmfold2/
pretrained.rs

1//! ESMFold2 pretrained model loading.
2//!
3//! Downloads the structure-head weights from HuggingFace and optionally wires
4//! in the ESMC-6B backbone for end-to-end structure prediction.
5//!
6//! ## Weight layout
7//!
8//! `biohub/ESMFold2-Fast` `model.safetensors` (~755 MB) contains the structure
9//! head only. The ESMC-6B backbone (~12 GB) must be loaded from a second repo
10//! (`biohub/ESMC-6B`).
11//!
12//! | Weight prefix          | Component                                 |
13//! |------------------------|-------------------------------------------|
14//! | `lm_encoder.*`         | LM adapter (4 blocks, 2560 → d_single=384) |
15//! | `folding_trunk.*`      | 24-layer Pairformer trunk                 |
16//! | `inputs.atom_encoder.*`| Atom encoder (3 blocks, d_atom=128→768)   |
17//! | `msa_encoder.*`        | MSA encoder (disabled in Fast)            |
18//! | `structure_head.*`     | Diffusion module (token 12-block + atom 3-block) |
19//! | `confidence_head.*`    | pLDDT / pAE / pDE / distogram heads       |
20//!
21//! ## Two loading modes
22//!
23//! - [`ESMFold2Runner::from_pretrained`] — structure head only (~755 MB).
24//!   Backbone hidden states are zero-initialised; coordinates will be
25//!   physically meaningless but all shapes are correct (useful for shape testing).
26//!
27//! - [`ESMFold2Runner::from_pretrained_with_backbone`] — structure head +
28//!   ESMC-6B backbone (~12 GB total). Required for real structure predictions.
29
30use super::config::ESMFold2Config;
31use super::model::ESMFold2Model;
32use super::output::ESMFold2Output;
33use crate::esmc::pretrained::{ESMCModels, ESMCRunner};
34use anyhow::Result;
35use candle_core::{DType, Device, Tensor};
36use candle_nn::VarBuilder;
37use hf_hub::HFClientSync;
38
39const ESMFOLD2_DTYPE: DType = DType::F32;
40
41// ── ESMFold2Models enum ───────────────────────────────────────────────────────
42
43/// Available ESMFold2 model variants hosted on HuggingFace.
44pub enum ESMFold2Models {
45    /// ESMFold2-Fast — single-sequence only, no MSA, optimized for speed.
46    /// Use `num_sampling_steps = 50` for quality, `14` for speed.
47    Fast,
48    /// ESMFold2 (Full) — optional MSA conditioning for higher accuracy.
49    Full,
50}
51
52impl ESMFold2Models {
53    /// Returns `(hf_repo_id, config)` for this variant.
54    pub fn model_info(&self) -> (&'static str, ESMFold2Config) {
55        match self {
56            Self::Fast => ("biohub/ESMFold2-Fast", ESMFold2Config::fast()),
57            Self::Full => ("biohub/ESMFold2", ESMFold2Config::fast()),
58        }
59    }
60}
61
62// ── ESMFold2Runner ────────────────────────────────────────────────────────────
63
64/// Wraps a loaded ESMFold2 model for all-atom structure prediction inference.
65///
66/// Create with:
67/// - [`from_pretrained`][Self::from_pretrained] — structure head only (~755 MB),
68///   backbone stubs with zeros (shape testing only).
69/// - [`from_pretrained_with_backbone`][Self::from_pretrained_with_backbone] —
70///   structure head + ESMC-6B backbone (~12 GB total, real predictions).
71pub struct ESMFold2Runner {
72    model: ESMFold2Model,
73    config: ESMFold2Config,
74    device: Device,
75    /// ESMC-6B backbone. `None` → hidden states are zeroed (stub mode).
76    backbone: Option<ESMCRunner>,
77}
78
79impl ESMFold2Runner {
80    /// Download and load the ESMFold2 structure head only (~755 MB).
81    ///
82    /// The ESMC-6B backbone is **not** loaded; hidden states passed to
83    /// `fold_protein` will be all-zeros, so predicted coordinates are
84    /// physically meaningless. Use this for shape/pipeline testing without
85    /// requiring 12 GB of disk space.
86    pub fn from_pretrained(model_variant: ESMFold2Models, device: Device) -> Result<Self> {
87        let (repo_id, config) = model_variant.model_info();
88        let model = Self::load_structure_head(repo_id, &config, &device)?;
89        Ok(Self { model, config, device, backbone: None })
90    }
91
92    /// Download and load the structure head (~755 MB) **and** the ESMC-6B
93    /// backbone (~12 GB).  Required for predictions with meaningful pLDDT.
94    pub fn from_pretrained_with_backbone(
95        model_variant: ESMFold2Models,
96        device: Device,
97    ) -> Result<Self> {
98        let (repo_id, config) = model_variant.model_info();
99        let model = Self::load_structure_head(repo_id, &config, &device)?;
100        eprintln!("ESMFold2Runner: loading ESMC-6B backbone (~12 GB)...");
101        let backbone = ESMCRunner::from_pretrained(ESMCModels::ESMC6B, device.clone())?;
102        eprintln!("ESMFold2Runner: backbone ready.");
103        Ok(Self { model, config, device, backbone: Some(backbone) })
104    }
105
106    fn load_structure_head(
107        repo_id: &str,
108        config: &ESMFold2Config,
109        device: &Device,
110    ) -> Result<ESMFold2Model> {
111        eprintln!("ESMFold2Runner: downloading structure head from {repo_id}...");
112        let (owner, name) = repo_id.split_once('/').unwrap_or(("", repo_id));
113        let client = HFClientSync::new()?;
114        let weights_path = client
115            .model(owner, name)
116            .download_file()
117            .filename("model.safetensors")
118            .send()?;
119        eprintln!("ESMFold2Runner: weights at {}", weights_path.display());
120
121        let vb = unsafe {
122            VarBuilder::from_mmaped_safetensors(&[&weights_path], ESMFOLD2_DTYPE, device)?
123        };
124        Ok(ESMFold2Model::load(vb, config.clone())?)
125    }
126
127    /// Fold a single protein sequence.
128    ///
129    /// # Arguments
130    /// * `sequence`           — amino-acid string (standard one-letter codes)
131    /// * `num_loops`          — recycling iterations (typically 3)
132    /// * `num_sampling_steps` — diffusion denoising steps (14 fast / 50 quality)
133    ///
134    /// When the runner was created with [`from_pretrained`][Self::from_pretrained]
135    /// (no backbone), hidden states are zeros and coordinates will be arbitrary.
136    /// Use [`from_pretrained_with_backbone`][Self::from_pretrained_with_backbone]
137    /// for real predictions.
138    pub fn fold_protein(
139        &self,
140        sequence: &str,
141        num_loops: usize,
142        num_sampling_steps: usize,
143    ) -> Result<ESMFold2Output> {
144        let l = sequence.len();
145        let device = &self.device;
146
147        // ESMC-6B backbone → (1, L, 2560).
148        // When backbone is None, use zeros (stub for shape testing).
149        let hidden_states = match &self.backbone {
150            Some(esmc) => {
151                // embed_sequence returns (1, L+2, d_model) with BOS and EOS tokens.
152                let embs = esmc.embed_sequence(sequence)?;
153                // Strip BOS (index 0) and EOS (index L+1), keeping L residues.
154                embs.narrow(1, 1, l)?
155            }
156            None => Tensor::zeros(
157                &[1, l, self.config.lm_d_model],
158                ESMFOLD2_DTYPE,
159                device,
160            )?,
161        };
162
163        let res_idx: Vec<f32> = (0..l).map(|i| i as f32).collect();
164        let residue_indices = Tensor::from_vec(res_idx, &[1, l], device)?;
165        let chain_ids = Tensor::zeros(&[1, l], ESMFOLD2_DTYPE, device)?;
166
167        Ok(self.model.forward(
168            &hidden_states,
169            &residue_indices,
170            &chain_ids,
171            num_loops,
172            num_sampling_steps,
173        )?)
174    }
175}
176
177// ── Tests ─────────────────────────────────────────────────────────────────────
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use candle_core::Device;
183
184    /// Shape smoke-test with stub backbone (no download needed).
185    /// Uses a tiny config so the test runs in milliseconds.
186    #[test]
187    fn test_fold_protein_stub_shapes() {
188        use super::super::config::ESMFold2Config;
189        use candle_nn::VarBuilder;
190        use candle_core::DType;
191
192        let device = Device::Cpu;
193        let cfg = ESMFold2Config {
194            lm_d_model: 64,
195            lm_encoder_n_layers: 1,
196            d_single: 32,
197            d_pair: 16,
198            trunk_n_layers: 1,
199            trunk_n_heads: 2,
200            c_token: 32,
201            c_atom: 16,
202            token_num_blocks: 1,
203            token_num_heads: 2,
204            fourier_dim: 16,
205            num_plddt_bins: 50,
206            num_pae_bins: 64,
207            num_pde_bins: 64,
208            distogram_bins: 39,
209            ..ESMFold2Config::fast()
210        };
211        let vb = VarBuilder::zeros(DType::F32, &device);
212        let model = ESMFold2Model::load(vb, cfg.clone()).unwrap();
213        let runner = ESMFold2Runner { model, config: cfg, device: device.clone(), backbone: None };
214
215        let seq = "ACDEFGHIK"; // 9 residues
216        let out = runner.fold_protein(seq, 1, 2).unwrap();
217
218        assert_eq!(out.sample_atom_coords.dims(), &[1, seq.len(), 3]);
219        assert_eq!(out.plddt.dims(), &[1, seq.len()]);
220        assert_eq!(out.ptm.dims(), &[1]);
221        assert_eq!(out.iptm.dims(), &[1]);
222    }
223
224    /// Integration test: downloads structure head (~755 MB) and folds ubiquitin.
225    /// The backbone is stubbed with zeros so coordinates are not meaningful,
226    /// but all shapes and the full pipeline must succeed.
227    ///
228    /// Run with:
229    ///   cargo test -p ferritin-plms test_esmfold2_fold_stub_integration -- --ignored
230    #[test]
231    #[ignore = "requires downloading biohub/ESMFold2-Fast weights (~755 MB)"]
232    fn test_esmfold2_fold_stub_integration() {
233        let device = Device::Cpu;
234        let runner =
235            ESMFold2Runner::from_pretrained(ESMFold2Models::Fast, device).expect("load failed");
236
237        // Ubiquitin (76 aa)
238        let seq = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG";
239        let out = runner.fold_protein(seq, 1, 2).expect("fold_protein failed");
240
241        assert_eq!(out.sample_atom_coords.dims(), &[1, seq.len(), 3]);
242        assert_eq!(out.plddt.dims(), &[1, seq.len()]);
243    }
244
245    /// Full end-to-end integration: ESMC-6B backbone + structure head.
246    /// Folds ubiquitin and asserts pLDDT > 0.5 on average (backbone gives
247    /// real signal; exact quality depends on diffusion steps).
248    ///
249    /// Run with:
250    ///   cargo test -p ferritin-plms test_esmfold2_fold_with_backbone -- --ignored
251    #[test]
252    #[ignore = "requires biohub/ESMFold2-Fast (~755 MB) + biohub/ESMC-6B (~12 GB)"]
253    fn test_esmfold2_fold_with_backbone() -> Result<()> {
254        let device = Device::Cpu;
255        let runner =
256            ESMFold2Runner::from_pretrained_with_backbone(ESMFold2Models::Fast, device)?;
257
258        let seq = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG";
259        let out = runner.fold_protein(seq, 3, 14)?;
260
261        assert_eq!(out.sample_atom_coords.dims(), &[1, seq.len(), 3]);
262        assert_eq!(out.plddt.dims(), &[1, seq.len()]);
263
264        let mean_plddt = out.plddt.mean_all()?.to_scalar::<f32>()?;
265        assert!(mean_plddt > 0.5, "expected mean pLDDT > 0.5, got {mean_plddt:.3}");
266
267        Ok(())
268    }
269}
ferritin_plms/esmfold2/pretrained.rs

ferritin_plms/esmfold2/
pretrained.rs
