pagesift_ocr/formula/
tokenizer.rs

1//! 分词器的解码方向（id → 文本），按 tokenizer.json：模型是 BPE，解码器是 ByteLevel
2//! （GPT-2 的字节到可见字符映射）。语义同 Hugging Face tokenizers 的
3//! `decode(ids, skip_special_tokens=True)`：特殊 token 跳过，各 token 的字符逐个映射回字节、
4//! 拼起来按 UTF-8 解码（非法序列换成 U+FFFD）；某个 token 里有映射表外的字符时，
5//! 这个 token 按它自身的 UTF-8 字节原样输出。归一化器（NFKC）和预分词只用于编码，解码不涉及。
6
7use std::path::Path;
8
9/// 解码用的分词表。
10pub struct Tokenizer {
11    /// 下标是 id：这个 token 解码出的字节；特殊 token 和词表里没有的 id 为 `None`（跳过）。
12    bytes: Vec<Option<Vec<u8>>>,
13}
14
15impl Tokenizer {
16    /// 读模型目录下的 tokenizer.json。
17    pub fn load(dir: &Path) -> Result<Tokenizer, String> {
18        Tokenizer::from_json(&crate::read_json(dir, "tokenizer.json")?)
19    }
20
21    pub fn from_json(j: &serde_json::Value) -> Result<Tokenizer, String> {
22        let decoder = j["decoder"]["type"].as_str().unwrap_or("");
23        if decoder != "ByteLevel" {
24            return Err(format!(
25                "unsupported decoder {decoder:?}, only ByteLevel is supported"
26            ));
27        }
28        let model = j["model"]["type"].as_str().unwrap_or("");
29        if model != "BPE" {
30            return Err(format!(
31                "unsupported tokenizer model {model:?}, only BPE is supported"
32            ));
33        }
34        let vocab = j["model"]["vocab"]
35            .as_object()
36            .ok_or("tokenizer.json has no model.vocab")?;
37        let table = byte_table();
38        let mut bytes: Vec<Option<Vec<u8>>> = Vec::new();
39        let mut put = |id: usize, v: Option<Vec<u8>>| {
40            if bytes.len() <= id {
41                bytes.resize(id + 1, None);
42            }
43            bytes[id] = v;
44        };
45        for (tok, id) in vocab {
46            let id = id.as_u64().ok_or("a vocabulary id is not an integer")? as usize;
47            put(id, Some(token_bytes(tok, &table)));
48        }
49        // 追加的 token 覆盖词表里的同 id 项；特殊的解码时跳过
50        for a in j["added_tokens"].as_array().into_iter().flatten() {
51            let id = a["id"]
52                .as_u64()
53                .ok_or("an added_tokens id is not an integer")? as usize;
54            let content = a["content"]
55                .as_str()
56                .ok_or("an added_tokens entry has no content")?;
57            let special = a["special"].as_bool().unwrap_or(false);
58            put(id, (!special).then(|| token_bytes(content, &table)));
59        }
60        Ok(Tokenizer { bytes })
61    }
62
63    /// 词表大小（最大 id + 1）。
64    pub fn len(&self) -> usize {
65        self.bytes.len()
66    }
67
68    pub fn is_empty(&self) -> bool {
69        self.bytes.is_empty()
70    }
71
72    /// id 序列 → 文本，跳过特殊 token 和未知 id。
73    pub fn decode(&self, ids: &[u32]) -> String {
74        let mut out = Vec::new();
75        for &id in ids {
76            if let Some(Some(b)) = self.bytes.get(id as usize) {
77                out.extend_from_slice(b);
78            }
79        }
80        String::from_utf8_lossy(&out).into_owned()
81    }
82}
83
84/// 一个 token 的字节：每个字符按映射表还原成一个字节；有表外字符时取它自身的 UTF-8 字节。
85fn token_bytes(tok: &str, table: &[Option<u8>]) -> Vec<u8> {
86    tok.chars()
87        .map(|c| table.get(c as usize).copied().flatten())
88        .collect::<Option<Vec<u8>>>()
89        .unwrap_or_else(|| tok.as_bytes().to_vec())
90}
91
92/// GPT-2 `bytes_to_unicode` 的逆表，下标是字符的码位：
93/// 可见字节 `!`–`~`、`¡`–`¬`、`®`–`ÿ` 映射到同码位的字符，其余 68 个字节按从小到大依次映射到
94/// U+0100 起的字符。
95fn byte_table() -> Vec<Option<u8>> {
96    // [协议] GPT-2 encoder.py 的 bytes_to_unicode：三段可见字节区间与表外字节的起始码位 256
97    let visible = |b: u8| (33..=126).contains(&b) || (161..=172).contains(&b) || b >= 174;
98    // [协议] 同上：表外字节从 U+0100 起
99    const SHIFT: usize = 256;
100    let mut table = vec![None; SHIFT * 2];
101    let mut n = 0;
102    for b in 0..=u8::MAX {
103        if visible(b) {
104            table[usize::from(b)] = Some(b);
105        } else {
106            table[SHIFT + n] = Some(b);
107            n += 1;
108        }
109    }
110    table.truncate(SHIFT + n);
111    table
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    /// 小分词表：条目取自 PP-FormulaNet_plus-L 的 tokenizer.json（id 是原文件里的）。
119    fn mini() -> Tokenizer {
120        let j = serde_json::json!({
121            "decoder": {"type": "ByteLevel", "add_prefix_space": true},
122            "model": {"type": "BPE", "vocab": {
123                "<s>": 0, "<pad>": 1, "</s>": 2, "(": 30, "d": 90, "\\": 82,
124                "ĉ": 220, "Ċ": 221, "ĠĠ": 304, "ĠâĪĴ": 612, "âĢĵ": 853, "âĪĴ": 1059,
125                "ĠÎ¼": 1173, "ä": 183, "¸": 139, "Ń": 278, "ä¸": 19598, "Ġown": 4158
126            }},
127            "added_tokens": [
128                {"id": 0, "content": "<s>", "special": true},
129                {"id": 1, "content": "<pad>", "special": true},
130                {"id": 2, "content": "</s>", "special": true}
131            ]
132        });
133        Tokenizer::from_json(&j).unwrap()
134    }
135
136    #[test]
137    fn decodes_byte_level_tokens() {
138        let t = mini();
139        assert_eq!(t.decode(&[612]), " −");
140        assert_eq!(t.decode(&[853, 1059]), "–−");
141        assert_eq!(t.decode(&[1173]), " μ");
142        assert_eq!(t.decode(&[304, 220, 221]), "  \t\n");
143        assert_eq!(t.decode(&[82, 90, 30, 4158]), "\\d( own");
144        // 一个汉字的三个字节分在三个 token 里，拼起来才是完整的 UTF-8
145        assert_eq!(t.decode(&[183, 139, 278]), "中");
146        // 只有前两个字节：非法序列换成 U+FFFD
147        assert_eq!(t.decode(&[19598]), "\u{fffd}");
148    }
149
150    #[test]
151    fn skips_special_and_unknown() {
152        let t = mini();
153        assert_eq!(t.decode(&[2, 0, 90, 1, 2, 99_999]), "d");
154    }
155
156    #[test]
157    fn byte_table_is_a_bijection() {
158        let t = byte_table();
159        let mut seen = [false; 256];
160        for b in t.iter().flatten() {
161            assert!(!seen[usize::from(*b)]);
162            seen[usize::from(*b)] = true;
163        }
164        assert!(seen.iter().all(|&s| s));
165        // 空格 → Ġ（U+0120），换行 → Ċ（U+010A）
166        assert_eq!(t['Ġ' as usize], Some(b' '));
167        assert_eq!(t['Ċ' as usize], Some(b'\n'));
168    }
169
170    /// 真实的 tokenizer.json（模型目录在本机缓存里才跑）：所有条目都能解码，抽查几条。
171    #[test]
172    fn real_tokenizer_if_present() {
173        let Some(dir) = crate::models::cache_root()
174            .map(|r| r.join("PP-FormulaNet_plus-L_safetensors"))
175            .filter(|d| d.join("tokenizer.json").exists())
176        else {
177            return;
178        };
179        let t = Tokenizer::load(&dir).unwrap();
180        assert_eq!(t.len(), 50_000);
181        assert_eq!(t.decode(&[612, 1059, 1173, 4158]), " −− μ own");
182        assert_eq!(t.decode(&[183, 139, 278]), "中");
183        assert_eq!(t.decode(&[0, 1, 2, 3, 22]), "");
184    }
185}
pagesift_ocr/formula/tokenizer.rs

pagesift_ocr/formula/
tokenizer.rs
