interscript_rs/imf/
tokens.rs

1//! The canonical ByT5 byte table (SPEC.md §4) — fixed, no vocab files.
2//!
3//! UTF-8 byte `b` maps to id `b + 3`; one trailing EOS (id 1) is
4//! appended on encode; pad = 0, unk = 2. Ids are NOT raw bytes —
5//! feeding `text.as_bytes()` directly produces silent garbage on real
6//! models (the classic ByT5 interop bug; SPEC.md §4 warning).
7
8/// Byte offset: id = byte + 3.
9pub const BYTE_OFFSET: i64 = 3;
10/// Special id 0: padding (skipped on decode).
11pub const PAD_ID: i64 = 0;
12/// Special id 1: end of sequence (stops decode; appended on encode).
13pub const EOS_ID: i64 = 1;
14/// Special id 2: unknown (skipped on decode).
15pub const UNK_ID: i64 = 2;
16
17/// `tokens.encode`: UTF-8 bytes b -> ids b+3, then one trailing EOS.
18pub fn encode(text: &str) -> Vec<i64> {
19    let mut ids: Vec<i64> = text
20        .as_bytes()
21        .iter()
22        .map(|&b| i64::from(b) + BYTE_OFFSET)
23        .collect();
24    ids.push(EOS_ID);
25    ids
26}
27
28/// `tokens.decode`: stop at EOS; skip pad/unk; `(id - 3) mod 256` per
29/// byte; reassemble as UTF-8 (lossy, like the reference loader's
30/// `errors="replace"`).
31pub fn decode(token_ids: &[i64]) -> String {
32    let mut bytes = Vec::with_capacity(token_ids.len());
33    for &token in token_ids {
34        if token == EOS_ID {
35            break;
36        }
37        if token == PAD_ID || token == UNK_ID {
38            continue;
39        }
40        bytes.push((token - BYTE_OFFSET).rem_euclid(256) as u8);
41    }
42    String::from_utf8_lossy(&bytes).into_owned()
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn encode_appends_offset_and_eos() {
51        // "a" -> [97 + 3, EOS]
52        assert_eq!(encode("a"), vec![100, EOS_ID]);
53        // empty string still carries the trailing EOS (the reference
54        // loader short-circuits on that shape).
55        assert_eq!(encode(""), vec![EOS_ID]);
56    }
57
58    #[test]
59    fn encode_multibyte_is_per_byte() {
60        // Khmer "ខ" = U+1781 = UTF-8 0xE1 0x9E 0x81
61        assert_eq!(encode("ខ"), vec![0xE1 + 3, 0x9E + 3, 0x81 + 3, EOS_ID]);
62    }
63
64    #[test]
65    fn decode_roundtrip_ascii_and_multibyte() {
66        let text = "ភាសា en ک";
67        assert_eq!(decode(&encode(text)), text);
68    }
69
70    #[test]
71    fn decode_stops_at_eos_and_skips_pad_unk() {
72        // pad 0 / unk 2 are dropped, EOS truncates.
73        // (token 100 -> byte 97 'a', token 101 -> byte 98 'b')
74        assert_eq!(decode(&[PAD_ID, 100, UNK_ID, 101, EOS_ID, 102]), "ab");
75    }
76
77    #[test]
78    fn decode_ids_below_offset_mod_256() {
79        // (id - 3) mod 256: id 324 wraps to byte 65 'A'.
80        assert_eq!(decode(&[324, EOS_ID]), "A");
81    }
82}
interscript_rs/imf/tokens.rs

interscript_rs/imf/
tokens.rs
