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//! Note the repo path: the checkpoints live under `data/weights/`, not at the
12//! repo root. Requesting the bare filename returns
13//! "Entry not found: esm3_sm_open_v1.pth" (ferritin-100.21).
14//!
15//! | Model attribute (Python) | VarBuilder prefix |
16//! |--------------------------|-------------------|
17//! | `encoder.*` | `encoder.*` |
18//! | `transformer.*` | `transformer.*` |
19//! | `output_heads.*` | `output_heads.*` |
20//!
21//! The structure encoder uses a separate checkpoint
22//! (`data/weights/esm3_structure_encoder_v0.pth`):
23//!
24//! | Python attribute | VarBuilder prefix |
25//! |-------------------------|--------------------|
26//! | `encoder.*` (blocks) | `encoder.*` |
27//! | `pre_vq_proj.weight` | `pre_vq_proj.*` |
28//! | `codebook.embeddings` | `codebook.*` |
29
30use crate::esm3::models::esm3::{ESM3, ESM3Config};
31use crate::esm3::models::vqvae::{StructureTokenEncoder, VqVaeConfig};
32use crate::esm3::tokenization::sequence::tokenize_sequence;
33use crate::esm3::utils::constants::SEQUENCE_PAD_TOKEN;
34use crate::loader::{LoadOptions, WeightSource};
35use crate::plm_runner::{
36 ModelMetadata, PlmRunner, SpecialTokenLayout, pad_token_batch, zero_padded_rows,
37};
38use crate::registry::{self, ModelCard};
39use anyhow::Result;
40use candle_core::{Device, Tensor};
41
42// ── ESM3Models enum ───────────────────────────────────────────────────────────
43
44/// Available ESM3 model variants.
45pub enum ESM3Models {
46 /// esm3-sm-open-v1 — 1.4B parameter open-access model (Cambrian Non-Commercial license).
47 SmOpen,
48}
49
50/// The structure encoder's registry id.
51pub const STRUCTURE_ENCODER_ID: &str = "esm3-structure-encoder-v0";
52
53impl ESM3Models {
54 /// This variant's registry id.
55 pub const fn registry_id(&self) -> &'static str {
56 match self {
57 Self::SmOpen => "esm3-sm-open-v1",
58 }
59 }
60
61 /// This variant's [`ModelCard`].
62 pub fn card(&self) -> &'static ModelCard {
63 registry::lookup(self.registry_id())
64 .expect("every ESM3Models variant must have a registry entry")
65 }
66
67 /// Weight source, `.pth` filename, and config for this variant.
68 ///
69 /// Repo and filename come from [`REGISTRY`][crate::registry::REGISTRY]
70 /// (ferritin-goh.1). Note the path: the checkpoints live under
71 /// `data/weights/`, not at the repo root (ferritin-100.21).
72 pub fn model_info(&self) -> (WeightSource, &'static str, ESM3Config) {
73 let card = self.card();
74 let config = match self {
75 Self::SmOpen => ESM3Config::sm_open(),
76 };
77 (card.source, card.file, config)
78 }
79}
80
81// ── ESM3Runner ────────────────────────────────────────────────────────────────
82
83/// Wraps a loaded ESM3 model for sequence embedding and multi-track inference.
84///
85/// ## HuggingFace access
86/// The ESM3 model is gated. Before calling `from_pretrained` the user must:
87/// 1. Accept the Cambrian Non-Commercial license at
88/// <https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1>
89/// 2. Run `huggingface-cli login` (or set `HF_TOKEN` env var)
90pub struct ESM3Runner {
91 model: ESM3,
92 device: Device,
93}
94
95impl ESM3Runner {
96 /// Download the ESM3 weights from HuggingFace and load the model.
97 ///
98 /// The `.pth` checkpoint is loaded directly via `candle_core::pickle::PthTensors`.
99 pub fn from_pretrained(variant: ESM3Models, device: Device) -> Result<Self> {
100 Self::from_pretrained_with(variant, &LoadOptions::new(device))
101 }
102
103 /// Load with an explicit device and dtype (ferritin-100.9).
104 pub fn from_pretrained_with(variant: ESM3Models, opts: &LoadOptions) -> Result<Self> {
105 let (source, filename, config) = variant.model_info();
106 let vb = source.var_builder(filename, opts)?;
107 let model = ESM3::load(vb, config)?;
108 Ok(Self {
109 model,
110 device: opts.device.clone(),
111 })
112 }
113
114 /// Tokenize `sequence` and return per-residue embeddings `(1, L, d_model)`.
115 ///
116 /// Only the sequence track is provided; all other tracks are `None`.
117 /// Embeddings are the pre-norm transformer hidden states (raw activations).
118 pub fn embed_sequence(&self, sequence: &str) -> Result<Tensor> {
119 let token_ids = tokenize_sequence(sequence, true);
120 let tokens = Tensor::new(token_ids.as_slice(), &self.device)?.unsqueeze(0)?; // (1, L)
121
122 let output = self.model.forward(
123 Some(&tokens), // sequence_tokens
124 None, // structure_tokens
125 None, // ss8_tokens
126 None, // sasa_tokens
127 None, // function_tokens
128 None, // residue_annotation_tokens
129 None, // average_plddt
130 None, // per_res_plddt
131 None, // sequence_id
132 None, // structure_coords
133 None, // chain_id
134 )?;
135
136 output
137 .embeddings
138 .ok_or_else(|| anyhow::anyhow!("ESM3 forward() returned no embeddings"))
139 }
140}
141
142impl PlmRunner for ESM3Runner {
143 fn embed(&self, sequence: &str) -> Result<Tensor> {
144 self.embed_sequence(sequence)
145 }
146
147 fn model_name(&self) -> &str {
148 "esm3"
149 }
150
151 /// `embed_sequence` calls `tokenize_sequence(.., true)`, which prepends
152 /// `SEQUENCE_BOS_TOKEN` and appends `SEQUENCE_EOS_TOKEN`.
153 fn special_tokens(&self) -> SpecialTokenLayout {
154 SpecialTokenLayout::BOS_EOS
155 }
156
157 fn metadata(&self) -> ModelMetadata {
158 let config = &self.model.config;
159 ModelMetadata {
160 d_model: config.d_model,
161 n_layers: config.n_layers,
162 vocab_size: config.d_sequence_vocab,
163 // Rotary positions: no hard architectural cap.
164 max_positions: None,
165 }
166 }
167
168 fn device(&self) -> &Device {
169 &self.device
170 }
171
172 /// Sequence-track logits `(1, L + 2, d_sequence_vocab)`.
173 ///
174 /// Only the sequence track is supplied; the other tracks are `None`.
175 fn logits(&self, sequence: &str) -> Result<Tensor> {
176 let token_ids = tokenize_sequence(sequence, true);
177 let tokens = Tensor::new(token_ids.as_slice(), &self.device)?.unsqueeze(0)?;
178 let output = self.model.forward(
179 Some(&tokens),
180 None,
181 None,
182 None,
183 None,
184 None,
185 None,
186 None,
187 None,
188 None,
189 None,
190 )?;
191 output
192 .sequence_logits
193 .ok_or_else(|| anyhow::anyhow!("ESM3 forward() returned no sequence logits"))
194 }
195
196 /// One batched forward pass over right-padded sequences (ferritin-100.12).
197 ///
198 /// ESM-3 reuses ESM-C's `MultiHeadAttention`, so `sequence_id` is the same
199 /// `where_cond` predicate (`1` keeps a key). Unlike ESM-C, `ESM3::forward`
200 /// has no fallback that derives one from the tokens: passing `None` — which
201 /// is what every single-sequence path here does — means no masking at all,
202 /// so a batch *must* supply it or the padding silently changes the real
203 /// residues' embeddings.
204 fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor> {
205 let rows: Vec<Vec<u32>> = sequences
206 .iter()
207 .map(|s| tokenize_sequence(s, true))
208 .collect();
209 let batch = pad_token_batch(&rows, SEQUENCE_PAD_TOKEN, &self.device)?;
210 let output = self.model.forward(
211 Some(&batch.ids), // sequence_tokens
212 None, // structure_tokens
213 None, // ss8_tokens
214 None, // sasa_tokens
215 None, // function_tokens
216 None, // residue_annotation_tokens
217 None, // average_plddt
218 None, // per_res_plddt
219 Some(&batch.mask), // sequence_id
220 None, // structure_coords
221 None, // chain_id
222 )?;
223 let embeddings = output
224 .embeddings
225 .ok_or_else(|| anyhow::anyhow!("ESM3 forward() returned no embeddings"))?;
226 zero_padded_rows(&embeddings, &batch.mask)
227 }
228}
229
230// ── StructureEncoderRunner ────────────────────────────────────────────────────
231
232/// Wraps the ESM3 structure token encoder for converting backbone coordinates to tokens.
233///
234/// Uses a separate checkpoint (`data/weights/esm3_structure_encoder_v0.pth`)
235/// from the main model.
236pub struct StructureEncoderRunner {
237 encoder: StructureTokenEncoder,
238 #[allow(dead_code)]
239 device: Device,
240}
241
242impl StructureEncoderRunner {
243 /// Load the VQ-VAE structure encoder (ferritin-100.22).
244 ///
245 /// This used to refuse: the port targeted a network the released
246 /// checkpoint does not contain. All of those differences are now modelled
247 /// — the stack roots at `transformer`, its blocks are geometric-only with
248 /// no `attn` and no trailing `norm`, `pre_vq_proj` carries its bias, and
249 /// the relative position embedding is loaded and used as the encoder's
250 /// initial hidden state.
251 pub fn from_pretrained(device: Device) -> Result<Self> {
252 Self::from_pretrained_with(&LoadOptions::new(device))
253 }
254
255 /// Load with an explicit device and dtype.
256 pub fn from_pretrained_with(opts: &LoadOptions) -> Result<Self> {
257 let card = registry::lookup(STRUCTURE_ENCODER_ID)
258 .expect("the structure encoder must have a registry entry");
259 let vb = card.source.var_builder(card.file, opts)?;
260 let encoder = StructureTokenEncoder::load(vb, VqVaeConfig::default())?;
261 Ok(Self {
262 encoder,
263 device: opts.device.clone(),
264 })
265 }
266
267 /// Encode backbone coordinates to structure tokens.
268 ///
269 /// - `coords`: `(B, L, 3, 3)` backbone `(N, CA, C)` atom positions.
270 ///
271 /// Returns `(B, L)` u32 structure token indices.
272 pub fn encode(&self, coords: &Tensor) -> Result<Tensor> {
273 self.encoder.encode(coords, None).map_err(Into::into)
274 }
275}