ferritin_plms/esmc/tokenization/
sequence_tokenizer.rs

1use crate::esmc::utils::constants::esm3::SEQUENCE_VOCAB;
2use anyhow::Result;
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokenizers::models::bpe::{BPE, BpeBuilder};
6use tokenizers::processors::PostProcessorWrapper;
7use tokenizers::processors::template::{Template, TemplateProcessing};
8use tokenizers::{AddedToken, Tokenizer};
9
10pub trait EsmTokenizerBase {
11    fn encode(&self) -> Result<()>;
12    fn decode(&self) -> Result<()>;
13    fn mask_token(&self) -> &str;
14    fn mask_token_id(&self) -> u32;
15    fn bos_token(&self) -> &str;
16    fn bos_token_id(&self) -> u32;
17    fn eos_token(&self) -> &str;
18    fn eos_token_id(&self) -> u32;
19    fn pad_token(&self) -> &str;
20    fn pad_token_id(&self) -> u32;
21    fn chain_break_token(&self) -> &str;
22    fn chain_break_token_id(&self) -> u32;
23    fn all_token_ids(&self) -> Vec<u32>;
24    fn special_token_ids(&self) -> Vec<u32>;
25}
26
27pub struct EsmSequenceTokenizer {
28    tokenizer: Arc<Tokenizer>,
29    cb_token: String,
30}
31
32impl EsmSequenceTokenizer {
33    pub fn new(
34        unk_token: &str,
35        cls_token: &str,
36        pad_token: &str,
37        mask_token: &str,
38        eos_token: &str,
39        chain_break_token: &str,
40    ) -> Result<Self> {
41        let mut token_to_id = HashMap::new();
42        for (i, tok) in SEQUENCE_VOCAB.iter().enumerate() {
43            token_to_id.insert(tok.to_string(), i);
44        }
45        let bpe_builder = BpeBuilder::new();
46        let bpe: BPE = bpe_builder
47            .unk_token(unk_token.to_string())
48            .build()
49            .map_err(|e| anyhow::anyhow!("Failed to build BPE tokenizer: {}", e))?;
50
51        let mut tokenizer = Tokenizer::new(bpe);
52        let special_tokens = vec![
53            AddedToken::from(cls_token, true),
54            AddedToken::from(pad_token, true),
55            AddedToken::from(mask_token, true),
56            AddedToken::from(eos_token, true),
57            AddedToken::from(chain_break_token, true),
58        ];
59
60        let _ = tokenizer.add_special_tokens(special_tokens);
61
62        let post_processor = TemplateProcessing::builder()
63            .try_single(Template::try_from(format!("{} $A {}", cls_token, eos_token)).unwrap())?
64            .special_tokens(vec![
65                (cls_token, tokenizer.token_to_id(cls_token).unwrap()),
66                (eos_token, tokenizer.token_to_id(eos_token).unwrap()),
67            ])
68            .build()?;
69
70        tokenizer.with_post_processor(Some(PostProcessorWrapper::Template(post_processor)));
71
72        Ok(Self {
73            tokenizer: Arc::new(tokenizer),
74            cb_token: chain_break_token.to_string(),
75        })
76    }
77}
78impl Default for EsmSequenceTokenizer {
79    fn default() -> Self {
80        Self::new("<unk>", "<cls>", "<pad>", "<mask>", "<eos>", "|")
81            .expect("Failed to create default tokenizer")
82    }
83}
84
85impl EsmSequenceTokenizer {
86    /// Tokenize an amino-acid sequence string into token IDs.
87    ///
88    /// Looks up each character in `SEQUENCE_VOCAB`. Unknown characters map to
89    /// the `<unk>` token (index 3). When `add_special_tokens` is true, prepends
90    /// BOS (`<cls>` = 0) and appends EOS (`<eos>` = 2).
91    pub fn tokenize_sequence(&self, sequence: &str, add_special_tokens: bool) -> Vec<u32> {
92        use std::collections::HashMap;
93        let vocab: HashMap<&str, u32> = SEQUENCE_VOCAB
94            .iter()
95            .enumerate()
96            .map(|(i, s)| (*s, i as u32))
97            .collect();
98        let unk_id = *vocab.get("<unk>").unwrap_or(&3);
99
100        let mut tokens = Vec::with_capacity(sequence.len() + 2);
101        if add_special_tokens {
102            tokens.push(*vocab.get("<cls>").unwrap_or(&0));
103        }
104        for ch in sequence.chars() {
105            let s = ch.to_string();
106            let id = vocab.get(s.as_str()).copied().unwrap_or(unk_id);
107            tokens.push(id);
108        }
109        if add_special_tokens {
110            tokens.push(*vocab.get("<eos>").unwrap_or(&2));
111        }
112        tokens
113    }
114
115    /// Decode token IDs back to an amino-acid sequence string.
116    ///
117    /// Skips the standard special tokens (BOS=0, PAD=1, EOS=2, MASK=32) and
118    /// concatenates the remaining vocabulary entries.
119    pub fn decode_sequence(&self, token_ids: &[u32]) -> String {
120        const SPECIAL: [u32; 4] = [0, 1, 2, 32]; // cls, pad, eos, mask
121        let mut result = String::new();
122        for &id in token_ids {
123            if SPECIAL.contains(&id) {
124                continue;
125            }
126            if let Some(tok) = SEQUENCE_VOCAB.get(id as usize) {
127                result.push_str(tok);
128            }
129        }
130        result
131    }
132}
133
134impl EsmTokenizerBase for EsmSequenceTokenizer {
135    fn encode(&self) -> Result<()> {
136        // Trait signature carries no input/output; use tokenize_sequence for real encoding.
137        Ok(())
138    }
139
140    fn decode(&self) -> Result<()> {
141        // Trait signature carries no input/output; use decode_sequence for real decoding.
142        Ok(())
143    }
144
145    fn mask_token(&self) -> &str {
146        "<mask>"
147    }
148
149    fn mask_token_id(&self) -> u32 {
150        // <mask> is at index 32 in SEQUENCE_VOCAB
151        self.tokenizer.token_to_id("<mask>").unwrap_or(32)
152    }
153
154    fn bos_token(&self) -> &str {
155        "<cls>"
156    }
157
158    fn bos_token_id(&self) -> u32 {
159        // <cls> is at index 0 in SEQUENCE_VOCAB
160        self.tokenizer.token_to_id("<cls>").unwrap_or(0)
161    }
162
163    fn eos_token(&self) -> &str {
164        "<eos>"
165    }
166
167    fn eos_token_id(&self) -> u32 {
168        // <eos> is at index 2 in SEQUENCE_VOCAB
169        self.tokenizer.token_to_id("<eos>").unwrap_or(2)
170    }
171
172    fn pad_token(&self) -> &str {
173        "<pad>"
174    }
175
176    fn pad_token_id(&self) -> u32 {
177        // <pad> is at index 1 in SEQUENCE_VOCAB
178        self.tokenizer.token_to_id("<pad>").unwrap_or(1)
179    }
180
181    fn chain_break_token(&self) -> &str {
182        &self.cb_token
183    }
184
185    fn chain_break_token_id(&self) -> u32 {
186        // "|" is at index 31 in SEQUENCE_VOCAB
187        self.tokenizer.token_to_id(&self.cb_token).unwrap_or(31)
188    }
189
190    fn all_token_ids(&self) -> Vec<u32> {
191        (0..SEQUENCE_VOCAB.len() as u32).collect()
192    }
193
194    fn special_token_ids(&self) -> Vec<u32> {
195        // cls=0, pad=1, eos=2, chain_break=31, mask=32
196        vec![
197            self.bos_token_id(),
198            self.pad_token_id(),
199            self.eos_token_id(),
200            self.chain_break_token_id(),
201            self.mask_token_id(),
202        ]
203    }
204}
ferritin_plms/esmc/tokenization/sequence_tokenizer.rs

ferritin_plms/esmc/tokenization/
sequence_tokenizer.rs
