logicaffeine_compile/model/tokenizer/
unicode.rs

1//! The codepoint facts the tokenizers read, as llama.cpp computes them (`src/unicode.cpp`): category flags, the
2//! lowercase map, the lenient UTF-8 decoder (an invalid byte decodes to U+FFFD and advances one byte), and GPT-2's
3//! byte-to-codepoint map. The tables in `unicode_tables` are generated from llama.cpp's running code and
4//! held equal to it for every codepoint by `model_tokenizer_lock`.
5//!
6//! The flags kept are the ones a tokenizer reads: the category and whitespace bits. llama.cpp's case and NFD bits
7//! never change a split (every codepoint carries exactly one category bit, so a flag word is non-zero either way).
8
9use std::sync::OnceLock;
10
11use super::unicode_tables::{FLAG_RANGES, LOWERCASE_RUNS};
12
13/// The number of codepoints the flag table covers; a larger value reads as undefined.
14pub const MAX_CODEPOINTS: u32 = 0x110000;
15
16/// A codepoint's flags, bit for bit llama.cpp's `unicode_cpt_flags`.
17#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
18pub struct Flags(pub u16);
19
20impl Flags {
21    pub const UNDEFINED: u16 = 0x0001;
22    pub const NUMBER: u16 = 0x0002;
23    pub const LETTER: u16 = 0x0004;
24    pub const SEPARATOR: u16 = 0x0008;
25    pub const ACCENT_MARK: u16 = 0x0010;
26    pub const PUNCTUATION: u16 = 0x0020;
27    pub const SYMBOL: u16 = 0x0040;
28    pub const CONTROL: u16 = 0x0080;
29    pub const MASK_CATEGORIES: u16 = 0x00FF;
30    pub const WHITESPACE: u16 = 0x0100;
31    /// Every bit the table keeps.
32    pub const KEPT: u16 = 0x01FF;
33
34    pub fn is_undefined(self) -> bool {
35        self.0 & Self::UNDEFINED != 0
36    }
37    pub fn is_number(self) -> bool {
38        self.0 & Self::NUMBER != 0
39    }
40    pub fn is_letter(self) -> bool {
41        self.0 & Self::LETTER != 0
42    }
43    pub fn is_separator(self) -> bool {
44        self.0 & Self::SEPARATOR != 0
45    }
46    pub fn is_accent_mark(self) -> bool {
47        self.0 & Self::ACCENT_MARK != 0
48    }
49    pub fn is_punctuation(self) -> bool {
50        self.0 & Self::PUNCTUATION != 0
51    }
52    pub fn is_symbol(self) -> bool {
53        self.0 & Self::SYMBOL != 0
54    }
55    pub fn is_control(self) -> bool {
56        self.0 & Self::CONTROL != 0
57    }
58    pub fn is_whitespace(self) -> bool {
59        self.0 & Self::WHITESPACE != 0
60    }
61
62    /// The category bits alone (`category_flag()`).
63    pub fn category(self) -> u16 {
64        self.0 & Self::MASK_CATEGORIES
65    }
66}
67
68/// The flags of `cpt`; past the table every codepoint is undefined. A table entry packs a range's first codepoint
69/// above its nine flag bits.
70pub fn flags(cpt: u32) -> Flags {
71    if cpt >= MAX_CODEPOINTS {
72        return Flags(Flags::UNDEFINED);
73    }
74    let i = FLAG_RANGES.partition_point(|&e| e >> 9 <= cpt) - 1;
75    Flags((FLAG_RANGES[i] & 0x1FF) as u16)
76}
77
78/// The flags of the codepoint at `pos` of `cpts[lo..hi]`, or none outside it (`_get_flags` in the splitters).
79pub fn flags_in(cpts: &[u32], lo: usize, hi: usize, pos: usize) -> Flags {
80    if lo <= pos && pos < hi {
81        flags(cpts[pos])
82    } else {
83        Flags(0)
84    }
85}
86
87/// `unicode_tolower`: the simple lowercase mapping, or the codepoint itself. The map is stored as runs
88/// `(first, count, stride, delta)`: `first + k·stride` lowercases to itself plus `delta` for `k < count`.
89pub fn to_lower(cpt: u32) -> u32 {
90    let i = LOWERCASE_RUNS.partition_point(|r| r.0 <= cpt);
91    if i == 0 {
92        return cpt;
93    }
94    let (first, count, stride, delta) = LOWERCASE_RUNS[i - 1];
95    let off = cpt - first;
96    if off % stride == 0 && off / stride < count {
97        (cpt as i64 + delta as i64) as u32
98    } else {
99        cpt
100    }
101}
102
103/// `unicode_cpt_is_han`: the CJK ideograph blocks the Kimi-K2 pre-tokenizer splits on.
104pub fn is_han(cpt: u32) -> bool {
105    matches!(cpt,
106        0x4E00..=0x9FFF
107        | 0x3400..=0x4DBF
108        | 0x20000..=0x2A6DF
109        | 0x2A700..=0x2B73F
110        | 0x2B740..=0x2B81F
111        | 0x2B820..=0x2CEAF
112        | 0x2CEB0..=0x2EBEF
113        | 0xF900..=0xFAFF
114        | 0x2F800..=0x2FA1F)
115}
116
117/// `unicode_len_utf8`: the sequence length a lead byte announces (a continuation byte counts as one).
118pub fn len_utf8(lead: u8) -> usize {
119    const LOOKUP: [usize; 16] = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4];
120    LOOKUP[(lead >> 4) as usize]
121}
122
123/// `unicode_cpt_from_utf8`: one codepoint, or `None` for an invalid sequence. Overlong forms and surrogates decode
124/// as written, the way llama.cpp's decoder reads them.
125fn cpt_from_utf8(b: &[u8], offset: &mut usize) -> Option<u32> {
126    let o = *offset;
127    let b0 = b[o] as u32;
128    let cont = |k: usize| b.get(o + k).is_some_and(|&c| c & 0xC0 == 0x80);
129    if b0 & 0x80 == 0 {
130        *offset += 1;
131        return Some(b0);
132    }
133    if b0 & 0x40 == 0 {
134        return None;
135    }
136    if b0 & 0x20 == 0 {
137        if !cont(1) {
138            return None;
139        }
140        *offset += 2;
141        return Some(((b0 & 0x1F) << 6) | (b[o + 1] as u32 & 0x3F));
142    }
143    if b0 & 0x10 == 0 {
144        if !(cont(1) && cont(2)) {
145            return None;
146        }
147        *offset += 3;
148        return Some(((b0 & 0x0F) << 12) | ((b[o + 1] as u32 & 0x3F) << 6) | (b[o + 2] as u32 & 0x3F));
149    }
150    if b0 & 0x08 == 0 {
151        if !(cont(1) && cont(2) && cont(3)) {
152            return None;
153        }
154        *offset += 4;
155        return Some(
156            ((b0 & 0x07) << 18) | ((b[o + 1] as u32 & 0x3F) << 12) | ((b[o + 2] as u32 & 0x3F) << 6) | (b[o + 3] as u32 & 0x3F),
157        );
158    }
159    None
160}
161
162/// `unicode_cpts_from_utf8`: every codepoint of `text`, an invalid byte reading as U+FFFD.
163pub fn cpts_from_utf8(text: &[u8]) -> Vec<u32> {
164    let mut out = Vec::with_capacity(text.len());
165    let mut offset = 0;
166    while offset < text.len() {
167        match cpt_from_utf8(text, &mut offset) {
168            Some(c) => out.push(c),
169            None => {
170                offset += 1;
171                out.push(0xFFFD);
172            }
173        }
174    }
175    out
176}
177
178/// A codepoint no UTF-8 sequence encodes: llama.cpp's `unicode_cpt_to_utf8` throws on it.
179#[derive(Clone, Copy, Debug, PartialEq, Eq)]
180pub struct InvalidCodepoint(pub u32);
181
182/// `unicode_cpt_to_utf8`, appending to `out`. Surrogates encode as their three bytes.
183pub fn push_utf8(cpt: u32, out: &mut Vec<u8>) -> Result<(), InvalidCodepoint> {
184    match cpt {
185        0..=0x7F => out.push(cpt as u8),
186        0x80..=0x7FF => out.extend_from_slice(&[0xC0 | ((cpt >> 6) & 0x1F) as u8, 0x80 | (cpt & 0x3F) as u8]),
187        0x800..=0xFFFF => out.extend_from_slice(&[
188            0xE0 | ((cpt >> 12) & 0x0F) as u8,
189            0x80 | ((cpt >> 6) & 0x3F) as u8,
190            0x80 | (cpt & 0x3F) as u8,
191        ]),
192        0x10000..=0x10FFFF => out.extend_from_slice(&[
193            0xF0 | ((cpt >> 18) & 0x07) as u8,
194            0x80 | ((cpt >> 12) & 0x3F) as u8,
195            0x80 | ((cpt >> 6) & 0x3F) as u8,
196            0x80 | (cpt & 0x3F) as u8,
197        ]),
198        _ => return Err(InvalidCodepoint(cpt)),
199    }
200    Ok(())
201}
202
203/// GPT-2's byte-level alphabet: the codepoint each byte is written as inside a byte-encoded token.
204fn byte_cpts() -> &'static [u32; 256] {
205    static MAP: OnceLock<[u32; 256]> = OnceLock::new();
206    MAP.get_or_init(|| {
207        let mut map = [u32::MAX; 256];
208        for ch in (0x21..=0x7E).chain(0xA1..=0xAC).chain(0xAE..=0xFF) {
209            map[ch as usize] = ch;
210        }
211        let mut n = 0;
212        for ch in 0..256 {
213            if map[ch] == u32::MAX {
214                map[ch] = 256 + n;
215                n += 1;
216            }
217        }
218        map
219    })
220}
221
222/// `unicode_byte_to_utf8`: the UTF-8 of the codepoint byte `b` is written as.
223pub fn byte_to_utf8(b: u8, out: &mut Vec<u8>) {
224    push_utf8(byte_cpts()[b as usize], out).expect("the byte alphabet is below U+0200");
225}
226
227/// `unicode_utf8_to_byte`: the byte a byte-level codepoint stands for, when it is one.
228pub fn utf8_to_byte(cpt: u32) -> Option<u8> {
229    static INVERSE: OnceLock<std::collections::HashMap<u32, u8>> = OnceLock::new();
230    INVERSE
231        .get_or_init(|| byte_cpts().iter().enumerate().map(|(b, &c)| (c, b as u8)).collect())
232        .get(&cpt)
233        .copied()
234}
logicaffeine_compile/model/tokenizer/unicode.rs

logicaffeine_compile/model/tokenizer/
unicode.rs
