ferritin_plms/esm3/tokenization/
ss8.rs

1//! ESM3 SS8 secondary-structure tokenization.
2//!
3//! Token layout:
4//!   0 = PAD, 1 = MASK, 2 = UNK
5//!   3..10 = 'G','H','I','T','E','B','S','C'  (DSSP 8-class)
6
7use crate::esm3::utils::constants::{SS8_UNK_TOKEN, SS8_VOCAB};
8
9/// Encode an SS8 secondary-structure string to token IDs.
10///
11/// Valid characters are "GHITEBSC". Unknown characters map to UNK (2).
12pub fn tokenize_ss8(ss8: &str) -> Vec<u32> {
13    ss8.chars()
14        .map(|c| match SS8_VOCAB.find(c) {
15            Some(idx) => idx as u32 + 3,
16            None => SS8_UNK_TOKEN,
17        })
18        .collect()
19}
20
21#[cfg(test)]
22mod tests {
23    use super::*;
24
25    #[test]
26    fn test_tokenize_ss8_known() {
27        // "G" is first in SS8_VOCAB → index 0 → token 3
28        assert_eq!(tokenize_ss8("G"), vec![3]);
29        // "C" is last → index 7 → token 10
30        assert_eq!(tokenize_ss8("C"), vec![10]);
31    }
32
33    #[test]
34    fn test_tokenize_ss8_unknown() {
35        assert_eq!(tokenize_ss8("X"), vec![2]); // UNK
36    }
37
38    #[test]
39    fn test_tokenize_ss8_full_vocab() {
40        let ids = tokenize_ss8(SS8_VOCAB);
41        assert_eq!(ids, vec![3, 4, 5, 6, 7, 8, 9, 10]);
42    }
43}
ferritin_plms/esm3/tokenization/ss8.rs

ferritin_plms/esm3/tokenization/
ss8.rs
