ferritin_plms/esmc/
pretrained.rs

1//! ESMC pretrained model loading.
2//!
3//! Downloads weights from `biohub/ESMC-{300M,600M,6B}` on HuggingFace and
4//! wraps the ESMC model for sequence embedding.
5//!
6//! ## Weight key layout (`biohub/ESMC-*` safetensors)
7//!
8//! The weights are saved from `ESMCForMaskedLM`, a HuggingFace wrapper that
9//! nests the backbone under an `esmc` attribute. All backbone keys are
10//! therefore prefixed with `esmc.`:
11//!
12//! | Python (HF safetensors key)                              | Rust VarBuilder path                          |
13//! |----------------------------------------------------------|-----------------------------------------------|
14//! | `esmc.embed.weight`                                      | `embed.weight`                                |
15//! | `esmc.transformer.blocks.{i}.attn.layernorm_qkv.0.*`   | `transformer.blocks.{i}.attn.layernorm_qkv.0.*` |
16//! | `esmc.transformer.blocks.{i}.attn.layernorm_qkv.1.*`   | `transformer.blocks.{i}.attn.layernorm_qkv.1.*` |
17//! | `esmc.transformer.blocks.{i}.attn.out_proj.weight`      | `transformer.blocks.{i}.attn.out_proj.weight`   |
18//! | `esmc.transformer.blocks.{i}.attn.q_ln.weight`          | `transformer.blocks.{i}.attn.q_ln.weight`       |
19//! | `esmc.transformer.blocks.{i}.attn.k_ln.weight`          | `transformer.blocks.{i}.attn.k_ln.weight`       |
20//! | `esmc.transformer.blocks.{i}.attn.rotary.*`             | `transformer.blocks.{i}.attn.rotary.*`          |
21//! | `esmc.transformer.blocks.{i}.ffn.0.*`                   | `transformer.blocks.{i}.ffn.0.*`                |
22//! | `esmc.transformer.blocks.{i}.ffn.1.weight`              | `transformer.blocks.{i}.ffn.1.weight`           |
23//! | `esmc.transformer.blocks.{i}.ffn.3.weight`              | `transformer.blocks.{i}.ffn.3.weight`           |
24//! | `esmc.transformer.norm.weight`                           | `transformer.norm.weight`                       |
25//! | `esmc.sequence_head.0.weight/bias`                       | `sequence_head.0.weight/bias`                   |
26//! | `esmc.sequence_head.2.weight`                            | `sequence_head.2.weight`                        |
27//! | `esmc.sequence_head.3.weight/bias`                       | `sequence_head.3.weight/bias`                   |
28//!
29//! The `from_pretrained` loader auto-detects whether the `esmc.` prefix is
30//! present and sets the VarBuilder root accordingly, so the same Rust model
31//! code works against both the wrapped and unwrapped formats.
32
33use crate::esmc::models::esmc::{ESMC, ESMCConfig};
34use crate::plm_runner::PlmRunner;
35use anyhow::Result;
36use candle_core::{DType, Device, Tensor};
37use candle_nn::VarBuilder;
38use hf_hub::HFClientSync;
39
40const ESMC_DTYPE: DType = DType::F32;
41
42/// Available ESMC model variants hosted on HuggingFace.
43pub enum ESMCModels {
44    /// ESMC 300M — 30 layers, d_model=960 (~1.3 GB weights)
45    ESMC300M,
46    /// ESMC 600M — 36 layers, d_model=1152
47    ESMC600M,
48    /// ESMC 6B — 80 layers, d_model=2560 (backbone for ESMFold2)
49    ESMC6B,
50}
51
52impl ESMCModels {
53    /// Returns `(hf_repo_id, config)` for this variant.
54    pub fn model_info(&self) -> (&'static str, ESMCConfig) {
55        match self {
56            Self::ESMC300M => ("biohub/ESMC-300M", ESMCConfig::esmc_300m()),
57            Self::ESMC600M => ("biohub/ESMC-600M", ESMCConfig::esmc_600m()),
58            Self::ESMC6B => ("biohub/ESMC-6B", ESMCConfig::esmc_6b()),
59        }
60    }
61}
62
63/// Wraps a loaded ESMC model for sequence embedding inference.
64pub struct ESMCRunner {
65    model: ESMC,
66}
67
68impl ESMCRunner {
69    /// Download weights from HuggingFace and load the model.
70    ///
71    /// Handles the `ESMCForMaskedLM` wrapper prefix (`esmc.`) transparently:
72    /// probes for `esmc.embed.weight` in the safetensors and sets the
73    /// VarBuilder root to `vb.pp("esmc")` when found, otherwise uses the
74    /// flat (unwrapped) layout.
75    pub fn from_pretrained(model: ESMCModels, device: Device) -> Result<Self> {
76        let (repo_id, config) = model.model_info();
77        let (owner, name) = repo_id.split_once('/').unwrap_or(("", repo_id));
78        let client = HFClientSync::new()?;
79        let weights_path = client
80            .model(owner, name)
81            .download_file()
82            .filename("model.safetensors")
83            .send()?;
84
85        let vb =
86            unsafe { VarBuilder::from_mmaped_safetensors(&[&weights_path], ESMC_DTYPE, &device)? };
87
88        // Detect whether weights use the HF ESMCForMaskedLM prefix "esmc."
89        // by probing for the embedding matrix at the prefixed path.
90        let vb_root = if vb
91            .get((config.embedding_dim, config.d_model), "esmc.embed.weight")
92            .is_ok()
93        {
94            vb.pp("esmc")
95        } else {
96            vb
97        };
98
99        let esmc = ESMC::load(vb_root, config)?;
100        Ok(Self { model: esmc })
101    }
102
103    /// Tokenize `sequence` and run a forward pass.
104    ///
105    /// Returns per-residue embeddings with shape `(1, L, d_model)` where
106    /// `L` includes the BOS and EOS tokens added by the tokenizer.
107    pub fn embed_sequence(&self, sequence: &str) -> Result<Tensor> {
108        let tokens = self.model.encode(sequence)?;
109        let tokens = tokens.unsqueeze(0)?;
110        let output = self.model.forward(&tokens, None, false)?;
111        output
112            .embeddings
113            .ok_or_else(|| anyhow::anyhow!("ESMC forward() returned no embeddings"))
114    }
115}
116
117impl PlmRunner for ESMCRunner {
118    /// Delegate to `embed_sequence`, returning per-residue embeddings `(1, L, d_model)`.
119    fn embed(&self, sequence: &str) -> Result<Tensor> {
120        self.embed_sequence(sequence)
121    }
122
123    fn model_name(&self) -> &str {
124        "esmc"
125    }
126}
ferritin_plms/esmc/pretrained.rs

ferritin_plms/esmc/
pretrained.rs
