ferritin_plms/esm2/
esm2_runner.rs

1//! ESM2 Runner
2//!
3//! Class for loading and running the ESM2 models
4use super::esm2::{ESM2, ESM2Config, ESM2Output};
5use crate::esm2::saprot_tokenizer::SaProtTokenizer;
6use crate::loader::{LoadOptions, WeightSource};
7use crate::plm_runner::{
8    ModelMetadata, PlmRunner, SpecialTokenLayout, pad_token_batch, zero_padded_rows,
9};
10use crate::registry::{self, ModelCard, TokenizerSpec, VocabAlphabet};
11use crate::types::PseudoProbability;
12use anyhow::{Error as E, Result, anyhow, bail};
13use candle_core::{DType, Device, Tensor};
14use candle_nn::ops::softmax;
15use serde_json;
16use tokenizers::Tokenizer;
17
18// ESM2 tokenizer: indices 4-23 are the 20 standard amino acids.
19const ESM2_STD_AA: [(usize, char); 20] = [
20    (4, 'L'),
21    (5, 'A'),
22    (6, 'G'),
23    (7, 'V'),
24    (8, 'S'),
25    (9, 'E'),
26    (10, 'R'),
27    (11, 'T'),
28    (12, 'I'),
29    (13, 'D'),
30    (14, 'P'),
31    (15, 'K'),
32    (16, 'Q'),
33    (17, 'N'),
34    (18, 'F'),
35    (19, 'Y'),
36    (20, 'M'),
37    (21, 'H'),
38    (22, 'W'),
39    (23, 'C'),
40];
41
42pub enum ESM2Models {
43    T6_8M,
44    T12_35M,
45    T30_150M,
46    T33_650M,
47    T36_3B,
48    T48_15B,
49    /// ESM-1v UR90S ensemble member 1 (ferritin-goh.4).
50    Esm1v1,
51    /// ESM-1v UR90S ensemble member 2 (ferritin-goh.4).
52    Esm1v2,
53    /// ESM-1v UR90S ensemble member 3 (ferritin-goh.4).
54    Esm1v3,
55    /// ESM-1v UR90S ensemble member 4 (ferritin-goh.4).
56    Esm1v4,
57    /// ESM-1v UR90S ensemble member 5 (ferritin-goh.4).
58    Esm1v5,
59    /// ESM-1b — learned absolute positions.
60    Esm1b,
61    /// SaProt 35M — ESM-2 architecture, structure-aware alphabet (ferritin-goh.3).
62    SaProt35M,
63    /// SaProt 650M — the variant most people use.
64    SaProt650M,
65    /// FastESM2 650M — ESM-2 650M republished as safetensors (ferritin-goh.12).
66    FastEsm2_650,
67    /// PepMLM 650M — ESM-2 650M fine-tuned for peptide binder design.
68    PepMlm650M,
69    /// DPLM 650M — a discrete diffusion model on ESM-2's architecture.
70    ///
71    /// Its hidden states are ordinary embeddings; its `lm_head` denoises
72    /// rather than scoring masked tokens, so do not read its logits as
73    /// masked-LM scores.
74    Dplm650M,
75}
76impl ESM2Models {
77    /// This variant's registry id.
78    pub const fn registry_id(&self) -> &'static str {
79        match self {
80            Self::T6_8M => "esm2-t6-8m",
81            Self::T12_35M => "esm2-t12-35m",
82            Self::T30_150M => "esm2-t30-150m",
83            Self::T33_650M => "esm2-t33-650m",
84            Self::T36_3B => "esm2-t36-3b",
85            Self::T48_15B => "esm2-t48-15b",
86            Self::Esm1v1 => "esm1v-t33-650m-ur90s-1",
87            Self::Esm1v2 => "esm1v-t33-650m-ur90s-2",
88            Self::Esm1v3 => "esm1v-t33-650m-ur90s-3",
89            Self::Esm1v4 => "esm1v-t33-650m-ur90s-4",
90            Self::Esm1v5 => "esm1v-t33-650m-ur90s-5",
91            Self::Esm1b => "esm1b-t33-650m-ur50s",
92            Self::SaProt35M => "saprot-35m-af2",
93            Self::SaProt650M => "saprot-650m-af2",
94            Self::FastEsm2_650 => "fastesm2-650",
95            Self::PepMlm650M => "pepmlm-650m",
96            Self::Dplm650M => "dplm-650m",
97        }
98    }
99
100    /// This variant's [`ModelCard`].
101    pub fn card(&self) -> &'static ModelCard {
102        registry::lookup(self.registry_id())
103            .expect("every ESM2Models variant must have a registry entry")
104    }
105
106    /// Where this variant's weights live, plus its built-in fallback config.
107    ///
108    /// The weight source comes from [`REGISTRY`][crate::registry::REGISTRY]
109    /// rather than being repeated here (ferritin-goh.1). The config stays a
110    /// separate hardcoded fallback because it is only used when the hub's
111    /// config.json is unreachable or unparseable.
112    pub fn model_info(&self) -> (WeightSource, ESM2Config) {
113        let config = match self {
114            Self::T6_8M => ESM2Config::t6_8m(),
115            Self::T12_35M => ESM2Config::t12_35m(),
116            Self::T30_150M => ESM2Config::t30_150m(),
117            Self::T33_650M => ESM2Config::t33_650m(),
118            Self::T36_3B => ESM2Config::t36_3b(),
119            Self::T48_15B => ESM2Config::t48_15b(),
120            Self::Esm1v1 => ESM2Config::esm1v_t33_650m(),
121            Self::Esm1v2 => ESM2Config::esm1v_t33_650m(),
122            Self::Esm1v3 => ESM2Config::esm1v_t33_650m(),
123            Self::Esm1v4 => ESM2Config::esm1v_t33_650m(),
124            Self::Esm1v5 => ESM2Config::esm1v_t33_650m(),
125            Self::Esm1b => ESM2Config::esm1b_t33_650m(),
126            Self::SaProt35M => ESM2Config::saprot_35m(),
127            Self::SaProt650M => ESM2Config::saprot_650m(),
128            // All three are ESM-2 650M's shape exactly — same hidden size,
129            // depth, head count, vocabulary and rotary positions — so they
130            // share its fallback config rather than duplicating it. This is
131            // only reached when the hub's config.json is unreachable; the
132            // published files were checked against these values (goh.12).
133            Self::FastEsm2_650 | Self::PepMlm650M | Self::Dplm650M => ESM2Config::t33_650m(),
134        };
135        (self.card().source, config)
136    }
137
138    /// Renamed to [`model_info`][Self::model_info] (ferritin-100.8).
139    #[deprecated(since = "0.4.0", note = "renamed to `model_info`, which takes &self")]
140    pub fn get_model_files(model: Self) -> (WeightSource, ESM2Config) {
141        model.model_info()
142    }
143}
144
145/// Which tokenizer an ESM-2-family model uses.
146///
147/// The architecture is shared, the alphabet is not: stock ESM-2 has a 33-token
148/// `tokenizer.json` compiled into the crate, while SaProt ships a bare 446-line
149/// `vocab.txt` and reads two characters per residue (ferritin-goh.3).
150enum SequenceTokenizer {
151    /// The embedded HuggingFace tokenizer, one byte per residue.
152    Esm(Box<Tokenizer>),
153    /// SaProt's (amino acid, 3Di) product alphabet, two bytes per residue.
154    SaProt(SaProtTokenizer),
155}
156
157impl SequenceTokenizer {
158    fn token_to_id(&self, token: &str) -> Option<u32> {
159        match self {
160            Self::Esm(t) => t.token_to_id(token),
161            Self::SaProt(t) => t.token_to_id(token),
162        }
163    }
164
165    fn encode(&self, sequence: &str) -> Result<Vec<u32>> {
166        Ok(match self {
167            Self::Esm(t) => t
168                .encode(sequence.to_string(), false)
169                .map_err(E::msg)?
170                .get_ids()
171                .to_vec(),
172            Self::SaProt(t) => t.encode(sequence),
173        })
174    }
175
176    fn residue_count(&self, sequence: &str) -> usize {
177        match self {
178            Self::Esm(_) => sequence.len(),
179            Self::SaProt(t) => t.residue_count(sequence),
180        }
181    }
182
183    /// Decode ids back to a sequence, dropping special tokens.
184    fn decode(&self, ids: &[u32]) -> Result<String> {
185        Ok(match self {
186            Self::Esm(t) => t
187                .decode(ids, true)
188                .map_err(|e| anyhow!("Failed to decode tokens: {e}"))?
189                .replace(' ', ""),
190            Self::SaProt(t) => t.decode(ids),
191        })
192    }
193}
194
195pub struct ESM2Runner {
196    model: ESM2,
197    tokenizer: SequenceTokenizer,
198    /// Retained so `PlmRunner::metadata` can report dimensions without
199    /// hardcoding them per variant.
200    config: ESM2Config,
201}
202impl ESM2Runner {
203    /// Load model from HuggingFace hub at F32, downloading config.json,
204    /// tokenizer files, and weights.
205    ///
206    /// Use [`from_pretrained_with`][Self::from_pretrained_with] to load at F16
207    /// or BF16 — `T36_3B` and `T48_15B` are not loadable at F32 on most
208    /// machines (~12 GB and ~60 GB respectively).
209    pub fn from_pretrained(modeltype: ESM2Models, device: Device) -> Result<ESM2Runner> {
210        Self::from_pretrained_with(modeltype, &LoadOptions::new(device))
211    }
212
213    /// Renamed to [`from_pretrained`][Self::from_pretrained] (ferritin-100.8).
214    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained`")]
215    pub fn load_model(modeltype: ESM2Models, device: Device) -> Result<ESM2Runner> {
216        Self::from_pretrained(modeltype, device)
217    }
218
219    /// Renamed to [`from_pretrained_with`][Self::from_pretrained_with] (ferritin-100.8).
220    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained_with`")]
221    pub fn load_model_with(modeltype: ESM2Models, opts: &LoadOptions) -> Result<ESM2Runner> {
222        Self::from_pretrained_with(modeltype, opts)
223    }
224
225    /// Load model with an explicit device and dtype.
226    ///
227    /// Reduced precision halves the memory footprint but changes the numerics;
228    /// see `tests/test_plm_dtype_parity.rs` for the tolerance each dtype
229    /// actually achieves against the Python reference (ferritin-100.9).
230    pub fn from_pretrained_with(modeltype: ESM2Models, opts: &LoadOptions) -> Result<ESM2Runner> {
231        let (source, fallback_config) = modeltype.model_info();
232        let config = Self::resolve_config(&source, fallback_config)?;
233        let card = modeltype.card();
234        let vb = source.var_builder(card.file, opts)?;
235        let model = ESM2::load(vb, config.clone())?;
236
237        // The alphabet follows the card, not the family: SaProt reuses this
238        // architecture with a 446-token vocab.txt (ferritin-goh.3). The match
239        // on `VocabAlphabet` is exhaustive on purpose — a new alphabet must
240        // say what this runner should do with it rather than falling through
241        // to ESM-2's own tokenizer (ferritin-goh.11).
242        let tokenizer = match card.tokenizer {
243            TokenizerSpec::HfVocabTxt(alphabet) => match alphabet {
244                VocabAlphabet::SaProtPairs => {
245                    let path = source.fetch("vocab.txt")?;
246                    let contents = std::fs::read_to_string(path)?;
247                    SequenceTokenizer::SaProt(SaProtTokenizer::from_vocab_txt(&contents)?)
248                }
249                VocabAlphabet::SingleResidue => {
250                    bail!(
251                        "{}: a single-residue vocab.txt alphabet has no tokenizer in the \
252                         ESM-2 runner; only SaProt's (amino acid, 3Di) pair alphabet is wired \
253                         here (ferritin-goh.11)",
254                        card.id
255                    )
256                }
257            },
258            _ => SequenceTokenizer::Esm(Box::new(ESM2::load_tokenizer()?)),
259        };
260        Ok(ESM2Runner {
261            model,
262            tokenizer,
263            config,
264        })
265    }
266    /// Load `config.json` from the hub, falling back to the built-in config.
267    ///
268    /// The fallback is **loud**. It used to be
269    /// `serde_json::from_str(..).unwrap_or(fallback_config)`, which discarded
270    /// the parse error without a word: a config the struct could not represent
271    /// silently became a different model's config. For a checkpoint like
272    /// `SaProt_35M_AF2` — 446-token vocabulary against ESM-2's 33 — the
273    /// VarBuilder would probably have errored on the embedding shape, but any
274    /// divergence that happened to be shape-compatible would have loaded
275    /// cleanly and produced wrong numbers (ferritin-goh.9).
276    fn resolve_config(source: &WeightSource, fallback: ESM2Config) -> Result<ESM2Config> {
277        let Some(path) = source.fetch_optional("config.json") else {
278            eprintln!(
279                "warning: {}: could not download config.json; using the built-in \
280                 config for this variant. Verify it matches the checkpoint.",
281                source.repo_id
282            );
283            return Ok(fallback);
284        };
285
286        let config_str = std::fs::read_to_string(path)?;
287        match serde_json::from_str::<ESM2Config>(&config_str) {
288            Ok(config) => Ok(config),
289            Err(e) => {
290                eprintln!(
291                    "warning: {}: config.json did not parse as an ESM2Config ({e}); \
292                     using the built-in config instead. If this checkpoint is not an \
293                     ESM-2 variant, the built-in config is probably wrong for it.",
294                    source.repo_id
295                );
296                Ok(fallback)
297            }
298        }
299    }
300
301    /// Encode a sequence to token ids **with** ESM-2's BOS (`<cls>`) and EOS
302    /// (`<eos>`) special tokens.
303    ///
304    /// The bundled `tokenizer.json` has no post-processor, so `encode(.., true)`
305    /// would still NOT add them. ESM-2 is trained with BOS/EOS wrapping the
306    /// sequence, and every consumer here (`predict_contacts`,
307    /// `get_pseudo_probabilities`, `embed`, `decode_logits`) assumes an `L + 2`
308    /// row output and strips those two rows. Running the model without them
309    /// feeds a different context and its logits diverge from the HF reference,
310    /// so they must be added explicitly.
311    fn encode_with_special(&self, sequence: &str) -> Result<Vec<u32>> {
312        let bos = self
313            .tokenizer
314            .token_to_id("<cls>")
315            .ok_or_else(|| anyhow!("ESM2 tokenizer missing <cls> token"))?;
316        let eos = self
317            .tokenizer
318            .token_to_id("<eos>")
319            .ok_or_else(|| anyhow!("ESM2 tokenizer missing <eos> token"))?;
320        let inner = self.tokenizer.encode(sequence)?;
321        let mut ids = Vec::with_capacity(inner.len() + 2);
322        ids.push(bos);
323        ids.extend_from_slice(&inner);
324        ids.push(eos);
325        Ok(ids)
326    }
327    pub fn run_forward(&self, prot_sequence: &str) -> Result<ESM2Output> {
328        let device = self.model.get_device();
329        let tokens = self.encode_with_special(prot_sequence)?;
330        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
331        let encoded = self.model.forward(&token_ids, None)?;
332        Ok(encoded)
333    }
334    /// Predict residue-residue contact probabilities for a single protein sequence.
335    ///
336    /// Returns a `(seq_len, seq_len)` contact probability matrix (BOS/EOS stripped,
337    /// so dimensions equal the number of amino acids in `prot_sequence`).
338    pub fn predict_contacts(&self, prot_sequence: &str) -> Result<Tensor> {
339        let device = self.model.get_device();
340        let tokens = self.encode_with_special(prot_sequence)?;
341        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
342        // squeeze batch dim: (1, L, L) → (L, L)
343        self.model
344            .predict_contacts(&token_ids, None)
345            .map_err(E::msg)?
346            .squeeze(0)
347            .map_err(E::msg)
348    }
349
350    pub fn decode_logits(&self, output: ESM2Output) -> Result<String> {
351        // Get the predicted token IDs by taking argmax along the vocabulary dimension
352        let predicted_token_ids = output.logits.argmax(2)?;
353        let predicted_token_ids = if predicted_token_ids.dims().len() > 1 {
354            predicted_token_ids.squeeze(0)?
355        } else {
356            predicted_token_ids
357        };
358        let token_ids: Vec<u32> = predicted_token_ids.to_vec1::<u32>()?;
359        self.tokenizer.decode(&token_ids)
360    }
361
362    /// Run ESM2 and return per-residue pseudo-probabilities for the 20 standard amino acids.
363    ///
364    /// BOS and EOS tokens are stripped before softmax is applied. Only amino acid / position
365    /// pairs with probability > 0.01 are included in the result.
366    pub fn get_pseudo_probabilities(&self, prot_sequence: &str) -> Result<Vec<PseudoProbability>> {
367        let output = self.run_forward(prot_sequence)?;
368        // logits: (1, L+2, vocab_size) — strip BOS at 0 and EOS at -1
369        let seq_len = output.logits.dim(1)? - 2;
370        let logits = output.logits.narrow(1, 1, seq_len)?;
371        let probs = softmax(&logits, 2)?;
372        let probs = probs.squeeze(0)?; // (L, vocab)
373        let probs_data: Vec<Vec<f32>> = probs.to_vec2()?;
374
375        let mut result = Vec::new();
376        for (pos, pos_probs) in probs_data.iter().enumerate() {
377            for (vocab_idx, aa_char) in ESM2_STD_AA.iter() {
378                let prob = pos_probs[*vocab_idx];
379                if prob > 0.01 {
380                    result.push(PseudoProbability {
381                        position: pos,
382                        pseudo_prob: prob,
383                        amino_acid: *aa_char,
384                    });
385                }
386            }
387        }
388        Ok(result)
389    }
390}
391
392impl PlmRunner for ESM2Runner {
393    /// Run the ESM2 transformer and return per-residue embeddings (pre-LM-head).
394    ///
395    /// Shape: `(1, L, hidden_size)` where `L` includes BOS and EOS tokens.
396    fn embed(&self, sequence: &str) -> Result<Tensor> {
397        let device = self.model.get_device();
398        let tokens = self.encode_with_special(sequence)?;
399        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
400        Ok(self.model.embed(&token_ids, None)?)
401    }
402
403    fn model_name(&self) -> &str {
404        "esm2"
405    }
406
407    /// ESM-2 wraps sequences in `<cls>` and `<eos>`; `encode_with_special`
408    /// adds them because the bundled `tokenizer.json` has no post-processor.
409    fn special_tokens(&self) -> SpecialTokenLayout {
410        SpecialTokenLayout::BOS_EOS
411    }
412
413    fn metadata(&self) -> ModelMetadata {
414        ModelMetadata {
415            d_model: self.config.hidden_size,
416            n_layers: self.config.num_hidden_layers as usize,
417            vocab_size: self.config.vocab_size as usize,
418            max_positions: Some(self.config.max_position_embeddings as usize),
419        }
420    }
421
422    fn device(&self) -> &Device {
423        self.model.get_device()
424    }
425
426    /// SaProt reads two characters per residue; stock ESM-2 reads one.
427    fn residue_count(&self, sequence: &str) -> usize {
428        self.tokenizer.residue_count(sequence)
429    }
430
431    /// Masked-LM logits `(1, L + 2, vocab_size)`.
432    fn logits(&self, sequence: &str) -> Result<Tensor> {
433        Ok(self.run_forward(sequence)?.logits)
434    }
435
436    /// One batched forward pass over right-padded sequences (ferritin-100.12).
437    ///
438    /// Padding is `config.pad_token_id`, which matters twice over: ESM-1b and
439    /// ESM-1v number their absolute positions by counting non-pad tokens, and
440    /// the `(batch, seq_len)` mask handed to `ESM2::embed` both keeps padded
441    /// keys out of self-attention and supplies the per-row sequence length that
442    /// token-dropout compensation divides by.
443    fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor> {
444        let device = self.model.get_device();
445        let rows = sequences
446            .iter()
447            .map(|s| self.encode_with_special(s))
448            .collect::<Result<Vec<_>>>()?;
449        let batch = pad_token_batch(&rows, self.config.pad_token_id as u32, device)?;
450        // F32 rather than the model dtype: the mask is summed to get per-row
451        // lengths, and F16 cannot count past 65504 exactly.
452        let mask = batch.mask_as(DType::F32)?;
453        let embeddings = self.model.embed(&batch.ids, Some(&mask))?;
454        zero_padded_rows(&embeddings, &batch.mask)
455    }
456}
ferritin_plms/esm2/esm2_runner.rs

ferritin_plms/esm2/
esm2_runner.rs
