ferritin_plms/amplify/
amplify_runner.rs

1//! Amplify RUnner
2//!
3//! Class for loading and running the AMPLIFY models
4
5use super::super::types::{ContactMap, PseudoProbability};
6use super::amplify::{AMPLIFY, AmplifyOutput};
7use super::config::AMPLIFYConfig;
8use crate::loader::{LoadOptions, WeightSource};
9use crate::plm_runner::{
10    ModelMetadata, PlmRunner, SpecialTokenLayout, additive_padding_mask, pad_token_batch,
11    zero_padded_rows,
12};
13use crate::registry::{self, ModelCard};
14use anyhow::{Error as E, Result, anyhow};
15use candle_core::{D, Device, Tensor};
16use candle_nn::ops;
17use tokenizers::Tokenizer;
18
19pub enum AmplifyModels {
20    AMP120M,
21    AMP350M,
22}
23impl AmplifyModels {
24    /// This variant's registry id.
25    pub const fn registry_id(&self) -> &'static str {
26        match self {
27            AmplifyModels::AMP120M => "amplify-120m",
28            AmplifyModels::AMP350M => "amplify-350m",
29        }
30    }
31
32    /// This variant's [`ModelCard`].
33    pub fn card(&self) -> &'static ModelCard {
34        registry::lookup(self.registry_id())
35            .expect("every AmplifyModels variant must have a registry entry")
36    }
37
38    /// Where this variant's weights live.
39    ///
40    /// Delegates to [`REGISTRY`][crate::registry::REGISTRY] rather than
41    /// repeating the repo string (ferritin-goh.1).
42    pub fn model_info(&self) -> WeightSource {
43        self.card().source
44    }
45
46    /// Renamed to [`model_info`][Self::model_info] (ferritin-100.8).
47    #[deprecated(since = "0.4.0", note = "renamed to `model_info`, which takes &self")]
48    pub fn get_model_files(model: Self) -> WeightSource {
49        model.model_info()
50    }
51}
52
53pub struct AmplifyRunner {
54    model: AMPLIFY,
55    tokenizer: Tokenizer,
56}
57impl AmplifyRunner {
58    /// Load model from HuggingFace hub at F32.
59    ///
60    /// Use [`from_pretrained_with`][Self::from_pretrained_with] for F16 or BF16.
61    pub fn from_pretrained(modeltype: AmplifyModels, device: Device) -> Result<AmplifyRunner> {
62        Self::from_pretrained_with(modeltype, &LoadOptions::new(device))
63    }
64
65    /// Renamed to [`from_pretrained`][Self::from_pretrained] (ferritin-100.8).
66    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained`")]
67    pub fn load_model(modeltype: AmplifyModels, device: Device) -> Result<AmplifyRunner> {
68        Self::from_pretrained(modeltype, device)
69    }
70
71    /// Renamed to [`from_pretrained_with`][Self::from_pretrained_with] (ferritin-100.8).
72    #[deprecated(since = "0.4.0", note = "renamed to `from_pretrained_with`")]
73    pub fn load_model_with(modeltype: AmplifyModels, opts: &LoadOptions) -> Result<AmplifyRunner> {
74        Self::from_pretrained_with(modeltype, opts)
75    }
76
77    /// Load model with an explicit device and dtype.
78    ///
79    /// Reduced precision halves the memory footprint but changes the numerics;
80    /// see `tests/test_plm_dtype_parity.rs` for the achievable tolerance
81    /// (ferritin-100.9).
82    pub fn from_pretrained_with(
83        modeltype: AmplifyModels,
84        opts: &LoadOptions,
85    ) -> Result<AmplifyRunner> {
86        let source = modeltype.model_info();
87        let config_filename = source.fetch("config.json")?;
88        let tokenizer_filename = source.fetch("tokenizer.json")?;
89        let config_str = std::fs::read_to_string(config_filename)?;
90        let config_str = config_str
91            .replace("SwiGLU", "swiglu")
92            .replace("Swiglu", "swiglu");
93        let config: AMPLIFYConfig = serde_json::from_str(&config_str)?;
94        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
95        let vb = source.var_builder("model.safetensors", opts)?;
96        let model = AMPLIFY::load(vb, &config)?;
97        Ok(AmplifyRunner { model, tokenizer })
98    }
99    pub fn run_forward(&self, prot_sequence: &str) -> Result<AmplifyOutput> {
100        let device = self.model.get_device();
101        let tokens = self
102            .tokenizer
103            .encode(prot_sequence.to_string(), true)
104            .map_err(E::msg)?
105            .get_ids()
106            .to_vec();
107        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
108        let encoded = self.model.forward(&token_ids, None, false, true)?;
109        Ok(encoded)
110    }
111    pub fn get_best_prediction(
112        &self,
113        prot_sequence: &str,
114    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
115        let model_output: AmplifyOutput = self.run_forward(prot_sequence)?;
116        let predictions = model_output.logits.argmax(D::Minus1)?;
117        let indices: Vec<u32> = predictions.to_vec2()?[0].to_vec();
118        let decoded = self.tokenizer.decode(indices.as_slice(), true)?;
119        let decoded = decoded.replace(" ", "");
120        Ok(decoded)
121    }
122    /// Per-residue pseudo-probabilities over the full AMPLIFY vocabulary.
123    ///
124    /// `position` is the residue index: BOS and EOS rows are stripped first,
125    /// so there is exactly one position per input residue. Previously these
126    /// rows were included, which both over-counted positions by two and
127    /// shifted every label by one — position 0 was BOS, not the first residue
128    /// (ferritin-100.18).
129    ///
130    /// Unlike [`crate::ESM2Runner::get_pseudo_probabilities`], every amino acid is
131    /// returned rather than only those above a probability threshold.
132    pub fn get_pseudo_probabilities(&self, prot_sequence: &str) -> Result<Vec<PseudoProbability>> {
133        let model_output: AmplifyOutput = self.run_forward(prot_sequence)?;
134        let layout = <Self as PlmRunner>::special_tokens(self);
135        let residue_rows = model_output.logits.dim(1)?.saturating_sub(layout.total());
136        let predictions = model_output
137            .logits
138            .narrow(1, layout.leading, residue_rows)?;
139        let outputs = self.extract_logits(&predictions)?;
140        Ok(outputs)
141    }
142    pub fn get_contact_map(&self, prot_sequence: &str) -> Result<Vec<ContactMap>> {
143        let model_output: AmplifyOutput = self.run_forward(prot_sequence)?;
144        let contact_map_tensor = model_output.get_contact_map()?.ok_or_else(|| {
145            anyhow!("AMPLIFY forward() returned no attentions for the contact map")
146        })?;
147        let averaged = contact_map_tensor.max_keepdim(D::Minus1)?;
148        let (position1, position2, val) = averaged.dims3()?;
149        let data = averaged.to_vec3::<f32>()?;
150
151        // Per-position residue labels. The contact map is BOS/EOS-stripped, so
152        // position `p` is the p-th residue; decode the actual token ids the
153        // model saw (encode adds BOS/EOS, so the residues are ids[1..len-1])
154        // rather than decoding the position index as if it were a token id.
155        let encoded = self
156            .tokenizer
157            .encode(prot_sequence.to_string(), true)
158            .map_err(E::msg)?;
159        let ids = encoded.get_ids();
160        let residue_ids = if ids.len() >= 2 {
161            &ids[1..ids.len() - 1]
162        } else {
163            ids
164        };
165        let label_at = |p: usize| -> char {
166            residue_ids
167                .get(p)
168                .and_then(|&id| self.tokenizer.decode(&[id], true).ok())
169                .and_then(|s| s.chars().next())
170                .unwrap_or('?')
171        };
172
173        let mut contacts = Vec::new();
174        #[allow(clippy::needless_range_loop)] // i/j/k index a 3-D contact tensor
175        for i in 0..position1 {
176            for j in 0..position2 {
177                for k in 0..val {
178                    contacts.push(ContactMap {
179                        position_1: i,
180                        amino_acid_1: label_at(i),
181                        position_2: j,
182                        amino_acid_2: label_at(j),
183                        contact_estimate: data[i][j][k],
184                        layer: 1,
185                    });
186                }
187            }
188        }
189        Ok(contacts)
190    }
191    /// Softmax `tensor` and flatten it into per-(position, amino acid) rows.
192    ///
193    /// `tensor` must already have its special-token rows stripped, so
194    /// `seq_pos` is a residue index.
195    fn extract_logits(&self, tensor: &Tensor) -> Result<Vec<PseudoProbability>> {
196        let tensor = ops::softmax(tensor, D::Minus1)?;
197        let data = tensor.to_vec3::<f32>()?;
198        let (_, seq_len, vocab_size) = tensor.dims3()?;
199        let mut logit_positions = Vec::with_capacity(seq_len * vocab_size);
200        #[allow(clippy::needless_range_loop)] // indexes a 3-D logits tensor
201        for seq_pos in 0..seq_len {
202            for vocab_idx in 0..vocab_size {
203                let score = data[0][seq_pos][vocab_idx];
204                let amino_acid_char = self
205                    .tokenizer
206                    .decode(&[vocab_idx as u32], false)
207                    .map_err(|e| anyhow!("Failed to decode: {}", e))?
208                    .chars()
209                    .next()
210                    .ok_or_else(|| anyhow!("Empty decoded string"))?;
211                logit_positions.push(PseudoProbability {
212                    position: seq_pos,
213                    amino_acid: amino_acid_char,
214                    pseudo_prob: score,
215                });
216            }
217        }
218        Ok(logit_positions)
219    }
220}
221
222impl PlmRunner for AmplifyRunner {
223    /// Run AMPLIFY and return the last-layer hidden states as per-residue embeddings.
224    ///
225    /// Shape: `(1, L, hidden_size)` where `L` includes BOS and EOS tokens.
226    fn embed(&self, sequence: &str) -> Result<Tensor> {
227        let device = self.model.get_device();
228        let tokens = self
229            .tokenizer
230            .encode(sequence.to_string(), true)
231            .map_err(E::msg)?
232            .get_ids()
233            .to_vec();
234        let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
235        let output = self.model.forward(&token_ids, None, true, false)?;
236        let mut hidden_states = output
237            .hidden_states
238            .ok_or_else(|| anyhow!("AMPLIFY forward() returned no hidden states"))?;
239        hidden_states
240            .pop()
241            .ok_or_else(|| anyhow!("AMPLIFY returned empty hidden states list"))
242    }
243
244    fn model_name(&self) -> &str {
245        "amplify"
246    }
247
248    /// AMPLIFY's `tokenizer.json` has a `TemplateProcessing` post-processor
249    /// that adds `<bos>` (3) and `<eos>` (4), so `encode(.., true)` wraps the
250    /// sequence.
251    fn special_tokens(&self) -> SpecialTokenLayout {
252        SpecialTokenLayout::BOS_EOS
253    }
254
255    fn metadata(&self) -> ModelMetadata {
256        let config = self.model.config();
257        ModelMetadata {
258            d_model: config.hidden_size,
259            n_layers: config.num_hidden_layers,
260            vocab_size: config.vocab_size,
261            max_positions: Some(config.max_length),
262        }
263    }
264
265    fn device(&self) -> &Device {
266        self.model.get_device()
267    }
268
269    /// Masked-LM logits `(1, L + 2, vocab_size)`.
270    fn logits(&self, sequence: &str) -> Result<Tensor> {
271        Ok(self.run_forward(sequence)?.logits)
272    }
273
274    /// One batched forward pass over right-padded sequences (ferritin-100.12).
275    ///
276    /// AMPLIFY's `pad_mask` is **additive**, not `1 = real` — it is added to
277    /// the attention scores — so the `(batch, seq_len)` mask goes through
278    /// `additive_padding_mask` in the model's dtype before it is passed in.
279    /// Handing `forward` a 0/1 mask instead raises no error and masks nothing:
280    /// real keys would get `+1.0` and pads `+0.0`, so an unpadded row shifts
281    /// uniformly and cancels in softmax, while a padded row leaves its pad
282    /// keys with the *highest* bias of all.
283    fn embed_batch(&self, sequences: &[&str]) -> Result<Tensor> {
284        let device = self.model.get_device();
285        let pad_id = self
286            .tokenizer
287            .token_to_id("<pad>")
288            .ok_or_else(|| anyhow!("AMPLIFY tokenizer has no <pad> token"))?;
289        let rows = sequences
290            .iter()
291            .map(|s| {
292                Ok(self
293                    .tokenizer
294                    .encode(s.to_string(), true)
295                    .map_err(E::msg)?
296                    .get_ids()
297                    .to_vec())
298            })
299            .collect::<Result<Vec<_>>>()?;
300        let batch = pad_token_batch(&rows, pad_id, device)?;
301        let dtype = self.model.dtype();
302        let pad_mask = additive_padding_mask(&batch.mask, dtype)?;
303
304        let output = self
305            .model
306            .forward(&batch.ids, Some(&pad_mask), true, false)?;
307        let hidden = output
308            .hidden_states
309            .ok_or_else(|| anyhow!("AMPLIFY forward() returned no hidden states"))?
310            .pop()
311            .ok_or_else(|| anyhow!("AMPLIFY returned empty hidden states list"))?;
312        zero_padded_rows(&hidden, &batch.mask)
313    }
314}
ferritin_plms/amplify/amplify_runner.rs

ferritin_plms/amplify/
amplify_runner.rs
