ferritin_plms/esm3/tokenization/
sequence.rs

1//! ESM3 sequence tokenization.
2
3use crate::esm3::utils::constants::{SEQUENCE_BOS_TOKEN, SEQUENCE_EOS_TOKEN, SEQUENCE_VOCAB};
4use std::collections::HashMap;
5
6fn vocab_map() -> HashMap<&'static str, u32> {
7    SEQUENCE_VOCAB
8        .iter()
9        .enumerate()
10        .map(|(i, s)| (*s, i as u32))
11        .collect()
12}
13
14/// Encode an amino-acid string to ESM3 sequence token IDs.
15///
16/// Unknown characters map to `<unk>` (index 3).
17/// When `add_special_tokens` is true, prepends BOS and appends EOS.
18pub fn tokenize_sequence(sequence: &str, add_special_tokens: bool) -> Vec<u32> {
19    let vocab = vocab_map();
20    let unk = *vocab.get("<unk>").unwrap_or(&3);
21
22    let mut tokens = Vec::with_capacity(sequence.len() + 2);
23    if add_special_tokens {
24        tokens.push(SEQUENCE_BOS_TOKEN);
25    }
26    for ch in sequence.chars() {
27        let s = ch.to_string();
28        tokens.push(vocab.get(s.as_str()).copied().unwrap_or(unk));
29    }
30    if add_special_tokens {
31        tokens.push(SEQUENCE_EOS_TOKEN);
32    }
33    tokens
34}
35
36/// Decode ESM3 sequence token IDs back to an amino-acid string.
37///
38/// Skips BOS (0), PAD (1), EOS (2), and MASK (32).
39pub fn decode_sequence(token_ids: &[u32]) -> String {
40    const SPECIAL: [u32; 4] = [0, 1, 2, 32];
41    let mut out = String::new();
42    for &id in token_ids {
43        if SPECIAL.contains(&id) {
44            continue;
45        }
46        if let Some(tok) = SEQUENCE_VOCAB.get(id as usize) {
47            out.push_str(tok);
48        }
49    }
50    out
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn test_tokenize_sequence_no_special() {
59        let ids = tokenize_sequence("MA", false);
60        // M is at index 20, A is at index 5
61        assert_eq!(ids, vec![20, 5]);
62    }
63
64    #[test]
65    fn test_tokenize_sequence_with_special() {
66        let ids = tokenize_sequence("G", true);
67        assert_eq!(ids[0], SEQUENCE_BOS_TOKEN);
68        assert_eq!(*ids.last().unwrap(), SEQUENCE_EOS_TOKEN);
69        assert_eq!(ids.len(), 3);
70    }
71
72    #[test]
73    fn test_decode_roundtrip() {
74        let seq = "ACDEFGHIKLMNPQRSTVWY";
75        let ids = tokenize_sequence(seq, true);
76        let decoded = decode_sequence(&ids);
77        assert_eq!(decoded, seq);
78    }
79
80    #[test]
81    fn test_unknown_char_maps_to_unk() {
82        let ids = tokenize_sequence("?", false);
83        assert_eq!(ids, vec![3]); // <unk>
84    }
85}
ferritin_plms/esm3/tokenization/sequence.rs

ferritin_plms/esm3/tokenization/
sequence.rs
