ferritin_plms/esm3/
pretrained.rs

1//! ESM3 pretrained weight loading.
2//!
3//! Downloads weights from `EvolutionaryScale/esm3-sm-open-v1` on HuggingFace (gated access —
4//! the user must accept the Cambrian Non-Commercial license and run `huggingface-cli login`).
5//!
6//! ## Weight layout
7//!
8//! The `.pth` file is a plain PyTorch state dict saved via `torch.save(model.state_dict(), ...)`.
9//! No wrapping key is used; `PthTensors::new(path, None)` loads it directly.
10//!
11//! | Model attribute (Python) | VarBuilder prefix |
12//! |--------------------------|-------------------|
13//! | `encoder.*`              | `encoder.*`       |
14//! | `transformer.*`          | `transformer.*`   |
15//! | `output_heads.*`         | `output_heads.*`  |
16//!
17//! The structure encoder uses a separate checkpoint (`esm3_structure_encoder_v0.pth`):
18//!
19//! | Python attribute        | VarBuilder prefix  |
20//! |-------------------------|--------------------|
21//! | `encoder.*` (blocks)    | `encoder.*`        |
22//! | `pre_vq_proj.weight`    | `pre_vq_proj.*`    |
23//! | `codebook.embeddings`   | `codebook.*`       |
24
25use crate::esm3::models::esm3::{ESM3, ESM3Config};
26use crate::esm3::models::vqvae::{StructureTokenEncoder, VqVaeConfig};
27use crate::esm3::tokenization::sequence::tokenize_sequence;
28use crate::plm_runner::PlmRunner;
29use anyhow::{Context, Result};
30use candle_core::pickle::PthTensors;
31use candle_core::{DType, Device, Tensor};
32use candle_nn::VarBuilder;
33use hf_hub::HFClientSync;
34
35const ESM3_DTYPE: DType = DType::F32;
36
37// ── ESM3Models enum ───────────────────────────────────────────────────────────
38
39/// Available ESM3 model variants.
40pub enum ESM3Models {
41    /// esm3-sm-open-v1 — 1.4B parameter open-access model (Cambrian Non-Commercial license).
42    SmOpen,
43}
44
45impl ESM3Models {
46    /// HuggingFace repo and `.pth` filename for this variant.
47    pub fn model_info(&self) -> (&'static str, &'static str, ESM3Config) {
48        match self {
49            Self::SmOpen => (
50                "EvolutionaryScale/esm3-sm-open-v1",
51                "esm3_sm_open_v1.pth",
52                ESM3Config::sm_open(),
53            ),
54        }
55    }
56}
57
58// ── ESM3Runner ────────────────────────────────────────────────────────────────
59
60/// Wraps a loaded ESM3 model for sequence embedding and multi-track inference.
61///
62/// ## HuggingFace access
63/// The ESM3 model is gated. Before calling `from_pretrained` the user must:
64/// 1. Accept the Cambrian Non-Commercial license at
65///    <https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1>
66/// 2. Run `huggingface-cli login` (or set `HF_TOKEN` env var)
67pub struct ESM3Runner {
68    model: ESM3,
69    device: Device,
70}
71
72impl ESM3Runner {
73    /// Download the ESM3 weights from HuggingFace and load the model.
74    ///
75    /// The `.pth` checkpoint is loaded directly via `candle_core::pickle::PthTensors`.
76    pub fn from_pretrained(variant: ESM3Models, device: Device) -> Result<Self> {
77        let (repo_id, filename, config) = variant.model_info();
78        let (owner, name) = repo_id.split_once('/').unwrap_or(("", repo_id));
79
80        let client = HFClientSync::new().context("failed to initialise HF client")?;
81        let weights_path = client
82            .model(owner, name)
83            .download_file()
84            .filename(filename)
85            .send()
86            .with_context(|| format!("failed to download {} from {}", filename, repo_id))?;
87
88        let pth = PthTensors::new(&weights_path, None)
89            .with_context(|| format!("failed to parse {}", weights_path.display()))?;
90        let vb = VarBuilder::from_backend(Box::new(pth), ESM3_DTYPE, device.clone());
91
92        let model = ESM3::load(vb, config)?;
93        Ok(Self { model, device })
94    }
95
96    /// Tokenize `sequence` and return per-residue embeddings `(1, L, d_model)`.
97    ///
98    /// Only the sequence track is provided; all other tracks are `None`.
99    /// Embeddings are the pre-norm transformer hidden states (raw activations).
100    pub fn embed_sequence(&self, sequence: &str) -> Result<Tensor> {
101        let token_ids = tokenize_sequence(sequence, true);
102        let tokens = Tensor::new(token_ids.as_slice(), &self.device)?.unsqueeze(0)?; // (1, L)
103
104        let output = self.model.forward(
105            Some(&tokens), // sequence_tokens
106            None,          // structure_tokens
107            None,          // ss8_tokens
108            None,          // sasa_tokens
109            None,          // function_tokens
110            None,          // residue_annotation_tokens
111            None,          // average_plddt
112            None,          // per_res_plddt
113            None,          // sequence_id
114            None,          // structure_coords
115            None,          // chain_id
116        )?;
117
118        output
119            .embeddings
120            .ok_or_else(|| anyhow::anyhow!("ESM3 forward() returned no embeddings"))
121    }
122}
123
124impl PlmRunner for ESM3Runner {
125    fn embed(&self, sequence: &str) -> Result<Tensor> {
126        self.embed_sequence(sequence)
127    }
128
129    fn model_name(&self) -> &str {
130        "esm3"
131    }
132}
133
134// ── StructureEncoderRunner ────────────────────────────────────────────────────
135
136/// Wraps the ESM3 structure token encoder for converting backbone coordinates to tokens.
137///
138/// Uses a separate checkpoint (`esm3_structure_encoder_v0.pth`) from the main model.
139pub struct StructureEncoderRunner {
140    encoder: StructureTokenEncoder,
141    device: Device,
142}
143
144impl StructureEncoderRunner {
145    /// Download and load the structure encoder from HuggingFace.
146    pub fn from_pretrained(device: Device) -> Result<Self> {
147        let repo_id = "EvolutionaryScale/esm3-sm-open-v1";
148        let filename = "esm3_structure_encoder_v0.pth";
149        let (owner, name) = repo_id.split_once('/').unwrap_or(("", repo_id));
150
151        let client = HFClientSync::new().context("failed to initialise HF client")?;
152        let weights_path = client
153            .model(owner, name)
154            .download_file()
155            .filename(filename)
156            .send()
157            .with_context(|| format!("failed to download {} from {}", filename, repo_id))?;
158
159        let pth = PthTensors::new(&weights_path, None)
160            .with_context(|| format!("failed to parse {}", weights_path.display()))?;
161        let vb = VarBuilder::from_backend(Box::new(pth), ESM3_DTYPE, device.clone());
162
163        let encoder = StructureTokenEncoder::load(vb, VqVaeConfig::default())?;
164        Ok(Self { encoder, device })
165    }
166
167    /// Encode backbone coordinates to structure tokens.
168    ///
169    /// - `coords`: `(B, L, 3, 3)` backbone `(N, CA, C)` atom positions.
170    ///
171    /// Returns `(B, L)` u32 structure token indices.
172    pub fn encode(&self, coords: &Tensor) -> Result<Tensor> {
173        self.encoder.encode(coords, None, None).map_err(Into::into)
174    }
175}
ferritin_plms/esm3/pretrained.rs

ferritin_plms/esm3/
pretrained.rs
