ferritin_plms/esm3/layers/
encode_inputs.rs

1//! ESM3 multi-track input embedding (EncodeInputs).
2//!
3//! Embeds up to 8 input tracks into a single `(B, L, d_model)` tensor by summing
4//! per-track contributions. All tracks are optional; missing tracks contribute zero.
5
6use crate::esm3::models::esm3::ESM3Config;
7use candle_core::{D, Module, Result, Tensor};
8use candle_nn::{self as nn, VarBuilder};
9
10// ── EmbeddingBag ─────────────────────────────────────────────────────────────
11
12/// Sum-mode embedding bag: looks up embeddings for multiple indices per position
13/// and sums them, zeroing out the padding index.
14pub struct EmbeddingBag {
15    embed: nn::Embedding,
16    padding_idx: u32,
17}
18
19impl EmbeddingBag {
20    pub fn load(
21        vb: VarBuilder,
22        vocab_size: usize,
23        embed_dim: usize,
24        padding_idx: u32,
25    ) -> Result<Self> {
26        Ok(Self {
27            embed: nn::embedding(vocab_size, embed_dim, vb)?,
28            padding_idx,
29        })
30    }
31
32    /// `indices`: `(*, K)` — K annotation IDs per position; `padding_idx` entries → zero.
33    /// Returns `(*, embed_dim)` via sum over K.
34    pub fn forward(&self, indices: &Tensor) -> Result<Tensor> {
35        // Look up: (*, K) → (*, K, embed_dim)
36        let embedded = self.embed.forward(indices)?;
37
38        // Mask out padding: where index == padding_idx, contribution = 0
39        let pad = self.padding_idx;
40        let mask = indices
41            .ne(pad)?
42            .unsqueeze(D::Minus1)?
43            .broadcast_as(embedded.shape())?
44            .to_dtype(embedded.dtype())?;
45        let masked = (embedded * mask)?;
46
47        // Sum over K dimension (second-to-last)
48        masked.sum(D::Minus2)
49    }
50}
51
52// ── RBF encoding ─────────────────────────────────────────────────────────────
53
54/// Radial Basis Function encoding of scalar values into `n_bins` features.
55///
56/// `values`: `(*)` float tensor in `[v_min, v_max]`.
57/// Returns `(*, n_bins)`.
58fn rbf(values: &Tensor, v_min: f64, v_max: f64, n_bins: usize) -> Result<Tensor> {
59    let device = values.device();
60    let dtype = values.dtype();
61
62    // Evenly spaced centers in [v_min, v_max]
63    let centers: Vec<f32> = (0..n_bins)
64        .map(|i| (v_min + (v_max - v_min) * i as f64 / (n_bins - 1) as f64) as f32)
65        .collect();
66    let centers = Tensor::new(centers.as_slice(), device)?
67        .to_dtype(dtype)?
68        .reshape((1, 1, n_bins))?; // (1, 1, n_bins) for broadcasting
69
70    let width = (v_max - v_min) / (n_bins - 1) as f64;
71    let denom = 2.0 * width * width;
72
73    let v = values.unsqueeze(D::Minus1)?; // (*, 1)
74    let diff = v.broadcast_sub(&centers)?; // (*, n_bins)
75    diff.sqr()?.affine(-1.0 / denom, 0.0)?.exp()
76}
77
78// ── EncodeInputs ──────────────────────────────────────────────────────────────
79
80pub struct EncodeInputs {
81    #[allow(dead_code)]
82    d_model: usize,
83    // Sequence track
84    sequence_embed: nn::Embedding,
85    // pLDDT tracks (projected from 16-bin RBF)
86    plddt_projection: nn::Linear,
87    structure_per_res_plddt_projection: nn::Linear,
88    // Structure token track
89    structure_tokens_embed: nn::Embedding,
90    // Secondary structure and SASA
91    ss8_embed: nn::Embedding,
92    sasa_embed: nn::Embedding,
93    // Function annotation tracks (8 separate embeddings concatenated)
94    function_embeds: Vec<nn::Embedding>,
95    n_function_tracks: usize,
96    // Residue (InterPro) annotation track (EmbeddingBag, sum mode)
97    residue_embed: EmbeddingBag,
98}
99
100impl EncodeInputs {
101    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
102        let d = config.d_model;
103        let n_rbf = config.n_rbf_bins;
104
105        let sequence_embed = nn::embedding(config.d_sequence_vocab, d, vb.pp("sequence_embed"))?;
106
107        // Both plddt projections are `nn.Linear(16, d_model)` upstream — with
108        // bias. Loading them `no_bias` silently dropped two tensors that are
109        // present in the checkpoint (ferritin-100.31), the same defect
110        // ferritin-100.27 found in the structure encoder's FFN. It hid here
111        // because nothing ever called these projections: the plddt tracks were
112        // only ever `None`, and an absent track was skipped entirely.
113        let plddt_projection = nn::linear(n_rbf, d, vb.pp("plddt_projection"))?;
114        let structure_per_res_plddt_projection =
115            nn::linear(n_rbf, d, vb.pp("structure_per_res_plddt_projection"))?;
116
117        // Structure vocab: 4096 codes + 5 special tokens
118        let structure_tokens_embed = nn::embedding(
119            config.d_structure_vocab + 5,
120            d,
121            vb.pp("structure_tokens_embed"),
122        )?;
123
124        let ss8_embed = nn::embedding(config.d_ss8_vocab, d, vb.pp("ss8_embed"))?;
125        let sasa_embed = nn::embedding(config.d_sasa_vocab, d, vb.pp("sasa_embed"))?;
126
127        // 8 function-track embeddings, each produces d_model // 8 features
128        let func_dim = d / config.n_function_tracks;
129        let mut function_embeds = Vec::with_capacity(config.n_function_tracks);
130        for i in 0..config.n_function_tracks {
131            function_embeds.push(nn::embedding(
132                config.d_function_vocab,
133                func_dim,
134                vb.pp(format!("function_embed.{}", i)),
135            )?);
136        }
137
138        let residue_embed = EmbeddingBag::load(
139            vb.pp("residue_embed"),
140            config.d_residue_vocab,
141            d,
142            0, // padding_idx
143        )?;
144
145        Ok(Self {
146            d_model: d,
147            sequence_embed,
148            plddt_projection,
149            structure_per_res_plddt_projection,
150            structure_tokens_embed,
151            ss8_embed,
152            sasa_embed,
153            function_embeds,
154            n_function_tracks: config.n_function_tracks,
155            residue_embed,
156        })
157    }
158
159    /// Embed all input tracks and sum them.
160    ///
161    /// All arguments are optional; present tracks contribute their embedding,
162    /// absent tracks contribute zero.
163    ///
164    /// - `sequence_tokens`:           `(B, L)` u32 sequence token IDs.
165    /// - `structure_tokens`:          `(B, L)` u32 structure (VQ-VAE) tokens.
166    /// - `ss8_tokens`:                `(B, L)` u32 secondary-structure tokens.
167    /// - `sasa_tokens`:               `(B, L)` u32 SASA-bin tokens.
168    /// - `function_tokens`:           `(B, L, n_tracks)` u32 function annotation tokens.
169    /// - `residue_annotation_tokens`: `(B, L, K)` u32 InterPro annotation IDs.
170    /// - `average_plddt`:             `(B, L)` f32 average per-structure pLDDT in `[0,1]`.
171    /// - `per_res_plddt`:             `(B, L)` f32 per-residue pLDDT in `[0,1]`.
172    ///
173    /// Returns `(B, L, d_model)`.
174    #[allow(clippy::too_many_arguments)] // one arg per ESM3 input track
175    pub fn forward(
176        &self,
177        sequence_tokens: Option<&Tensor>,
178        structure_tokens: Option<&Tensor>,
179        ss8_tokens: Option<&Tensor>,
180        sasa_tokens: Option<&Tensor>,
181        function_tokens: Option<&Tensor>,
182        residue_annotation_tokens: Option<&Tensor>,
183        average_plddt: Option<&Tensor>,
184        per_res_plddt: Option<&Tensor>,
185    ) -> Result<Tensor> {
186        let n_rbf = 16usize;
187
188        // Accumulate embeddings into x; first present track initialises x.
189        let mut x: Option<Tensor> = None;
190        let mut add = |t: Tensor| -> Result<()> {
191            x = Some(match x.take() {
192                None => t,
193                Some(acc) => acc.add(&t)?,
194            });
195            Ok(())
196        };
197
198        if let Some(seq) = sequence_tokens {
199            add(self.sequence_embed.forward(seq)?)?;
200        }
201
202        if let Some(st) = structure_tokens {
203            add(self.structure_tokens_embed.forward(st)?)?;
204        }
205
206        if let Some(ss8) = ss8_tokens {
207            add(self.ss8_embed.forward(ss8)?)?;
208        }
209
210        if let Some(sasa) = sasa_tokens {
211            add(self.sasa_embed.forward(sasa)?)?;
212        }
213
214        // The plddt tracks are the only float inputs, so they are the only ones
215        // that can arrive in a dtype the weights do not share. Casting here is
216        // what keeps a half-precision ESM3 working: the upstream SDK omits this
217        // and dies with "expected m1 and m2 to have the same dtype" on a BF16
218        // checkpoint (ferritin-100.29).
219        if let Some(plddt) = average_plddt {
220            let enc =
221                rbf(plddt, 0.0, 1.0, n_rbf)?.to_dtype(self.plddt_projection.weight().dtype())?;
222            add(self.plddt_projection.forward(&enc)?)?;
223        }
224
225        if let Some(per_res) = per_res_plddt {
226            let enc = rbf(per_res, 0.0, 1.0, n_rbf)?
227                .to_dtype(self.structure_per_res_plddt_projection.weight().dtype())?;
228            add(self.structure_per_res_plddt_projection.forward(&enc)?)?;
229        }
230
231        if let Some(func) = function_tokens {
232            // func: (B, L, n_tracks); each track uses its own embedding
233            let mut parts: Vec<Tensor> = Vec::with_capacity(self.n_function_tracks);
234            for i in 0..self.n_function_tracks {
235                let track = func.narrow(D::Minus1, i, 1)?.squeeze(D::Minus1)?; // (B, L)
236                parts.push(self.function_embeds[i].forward(&track)?); // (B, L, func_dim)
237            }
238            let func_emb = Tensor::cat(&parts, D::Minus1)?; // (B, L, d_model)
239            add(func_emb)?;
240        }
241
242        if let Some(res) = residue_annotation_tokens {
243            add(self.residue_embed.forward(res)?)?;
244        }
245
246        // If no tracks provided, return zeros. In practice at least sequence_tokens is present.
247        match x {
248            Some(t) => Ok(t),
249            None => candle_core::bail!("EncodeInputs: at least one input track must be provided"),
250        }
251    }
252}
ferritin_plms/esm3/layers/encode_inputs.rs

ferritin_plms/esm3/layers/
encode_inputs.rs
