ferritin_plms/esm3/tokenization/
structure.rs

1//! ESM3 structure tokenization.
2//!
3//! Structure tokens are VQ-VAE codebook indices (0..4095). This module
4//! handles wrapping them with BOS/EOS special tokens for model input.
5
6use crate::esm3::utils::constants::{STRUCTURE_BOS_TOKEN, STRUCTURE_EOS_TOKEN};
7
8/// Wrap VQ-VAE codebook indices with optional BOS/EOS special tokens.
9///
10/// `codes` should contain raw VQ-VAE indices in [0, 4095]. The function
11/// does not validate range — that is left to the VQ-VAE encoder output.
12pub fn tokenize_structure(codes: &[u32], add_special_tokens: bool) -> Vec<u32> {
13    let mut tokens = Vec::with_capacity(codes.len() + 2);
14    if add_special_tokens {
15        tokens.push(STRUCTURE_BOS_TOKEN);
16    }
17    tokens.extend_from_slice(codes);
18    if add_special_tokens {
19        tokens.push(STRUCTURE_EOS_TOKEN);
20    }
21    tokens
22}
23
24#[cfg(test)]
25mod tests {
26    use super::*;
27    use crate::esm3::utils::constants::{STRUCTURE_BOS_TOKEN, STRUCTURE_EOS_TOKEN};
28
29    #[test]
30    fn test_tokenize_structure_no_special() {
31        let codes = vec![0u32, 100, 4095];
32        assert_eq!(tokenize_structure(&codes, false), codes);
33    }
34
35    #[test]
36    fn test_tokenize_structure_with_special() {
37        let codes = vec![42u32];
38        let tokens = tokenize_structure(&codes, true);
39        assert_eq!(tokens, vec![STRUCTURE_BOS_TOKEN, 42, STRUCTURE_EOS_TOKEN]);
40    }
41
42    #[test]
43    fn test_tokenize_structure_empty() {
44        let tokens = tokenize_structure(&[], true);
45        assert_eq!(tokens, vec![STRUCTURE_BOS_TOKEN, STRUCTURE_EOS_TOKEN]);
46    }
47}
ferritin_plms/esm3/tokenization/structure.rs

ferritin_plms/esm3/tokenization/
structure.rs
