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::plm_runner::PlmRunner;
6use crate::types::PseudoProbability;
7use anyhow::{Error as E, Result, anyhow};
8use candle_core::{DType, Device, Tensor};
9use candle_nn::ops::softmax;
10use candle_nn::VarBuilder;
11use hf_hub::HFClientSync;
12use serde_json;
13use tokenizers::Tokenizer;
14
15const ESM2_DTYPE: DType = DType::F32;
16
17// ESM2 tokenizer: indices 4-23 are the 20 standard amino acids.
18const ESM2_STD_AA: [(usize, char); 20] = [
19    (4, 'L'), (5, 'A'), (6, 'G'), (7, 'V'), (8, 'S'),
20    (9, 'E'), (10, 'R'), (11, 'T'), (12, 'I'), (13, 'D'),
21    (14, 'P'), (15, 'K'), (16, 'Q'), (17, 'N'), (18, 'F'),
22    (19, 'Y'), (20, 'M'), (21, 'H'), (22, 'W'), (23, 'C'),
23];
24
25pub enum ESM2Models {
26    T6_8M,
27    T12_35M,
28    T30_150M,
29    T33_650M,
30    T36_3B,
31    T48_15B,
32}
33impl ESM2Models {
34    pub fn get_model_files(model: Self) -> (&'static str, &'static str, ESM2Config) {
35        match model {
36            Self::T6_8M => ("facebook/esm2_t6_8M_UR50D", "main", ESM2Config::t6_8m()),
37            Self::T12_35M => ("facebook/esm2_t12_35M_UR50D", "main", ESM2Config::t12_35m()),
38            Self::T30_150M => (
39                "facebook/esm2_t30_150M_UR50D",
40                "main",
41                ESM2Config::t30_150m(),
42            ),
43            Self::T33_650M => (
44                "facebook/esm2_t33_650M_UR50D",
45                "main",
46                ESM2Config::t33_650m(),
47            ),
48            Self::T36_3B => ("facebook/esm2_t36_3B_UR50D", "main", ESM2Config::t36_3b()),
49            Self::T48_15B => ("facebook/esm2_t48_15B_UR50D", "main", ESM2Config::t48_15b()),
50        }
51    }
52}
53
54pub struct ESM2Runner {
55    model: ESM2,
56    tokenizer: Tokenizer,
57}
58impl ESM2Runner {
59    /// Load model from HuggingFace hub, downloading config.json, tokenizer files, and weights.
60    pub fn load_model(modeltype: ESM2Models, device: Device) -> Result<ESM2Runner> {
61        let (model_id, revision, fallback_config) = ESM2Models::get_model_files(modeltype);
62        let (owner, name) = model_id.split_once('/').unwrap_or(("", model_id));
63        let client = HFClientSync::new()?;
64        let repo = client.model(owner, name);
65        // Try to load config from HF hub; fall back to hardcoded config if unavailable.
66        let config = match repo
67            .download_file()
68            .filename("config.json")
69            .revision(revision)
70            .send()
71        {
72            Ok(config_path) => {
73                let config_str = std::fs::read_to_string(config_path)?;
74                serde_json::from_str::<ESM2Config>(&config_str).unwrap_or(fallback_config)
75            }
76            Err(_) => fallback_config,
77        };
78        let weights_filename = repo
79            .download_file()
80            .filename("model.safetensors")
81            .revision(revision)
82            .send()?;
83        let vb = unsafe {
84            VarBuilder::from_mmaped_safetensors(&[weights_filename], ESM2_DTYPE, &device)?
85        };
86        let model = ESM2::load(vb, config)?;
87        let tokenizer = ESM2::load_tokenizer()?;
88        Ok(ESM2Runner { model, tokenizer })
89    }
90    pub fn run_forward(&self, prot_sequence: &str) -> Result<ESM2Output> {
91        let device = self.model.get_device();
92        let tokens = self
93            .tokenizer
94            .encode(prot_sequence.to_string(), false)
95            .map_err(E::msg)?
96            .get_ids()
97            .to_vec();
98        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
99        let encoded = self.model.forward(&token_ids, None)?;
100        Ok(encoded)
101    }
102    /// Predict residue-residue contact probabilities for a single protein sequence.
103    ///
104    /// Returns a `(seq_len, seq_len)` contact probability matrix (BOS/EOS stripped,
105    /// so dimensions equal the number of amino acids in `prot_sequence`).
106    pub fn predict_contacts(&self, prot_sequence: &str) -> Result<Tensor> {
107        let device = self.model.get_device();
108        let tokens = self
109            .tokenizer
110            .encode(prot_sequence.to_string(), false)
111            .map_err(E::msg)?
112            .get_ids()
113            .to_vec();
114        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
115        // squeeze batch dim: (1, L, L) → (L, L)
116        self.model
117            .predict_contacts(&token_ids, None)
118            .map_err(E::msg)?
119            .squeeze(0)
120            .map_err(E::msg)
121    }
122
123    pub fn decode_logits(&self, output: ESM2Output) -> Result<String> {
124        // Get the predicted token IDs by taking argmax along the vocabulary dimension
125        let predicted_token_ids = output.logits.argmax(2)?;
126        let predicted_token_ids = if predicted_token_ids.dims().len() > 1 {
127            predicted_token_ids.squeeze(0)?
128        } else {
129            predicted_token_ids
130        };
131        let token_ids: Vec<u32> = predicted_token_ids.to_vec1::<u32>()?;
132        let decoded_sequence = self
133            .tokenizer
134            .decode(&token_ids, true) // set skip_special_tokens to true
135            .map_err(|e| anyhow!("Failed to decode tokens: {}", e))?
136            .replace(" ", "");
137        Ok(decoded_sequence)
138    }
139
140    /// Run ESM2 and return per-residue pseudo-probabilities for the 20 standard amino acids.
141    ///
142    /// BOS and EOS tokens are stripped before softmax is applied. Only amino acid / position
143    /// pairs with probability > 0.01 are included in the result.
144    pub fn get_pseudo_probabilities(&self, prot_sequence: &str) -> Result<Vec<PseudoProbability>> {
145        let output = self.run_forward(prot_sequence)?;
146        // logits: (1, L+2, vocab_size) — strip BOS at 0 and EOS at -1
147        let seq_len = output.logits.dim(1)? - 2;
148        let logits = output.logits.narrow(1, 1, seq_len)?;
149        let probs = softmax(&logits, 2)?;
150        let probs = probs.squeeze(0)?; // (L, vocab)
151        let probs_data: Vec<Vec<f32>> = probs.to_vec2()?;
152
153        let mut result = Vec::new();
154        for (pos, pos_probs) in probs_data.iter().enumerate() {
155            for (vocab_idx, aa_char) in ESM2_STD_AA.iter() {
156                let prob = pos_probs[*vocab_idx];
157                if prob > 0.01 {
158                    result.push(PseudoProbability {
159                        position: pos,
160                        pseudo_prob: prob,
161                        amino_acid: *aa_char,
162                    });
163                }
164            }
165        }
166        Ok(result)
167    }
168}
169
170impl PlmRunner for ESM2Runner {
171    /// Run the ESM2 transformer and return per-residue embeddings (pre-LM-head).
172    ///
173    /// Shape: `(1, L, hidden_size)` where `L` includes BOS and EOS tokens.
174    fn embed(&self, sequence: &str) -> Result<Tensor> {
175        let device = self.model.get_device();
176        let tokens = self
177            .tokenizer
178            .encode(sequence.to_string(), false)
179            .map_err(E::msg)?
180            .get_ids()
181            .to_vec();
182        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
183        Ok(self.model.embed(&token_ids, None)?)
184    }
185
186    fn model_name(&self) -> &str {
187        "esm2"
188    }
189}
ferritin_plms/esm2/esm2_runner.rs

ferritin_plms/esm2/
esm2_runner.rs
