Skip to main content

logicaffeine_compile/model/hf/
vocab.rs

1//! The tokenizer keys llama.cpp's converter writes from a checkpoint directory's tokenizer files, computed as the
2//! converter computes them: a byte-pair vocabulary from `tokenizer.json` (`get_vocab_base`), a SentencePiece one from
3//! `tokenizer.model` (`_create_vocab_sentencepiece`), and the special tokens, added-token flags, merges, normalizer
4//! flags and chat template `gguf.SpecialVocab` reads beside them.
5//!
6//! The converter names a byte-pair vocabulary's pre-tokenizer by the digest of the ids its Hugging Face tokenizer gives
7//! a fixed probe text ([`PRE_TOKENIZER_PROBE`]), looked up in a table of the digests of known models
8//! ([`PRE_TOKENIZER_DIGESTS`]). This reader encodes the probe with the directory's own `tokenizer.json`, run as the
9//! `tokenizers` library runs it ([`HfTokenizer`]), and looks the digest up in the same table.
10
11use std::collections::{BTreeMap, HashMap};
12use std::io;
13use std::path::Path;
14
15use serde_json::{Map, Value};
16
17use super::encode::HfTokenizer;
18use super::tiktoken::{byte_level, Tiktoken};
19use super::{invalid, Family, Hparams};
20use crate::crypto_accel::sha256;
21use crate::model::gguf::MetaValue;
22
23/// The text whose tokenization names a byte-pair pre-tokenizer (the converter's `chktxt`).
24pub const PRE_TOKENIZER_PROBE: &str = "\n \n\n \n\n\n \t \t\t \t\n  \n   \n    \n     \n🚀 (normal) 😶\u{200d}🌫️ (multiple emojis concatenated) ✅ 🦙🦙 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច😁 ?我想在apple工作1314151天~ ------======= нещо на Български ''''''```````\"\"\"\"......!!!!!!?????? I've been 'told he's there, 'RE you sure? 'M not sure I'll make it, 'D you like some tea? We'Ve a'lL";
25
26/// The converter's `get_vocab_base_pre` table: the SHA-256 of the probe's ids, written as a Python list, and the
27/// pre-tokenizer name it writes for them, in the converter's order (a digest listed twice names what its last entry
28/// says).
29pub const PRE_TOKENIZER_DIGESTS: &[(&str, &str)] = &[
30    ("b6e8e1518dc4305be2fe39c313ed643381c4da5db34a98f6a04c093f8afbe99b", "chatglm-bpe"),
31    ("81d72c7348a9f0ebe86f23298d37debe0a5e71149e29bd283904c02262b27516", "chatglm-bpe"),
32    ("a1336059768a55c99a734006ffb02203cd450fed003e9a71886c88acf24fdbc2", "glm4"),
33    ("9ca2dd618e8afaf09731a7cf6e2105b373ba6a1821559f258b272fe83e6eb902", "glm4"),
34    ("cdf5f35325780597efd76153d4d1c16778f766173908894c04afc20108536267", "glm4"),
35    ("1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35", "minerva-7b"),
36    ("7e57df22b1fe23a7b1e1c7f3dc4e3f96d43a4eb0836d0c6bdc3436d7b2f1c664", "hunyuan"),
37    ("bba3b3366b646dbdded5dbc42d59598b849371afc42f7beafa914afaa5b70aa6", "hunyuan-dense"),
38    ("e6ddf9c6686791c12d698d34c31ab9be1fea9af5a3d9a6909783ab382198ae1c", "hy_v4"),
39    ("a6b57017d60e6edb4d88ecc2845188e0eb333a70357e45dcc9b53964a73bbae6", "falcon-h1"),
40    ("60476e1243776c4fb1b993dbd7a5f15ac22f83c80afdf425fa5ae01c8d44ef86", "falcon-h1"),
41    ("3eda48b4c4dc7de733d1a8b3e3b4a85243dbbf704da2ee9d42c6beced8897896", "falcon-h1"),
42    ("48f8e02c0359c0bbdd82f26909171fac1c18a457bb47573ed1fe3bbb2c1cfd4b", "falcon-h1"),
43    ("81212dc7cdb7e0c1074ca62c5aeab0d43c9f52b8a737be7b12a777c953027890", "kimi-k2"),
44    ("d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c", "qwen2"),
45    ("1444df51289cfa8063b96f0e62b1125440111bc79a52003ea14b6eac7016fd5f", "qwen35"),
46    ("66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273", "grok-2"),
47    ("b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df", "jina-v2-de"),
48    ("0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4", "gpt-2"),
49    ("9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7", "lfm2"),
50    ("846deafc5b0fa786186fa4ae6c7b49903cf2f1d1895bdb80b9120d60be135252", "gemma4"),
51    ("0a766d034107bc736a3f2dc4968fd62e54a3570f1454443e0c5a4cc6bd7941ed", "spark2_5"),
52    ("1f9825a388f700a6b591722f17d470cbbcf10973ece35d2fd14239a14110ae1a", "k2-horizon"),
53    ("0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5", "llama-bpe"),
54    ("049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754", "deepseek-llm"),
55    ("347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821", "deepseek-coder"),
56    ("8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed", "falcon"),
57    ("0876d13b50744004aa9aeae05e7b0647eac9d801b5ba4668afc01e709c15e19f", "bert-bge"),
58    ("9d032fcbd5501f4a38150912590928bfb36091efb5df11b8e2124b0390e3fb1e", "falcon3"),
59    ("8e62295832751ca1e8f92f2226f403dea30dc5165e448b5bfa05af5340c64ec7", "bert-bge-large"),
60    ("b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166", "mpt"),
61    ("35d91631860c815f952d711435f48d356ebac988362536bed955d43bfa436e34", "starcoder"),
62    ("3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454", "gpt-2"),
63    ("32d85c31273f8019248f2559fed492d929ea28b17e51d81d3bb36fff23ca72b3", "stablelm2"),
64    ("6221ad2852e85ce96f791f476e0b390cf9b474c9e3d1362f53a24a06dc8220ff", "refact"),
65    ("9c2227e4dd922002fb81bde4fc02b0483ca4f12911410dee2255e4987644e3f8", "command-r"),
66    ("d772b220ace2baec124bed8cfafce0ead7d6c38a4b65ef11261cf9d5d62246d1", "tiny_aya"),
67    ("52df12b4c8d4176e7481aab4b6e8454d1fd0a210a04a574f6d4e067d10e23c3e", "cohere2moe"),
68    ("e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea", "qwen2"),
69    ("b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166", "olmo"),
70    ("a8594e3edff7c29c003940395316294b2c623e09894deebbc65f33f1515df79e", "dbrx"),
71    ("c7699093ba4255a91e702aa38a596aa81669f3525dae06c2953267dde580f448", "jina-v1-en"),
72    ("0876d13b50744004aa9aeae05e7b0647eac9d801b5ba4668afc01e709c15e19f", "jina-v2-en"),
73    ("171aeeedd6fb548d418a7461d053f11b6f1f1fc9b387bd66640d28a4b9f5c643", "jina-v2-es"),
74    ("27949a2493fc4a9f53f5b9b029c82689cfbe5d3a1929bb25e043089e28466de6", "jina-v2-de"),
75    ("a023e9fdc5a11f034d3ef515b92350e56fb2af1f66c6b6811a4444ea9bf8763d", "jina-v5-nano"),
76    ("c136ed14d01c2745d4f60a9596ae66800e2b61fa45643e72436041855ad4089d", "smaug-bpe"),
77    ("c7ea5862a53e4272c035c8238367063e2b270d51faa48c0f09e9d5b54746c360", "poro-chat"),
78    ("7967bfa498ade6b757b064f31e964dddbb80f8f9a4d68d4ba7998fcf281c531a", "jina-v2-code"),
79    ("7fc505bd3104ca1083b150b17d088b59534ede9bde81f0dd2090967d7fe52cee", "viking"),
80    ("b53802fb28e26d645c3a310b34bfe07da813026ec7c7716883404d5e0f8b1901", "jais"),
81    ("bc5108ee1eb6a3d600cadd065f63190fbd0554dbc9e4bbd6a0d977970afc8d2a", "jais-2"),
82    ("7b3e7548e4308f52a76e8229e4e6cc831195d0d1df43aed21ac6c93da05fec5f", "codeshell"),
83    ("63b97e4253352e6f357cc59ea5b583e3a680eaeaf2632188c2b952de2588485e", "tekken"),
84    ("855059429035d75a914d1eda9f10a876752e281a054a7a3d421ef0533e5b6249", "smollm"),
85    ("3c30d3ad1d6b64202cd222813e7736c2db6e1bd6d67197090fc1211fbc612ae7", "bloom"),
86    ("bc01ce58980e1db43859146dc51b1758b3b88729b217a74792e9f8d43e479d21", "gpt3-finnish"),
87    ("4e2b24cc4770243d65a2c9ec19770a72f08cffc161adbb73fcbb6b7dd45a0aae", "exaone"),
88    ("fcace8b9cac38ce847670c970cd5892031a753a1ef381abd1d9af00f713da085", "phi-2"),
89    ("60824e3c0d9401f89943cbb2fff727f0e2d4c545ba4df2d6e4f09a6db0f5b450", "chameleon"),
90    ("8b5a93ed704057481f240da0be7e7dca721d7f8f4755263b6807227a2cbeae65", "roberta-bpe"),
91    ("ad851be1dba641f2e3711822f816db2c265f788b37c63b4e1aeacb9ee92de8eb", "gigachat"),
92    ("d4c8f286ea6b520b3d495c4455483cfa2302c0cfcd4be05d781b6a8a0a7cdaf1", "megrez"),
93    ("877081d19cf6996e2c4ff0e1236341e9b7bde288f5311a56a937f0afbbb3aeb5", "deepseek-v3"),
94    ("b3f499bb4255f8ca19fccd664443283318f2fd2414d5e0b040fbdd0cc195d6c5", "deepseek-r1-qwen"),
95    ("ccc2ef013c104be7bae2965776d611e1d7a8a2a9c547dd93a682c9a9fc80352e", "gpt-4o"),
96    ("7dec86086fcc38b66b7bc1575a160ae21cf705be7718b9d5598190d7c12db76f", "superbpe"),
97    ("1994ffd01900cfb37395608534236ecd63f2bd5995d6cb1004dda1af50240f15", "trillion"),
98    ("96a5f08be6259352137b512d4157e333e21df7edd3fcd152990608735a65b224", "bailingmoe"),
99    ("d353350c764d8c3b39c763113960e4fb4919bea5fbf208a0e3b22e8469dc7406", "llama4"),
100    ("0e9433cbbb161f89e264eb32e8e64bfe69e834973ffca5d41d3948a604a3e2a3", "pixtral"),
101    ("d5f1dd6f980fec569fb218a81a7658ac45fc56b38c5a0adeb1c232fbe04ef5ec", "seed-coder"),
102    ("b0a6b1c0bd5998ebd9df08611efde34a4ff03faed45ae09c43e6b31ebd4b94cf", "a.x-4.0"),
103    ("f6791d196f87ce6b56a7d234be618e0d58f8cda3549416635b2bebcd22cd95c4", "midm-2.0"),
104    ("169bf0296a13c4d9b7672313f749eb36501d931022de052aad6e36f2bf34dd51", "lfm2"),
105    ("2085e1638f6c377a0aa4ead21b27bb4cb941bf800df86ed391011769c1758dfb", "exaone4"),
106    ("a1e163ecab2e718a4c829d1148b6e86824ec36163bb71941c3dca9cd5ac25756", "mellum"),
107    ("a0b64b4385f123663873756336c085744376d015ff328bb1d901598f63c44152", "modern-bert"),
108    ("49fc0303c9e0d2c2c565c510f64b2d9b271276acdcdadff733249eda9f7d59df", "afmoe"),
109    ("9b1be57e70d20d9501b2b3186e792d81181ae36ada3903c26f9fea418cf87206", "bailingmoe2"),
110    ("53e325976a6e142379c19b09afcae354f2f496f147afa8f9e189a33fe4e3024e", "granite-docling"),
111    ("f4f37b6c8eb9ea29b3eac6bb8c8487c5ab7885f8d8022e67edc1c68ce8403e95", "minimax-m2"),
112    ("4a2e2abae11ca2b86d570fc5b44be4d5eb5e72cc8f22dd136a94b37da83ab665", "kormo"),
113    ("9d70134b369a70e5735009b6de918f7581b5211f7c074d1f89f753aea8248af1", "youtu"),
114    ("16389f0a1f51ee53e562ffd51c371dc508639ab0e4261502071836e50e223e91", "solar-open"),
115    ("6c81ce329e0802883b22eabab0d3fa48357337ef1ecb45443828bf1f6254833f", "exaone-moe"),
116    ("d30d75d9059f1aa2c19359de71047b3ae408c70875e8a3ccf8c5fba56c9d8af4", "qwen35"),
117    ("b4b8ca1f9769494fbd956ebc4c249de6131fb277a4a3345a7a92c7dd7a55808d", "joyai-llm"),
118    ("e4d54df1ebc1f2b91acd986c5b51aa50837d5faf7c7398e73c1f9e9ee5d19869", "kanana2"),
119    ("862f827721df956049dff5ca81a57f29e575280bc622e290d3bf4e35eca29015", "f2llmv2"),
120    ("62f6fb0a6fd5098caeabb19b07a5c1099cafc8b9c40eab6ea89ece4ec02fbc57", "sarvam-moe"),
121    ("f728162c1315c26e40249849799b4ba3fe584c32084b4795b03eb295e63cb5af", "talkie"),
122    ("36f3066e97b7f3994b379aaacde306c1444c6ae84e81a5ae3cd2b7ed3b8c42d4", "minicpm5"),
123    ("f241072145675bf8322086f115aebad05e9f869557a238bf2150a2a417d1bf60", "granite-embed-multi-97m"),
124    ("789696f5946cc0fc59371f39f6097cafed196b3acded6140432f26bbb1ae1669", "granite-embed-multi-311m"),
125    ("9dcf830ee9990cdbf78cc523a5f7bd9ad8f3f9890c2d3581d2785ad10f07049d", "mellum2"),
126    ("972da7b59cec44d1f0a490a86c96df53859e486e481563e5dddac155013d87ac", "laguna"),
127    ("653660222fb704f61cbf2b618a8ae6502b7f8b20c980f9a5de07ed78e13319cd", "ufakzeka"),
128    ("4b05e02dad1c5ae07d266fd3342ddb644c6f6be058d728bc0a33af31a1d6ee66", "mmbert"),
129    ("a9af07a84191f55098b248ae6f3dfe9e32d3190bebe8eafd91c1ddec9bc3449f", "k2-horizon"),
130];
131
132/// GGUF value types of the arrays the converter writes.
133const STRING: u32 = 8;
134const INT32: u32 = 5;
135const FLOAT32: u32 = 6;
136
137/// `gguf.TokenType`.
138const NORMAL: i32 = 1;
139const UNKNOWN: i32 = 2;
140const CONTROL: i32 = 3;
141const USER_DEFINED: i32 = 4;
142const UNUSED: i32 = 5;
143const BYTE: i32 = 6;
144
145/// The special-token kinds `SpecialVocab` reads by default.
146const SPECIAL_TYPES: [&str; 7] = ["bos", "eos", "unk", "sep", "pad", "cls", "mask"];
147
148/// The keys as the converter's writer leaves them: a later write of a key replaces the earlier.
149#[derive(Default)]
150struct Keys(BTreeMap<String, MetaValue>);
151
152impl Keys {
153    fn put(&mut self, key: &str, v: MetaValue) {
154        self.0.insert(key.to_string(), v);
155    }
156    fn str(&mut self, key: &str, v: &str) {
157        self.put(key, MetaValue::Str(v.as_bytes().to_vec()));
158    }
159    fn u32(&mut self, key: &str, v: u32) {
160        self.put(key, MetaValue::U32(v));
161    }
162    fn bool(&mut self, key: &str, v: bool) {
163        self.put(key, MetaValue::Bool(v));
164    }
165    fn strings(&mut self, key: &str, items: &[Vec<u8>]) {
166        self.put(key, MetaValue::Array { ty: STRING, items: items.iter().map(|s| MetaValue::Str(s.clone())).collect() });
167    }
168}
169
170fn read_json(path: &Path) -> io::Result<Option<Value>> {
171    if !path.is_file() {
172        return Ok(None);
173    }
174    let text = std::fs::read_to_string(path).map_err(|e| io::Error::new(e.kind(), format!("{}: {e}", path.display())))?;
175    serde_json::from_str(&text).map(Some).map_err(|e| invalid(format!("{}: {e}", path.display())))
176}
177
178/// Python's truth value of a JSON value: absent, `null`, `false`, zero and empty strings, arrays and objects are false.
179fn truthy(v: Option<&Value>) -> bool {
180    match v {
181        None | Some(Value::Null) => false,
182        Some(Value::Bool(b)) => *b,
183        Some(Value::Number(n)) => n.as_f64().is_some_and(|x| x != 0.0),
184        Some(Value::String(s)) => !s.is_empty(),
185        Some(Value::Array(a)) => !a.is_empty(),
186        Some(Value::Object(o)) => !o.is_empty(),
187    }
188}
189
190/// Whether a JSON value is the string `s` (Python's `==` between a string and whatever the file holds).
191fn is(v: &Option<Value>, s: &str) -> bool {
192    matches!(v, Some(Value::String(x)) if x == s)
193}
194
195/// `ModelBase.does_token_look_special`: added tokens some checkpoints leave unmarked that are control tokens.
196fn looks_special(t: &str) -> bool {
197    matches!(t, "<pad>" | "<mask>" | "<2mass>" | "[@BOS@]")
198        || (t.starts_with("<|") && t.ends_with("|>"))
199        || (t.starts_with("<|") && t.ends_with("|>"))
200        || (t.starts_with("<unused") && t.ends_with('>'))
201}
202
203/// The tokenizer keys of the checkpoint in `dir`, with `general.architecture`, as the converter writes them for
204/// family `f`.
205pub(super) fn tokenizer_meta(dir: &Path, h: &Hparams, f: Family) -> io::Result<BTreeMap<String, MetaValue>> {
206    let mut w = Keys::default();
207    w.str("general.architecture", f.arch());
208    let sentencepiece_file = dir.join("tokenizer.model").is_file();
209    match f {
210        Family::Llama => llama(dir, h, &mut w)?,
211        // `Qwen2Model.set_vocab`, which Qwen3's class inherits; the mixtures' and Qwen3.5's classes read tokenizer.json.
212        Family::Qwen2 | Family::Qwen3 if sentencepiece_file => sentencepiece(dir, h, &mut w)?,
213        Family::Gemma3 if sentencepiece_file => {
214            sentencepiece(dir, h, &mut w)?;
215            w.bool("tokenizer.ggml.add_space_prefix", false);
216        }
217        // `Glm4Model.set_vocab` and `_set_vocab_glm` (GLM-4.5's class): the byte-pair vocabulary with special tokens
218        // named by their text, each kept only where the files name none of its kind.
219        Family::Glm4 => gpt2(dir, h, &mut w, &[("eos", "<|endoftext|>"), ("eot", "<|user|>"), ("unk", "<|endoftext|>"), ("bos", "<|endoftext|>")])?,
220        Family::Glm4Moe => gpt2(dir, h, &mut w, &[("bos", "[gMASK]"), ("eot", "<|user|>"), ("unk", "<|endoftext|>"), ("eom", "<|observation|>")])?,
221        // `DeepseekV2Model.set_vocab`: a directory with no `tokenizer.json` is read through its tokenizer code, and the
222        // converter writes only Kimi K2's tiktoken vocabulary.
223        Family::Deepseek2 if !dir.join("tokenizer.json").is_file() && dir.join("tiktoken.model").is_file() => kimi(dir, h, &mut w)?,
224        _ => gpt2(dir, h, &mut w, &[])?,
225    }
226    Ok(w.0)
227}
228
229/// `LlamaModel.set_vocab`: a SentencePiece model when the checkpoint ships one, else its byte-pair `tokenizer.json`;
230/// Code Llama's infill tokens and the small Granite models' unadded beginning token on top.
231fn llama(dir: &Path, h: &Hparams, w: &mut Keys) -> io::Result<()> {
232    if let Some(tc) = read_json(&dir.join("tokenizer_config.json"))? {
233        match tc.get("add_prefix_space") {
234            None | Some(Value::Null) => {}
235            Some(Value::Bool(b)) => w.bool("tokenizer.ggml.add_space_prefix", *b),
236            Some(other) => return Err(invalid(format!("tokenizer_config.json: add_prefix_space is {other}"))),
237        }
238        if tc.get("tokenizer_class").and_then(Value::as_str) == Some("HybridDNATokenizer") {
239            return Err(invalid("tokenizer_config.json: a HybridDNA tokenizer, which this reader does not read"));
240        }
241    }
242    if dir.join("tekken.json").is_file() && !dir.join("tokenizer.json").is_file() {
243        return Err(invalid(format!("{}: a Mistral tekken tokenizer, which this reader does not read", dir.display())));
244    }
245    if dir.join("tokenizer.model").is_file() {
246        sentencepiece(dir, h, w)?;
247    } else {
248        // `LlamaHfVocab` takes a byte-fallback `tokenizer.json` with a sequence decoder; Llama 3's and every other
249        // byte-pair vocabulary go to `_set_vocab_gpt2`.
250        if let Some(tj) = read_json(&dir.join("tokenizer.json"))? {
251            let m = &tj["model"];
252            let bpe = m["type"].as_str() == Some("BPE");
253            let llama3 = bpe && truthy(m.get("ignore_merges")) && !truthy(Some(m.get("byte_fallback").unwrap_or(&Value::Bool(true))));
254            let hf_spm = !llama3 && bpe && truthy(m.get("byte_fallback")) && tj["decoder"]["type"].as_str() == Some("Sequence");
255            if hf_spm {
256                return Err(invalid(format!(
257                    "{}/tokenizer.json: a byte-fallback vocabulary the converter reads through LlamaHfVocab, which this reader does not read",
258                    dir.display()
259                )));
260            }
261        }
262        gpt2(dir, h, w, &[])?;
263    }
264    let vocab_size = h.u64(&["vocab_size"]).unwrap_or(32000);
265    if vocab_size == 32016 {
266        let mut sv = SpecialVocab::load(dir, false, None, &["prefix", "suffix", "middle", "eot"])?;
267        for (typ, id) in [("prefix", 32007), ("suffix", 32008), ("middle", 32009), ("eot", 32010)] {
268            sv.set_id(typ, &Value::from(id))?;
269        }
270        sv.write(w);
271    }
272    if vocab_size == 49152 {
273        w.bool("tokenizer.ggml.add_bos_token", false);
274    }
275    Ok(())
276}
277
278/// One `SentencePiece` of a `ModelProto`: its piece, score and type.
279struct Piece {
280    text: Vec<u8>,
281    score: f32,
282    kind: u64,
283}
284
285/// The pieces of a serialized SentencePiece `ModelProto` (field 1; each a piece, a score and a type, `NORMAL` when
286/// unstated).
287fn sentencepiece_pieces(bytes: &[u8]) -> io::Result<Vec<Piece>> {
288    fn varint(b: &[u8], at: &mut usize) -> io::Result<u64> {
289        let mut v = 0u64;
290        for shift in (0..64).step_by(7) {
291            let byte = *b.get(*at).ok_or_else(|| invalid("tokenizer.model ends inside a varint"))?;
292            *at += 1;
293            v |= u64::from(byte & 0x7f) << shift;
294            if byte & 0x80 == 0 {
295                return Ok(v);
296            }
297        }
298        Err(invalid("tokenizer.model holds an overlong varint"))
299    }
300    /// The next field of a message: its number, wire type and, for a length-delimited field, its bytes.
301    fn field<'a>(b: &'a [u8], at: &mut usize) -> io::Result<(u64, u64, &'a [u8], u64)> {
302        let key = varint(b, at)?;
303        let (num, wire) = (key >> 3, key & 7);
304        let (bytes, value) = match wire {
305            0 => (&b[0..0], varint(b, at)?),
306            1 | 5 => {
307                let n = if wire == 1 { 8 } else { 4 };
308                let s = b.get(*at..*at + n).ok_or_else(|| invalid("tokenizer.model ends inside a fixed field"))?;
309                *at += n;
310                (s, 0)
311            }
312            2 => {
313                let n = varint(b, at)? as usize;
314                let s = b.get(*at..*at + n).ok_or_else(|| invalid("tokenizer.model ends inside a field"))?;
315                *at += n;
316                (s, 0)
317            }
318            w => return Err(invalid(format!("tokenizer.model: wire type {w}"))),
319        };
320        Ok((num, wire, bytes, value))
321    }
322    let mut pieces = Vec::new();
323    let mut at = 0;
324    while at < bytes.len() {
325        let (num, wire, body, _) = field(bytes, &mut at)?;
326        if num != 1 || wire != 2 {
327            continue;
328        }
329        let mut p = Piece { text: Vec::new(), score: 0.0, kind: 1 };
330        let mut i = 0;
331        while i < body.len() {
332            let (n, w, b, v) = field(body, &mut i)?;
333            match (n, w) {
334                (1, 2) => p.text = b.to_vec(),
335                (2, 5) => p.score = f32::from_le_bytes(b.try_into().unwrap()),
336                (3, 0) => p.kind = v,
337                _ => {}
338            }
339        }
340        pieces.push(p);
341    }
342    Ok(pieces)
343}
344
345/// `_set_vocab_sentencepiece`: the pieces of `tokenizer.model` (an id past the model's own vocabulary a `[PAD<id>]`),
346/// the added tokens of `added_tokens.json` and `tokenizer_config.json` over them, and the special tokens.
347fn sentencepiece(dir: &Path, h: &Hparams, w: &mut Keys) -> io::Result<()> {
348    let path = dir.join("tokenizer.model");
349    let pieces = sentencepiece_pieces(&std::fs::read(&path)?)?;
350    let vocab_size = match h.u64(&["vocab_size_per_layer_input", "vocab_size"]) {
351        Some(n) if n > 0 => n as usize,
352        _ => pieces.len(),
353    };
354    let mut tokens: Vec<Vec<u8>> = (0..vocab_size).map(|i| format!("[PAD{i}]").into_bytes()).collect();
355    let mut scores = vec![-10000.0f32; vocab_size];
356    let mut types = vec![UNUSED; vocab_size];
357    for (id, p) in pieces.iter().enumerate().take(vocab_size) {
358        tokens[id] = p.text.clone();
359        scores[id] = p.score;
360        types[id] = match p.kind {
361            2 => UNKNOWN,
362            3 => CONTROL,
363            5 => UNUSED,
364            6 => BYTE,
365            _ => NORMAL,
366        };
367    }
368    if let Some(Value::Object(added)) = read_json(&dir.join("added_tokens.json"))? {
369        for (text, id) in added {
370            let id = id.as_u64().ok_or_else(|| invalid(format!("added_tokens.json: {text} has no id")))? as usize;
371            if id < vocab_size {
372                tokens[id] = text.into_bytes();
373                scores[id] = -1000.0;
374                types[id] = USER_DEFINED;
375            }
376        }
377    }
378    if let Some(tc) = read_json(&dir.join("tokenizer_config.json"))? {
379        if let Some(Value::Object(decoder)) = tc.get("added_tokens_decoder") {
380            for (id, data) in decoder {
381                let id: usize = id.parse().map_err(|_| invalid(format!("tokenizer_config.json: added token id `{id}`")))?;
382                if id >= vocab_size {
383                    continue;
384                }
385                let content = data["content"].as_str().ok_or_else(|| invalid(format!("tokenizer_config.json: added token {id} has no content")))?;
386                let text = if truthy(data.get("special")) || looks_special(content) {
387                    types[id] = CONTROL;
388                    content.to_string()
389                } else {
390                    types[id] = USER_DEFINED;
391                    content.replace('\u{2581}', " ")
392                };
393                scores[id] = -1000.0;
394                tokens[id] = text.into_bytes();
395            }
396        }
397    }
398    w.str("tokenizer.ggml.model", "llama");
399    w.str("tokenizer.ggml.pre", "default");
400    w.strings("tokenizer.ggml.tokens", &tokens);
401    w.put("tokenizer.ggml.scores", MetaValue::Array { ty: FLOAT32, items: scores.into_iter().map(MetaValue::F32).collect() });
402    w.put("tokenizer.ggml.token_type", MetaValue::Array { ty: INT32, items: types.into_iter().map(MetaValue::I32).collect() });
403    SpecialVocab::load(dir, false, Some(tokens.len()), &SPECIAL_TYPES)?.write(w);
404    Ok(())
405}
406
407/// An added token as the fast tokenizer holds it.
408struct Added {
409    content: String,
410    special: bool,
411    normalized: bool,
412}
413
414/// `_set_vocab_gpt2`: the byte-pair vocabulary of `tokenizer.json` (`get_vocab_base`), its pre-tokenizer named by
415/// the probe digest, and the special tokens and merges. `named` are the converter's `_set_special_token` calls after the
416/// files are read: each kind's id the added token of that text, where the files set no id of the kind.
417fn gpt2(dir: &Path, h: &Hparams, w: &mut Keys, named: &[(&str, &str)]) -> io::Result<()> {
418    let tj_path = dir.join("tokenizer.json");
419    let tj = read_json(&tj_path)?
420        .ok_or_else(|| invalid(format!("{}: no tokenizer.json (a vocabulary only a slow tokenizer reads)", dir.display())))?;
421    let tc = read_json(&dir.join("tokenizer_config.json"))?;
422    let model = &tj["model"];
423    if model["type"].as_str() != Some("BPE") {
424        return Err(invalid(format!("{}: a `{}` model where a byte-pair vocabulary is read", tj_path.display(), model["type"])));
425    }
426    let mut added: BTreeMap<u64, Added> = BTreeMap::new();
427    for a in tj["added_tokens"].as_array().into_iter().flatten() {
428        let id = a["id"].as_u64().ok_or_else(|| invalid(format!("{}: an added token without an id", tj_path.display())))?;
429        let content = a["content"].as_str().unwrap_or_default().to_string();
430        added.insert(id, Added { content, special: truthy(a.get("special")), normalized: truthy(a.get("normalized")) });
431    }
432    if let Some(Value::Object(decoder)) = tc.as_ref().and_then(|c| c.get("added_tokens_decoder")) {
433        for (id, d) in decoder {
434            let Ok(id) = id.parse::<u64>() else { continue };
435            added.entry(id).or_insert_with(|| Added {
436                content: d["content"].as_str().unwrap_or_default().to_string(),
437                special: truthy(d.get("special")),
438                normalized: truthy(d.get("normalized")),
439            });
440        }
441    }
442    // `tokenizer.vocab`: the model's vocabulary with the added tokens over it, keyed by text.
443    let mut vocab: HashMap<String, u64> = HashMap::new();
444    for (text, id) in model["vocab"].as_object().into_iter().flatten() {
445        vocab.insert(text.clone(), id.as_u64().ok_or_else(|| invalid(format!("{}: `{text}` has no id", tj_path.display())))?);
446    }
447    let added_vocab: HashMap<&str, u64> = added.iter().map(|(&id, a)| (a.content.as_str(), id)).collect();
448    for (&text, &id) in &added_vocab {
449        vocab.insert(text.to_string(), id);
450    }
451    let vocab_size = h.u64(&["vocab_size"]).unwrap_or(vocab.len() as u64);
452    if let Some(top) = vocab.values().max().filter(|&&m| m >= vocab_size) {
453        return Err(invalid(format!("{}: token id {top} past the vocabulary of {vocab_size}", tj_path.display())));
454    }
455    let reverse: HashMap<u64, &str> = vocab.iter().map(|(t, &id)| (id, t.as_str())).collect();
456    let clean_up = tc.as_ref().and_then(|c| c.get("clean_up_tokenization_spaces")).and_then(Value::as_bool).unwrap_or(false);
457    let mut tokens = Vec::with_capacity(vocab_size as usize);
458    let mut types = Vec::with_capacity(vocab_size as usize);
459    for i in 0..vocab_size {
460        let Some(&text) = reverse.get(&i) else {
461            tokens.push(format!("[PAD{i}]").into_bytes());
462            types.push(UNUSED);
463            continue;
464        };
465        let mut token = text.to_string();
466        if added_vocab.contains_key(text) {
467            let a = added.get(&i).ok_or_else(|| invalid(format!("{}: `{text}` is an added token under another id", tj_path.display())))?;
468            if !a.normalized && clean_up {
469                // The converter writes `decode(encode(token))`: the added token itself, with the decoder's clean-up.
470                token = clean_up_tokenization(&token);
471            }
472            if a.special || looks_special(&token) {
473                types.push(CONTROL);
474            } else {
475                token = token.replace('\u{2581}', " ");
476                types.push(USER_DEFINED);
477            }
478        } else {
479            types.push(NORMAL);
480        }
481        tokens.push(token.into_bytes());
482    }
483    let mut specials = SpecialVocab::load(dir, true, None, &SPECIAL_TYPES)?;
484    for &(typ, text) in named {
485        let id = added_vocab
486            .get(text)
487            .ok_or_else(|| invalid(format!("{}: no added token `{text}`, the converter's {typ} token", tj_path.display())))?;
488        specials.set_id(typ, &Value::from(*id))?;
489    }
490    let pre = pre_tokenizer_name(&tj, &tj_path)?;
491
492    w.str("tokenizer.ggml.model", "gpt2");
493    w.str("tokenizer.ggml.pre", &pre);
494    w.strings("tokenizer.ggml.tokens", &tokens);
495    w.put("tokenizer.ggml.token_type", MetaValue::Array { ty: INT32, items: types.into_iter().map(MetaValue::I32).collect() });
496    specials.write(w);
497    Ok(())
498}
499
500/// `DeepseekV2Model.set_vocab`'s Kimi K2 path: the tiktoken vocabulary written as a byte-pair one — each mergeable token
501/// in GPT-2's byte alphabet with its merge, the reserved special tokens after them, `[PAD<id>]` for an id neither names
502/// — named by the probe digest of the tokenizer's own encoding, which must be Kimi K2's; the special-token ids and
503/// template as `SpecialVocab` reads them without merges.
504fn kimi(dir: &Path, h: &Hparams, w: &mut Keys) -> io::Result<()> {
505    let t = Tiktoken::open(dir)?;
506    let ids = t.encode(PRE_TOKENIZER_PROBE);
507    let listed = format!("[{}]", ids.iter().map(u32::to_string).collect::<Vec<_>>().join(", "));
508    let digest: String = sha256(listed.as_bytes()).iter().map(|b| format!("{b:02x}")).collect();
509    let pre = pre_tokenizer_named(&digest)
510        .ok_or_else(|| invalid(format!("{}: a tiktoken pre-tokenizer the converter does not recognize (probe digest {digest})", dir.display())))?;
511    if pre != "kimi-k2" {
512        return Err(invalid(format!("{}: a tiktoken vocabulary whose pre-tokenizer is `{pre}`; the converter writes only kimi-k2's", dir.display())));
513    }
514    // `{**vocab, **special_tokens}`: each text keeps its first place and takes its last id; the reverse map, in that
515    // order, keeps the last text of an id.
516    let mut merges: Vec<Vec<u8>> = Vec::new();
517    let mut combined: Vec<(String, u32)> = Vec::new();
518    let mut place: HashMap<String, usize> = HashMap::new();
519    let mut put = |text: String, id: u32, combined: &mut Vec<(String, u32)>| match place.get(&text) {
520        Some(&i) => combined[i].1 = id,
521        None => {
522            place.insert(text.clone(), combined.len());
523            combined.push((text, id));
524        }
525    };
526    for (bytes, rank) in t.tokens() {
527        put(byte_level(bytes), *rank, &mut combined);
528        if bytes.len() > 1 {
529            if let Some((a, b)) = t.merge_of(bytes, *rank) {
530                merges.push(format!("{} {}", byte_level(&a), byte_level(&b)).into_bytes());
531            }
532        }
533    }
534    for (text, id) in t.specials() {
535        put(text.clone(), *id, &mut combined);
536    }
537    let special: std::collections::HashSet<u32> = t.specials().iter().map(|&(_, id)| id).collect();
538    let mut reverse: HashMap<u32, &str> = HashMap::new();
539    for (text, id) in &combined {
540        reverse.insert(*id, text.as_str());
541    }
542    let vocab_size = h.u64(&["vocab_size"]).ok_or_else(|| invalid("config.json: no vocab_size, which the Kimi vocabulary reads"))?;
543    let (mut tokens, mut types) = (Vec::with_capacity(vocab_size as usize), Vec::with_capacity(vocab_size as usize));
544    for i in 0..vocab_size as u32 {
545        match reverse.get(&i) {
546            None => {
547                tokens.push(format!("[PAD{i}]").into_bytes());
548                types.push(UNUSED);
549            }
550            Some(text) => {
551                tokens.push(text.as_bytes().to_vec());
552                types.push(if special.contains(&i) { CONTROL } else { NORMAL });
553            }
554        }
555    }
556    w.str("tokenizer.ggml.model", "gpt2");
557    w.str("tokenizer.ggml.pre", pre);
558    w.strings("tokenizer.ggml.tokens", &tokens);
559    w.put("tokenizer.ggml.token_type", MetaValue::Array { ty: INT32, items: types.into_iter().map(MetaValue::I32).collect() });
560    w.strings("tokenizer.ggml.merges", &merges);
561    SpecialVocab::load(dir, false, None, &SPECIAL_TYPES)?.write(w);
562    Ok(())
563}
564
565/// `PreTrainedTokenizerBase.clean_up_tokenization`.
566fn clean_up_tokenization(s: &str) -> String {
567    s.replace(" .", ".")
568        .replace(" ?", "?")
569        .replace(" !", "!")
570        .replace(" ,", ",")
571        .replace(" ' ", "'")
572        .replace(" n't", "n't")
573        .replace(" 'm", "'m")
574        .replace(" 's", "'s")
575        .replace(" 've", "'ve")
576        .replace(" 're", "'re")
577}
578
579/// `get_vocab_base_pre`: the name the converter's table gives the digest of the probe's ids, the probe encoded as the
580/// Hugging Face tokenizer encodes it (its special tokens added) and the ids written as a Python list.
581fn pre_tokenizer_name(tj: &Value, path: &Path) -> io::Result<String> {
582    let tokenizer = HfTokenizer::from_json(tj).map_err(|e| invalid(format!("{}: {e}", path.display())))?;
583    let ids = tokenizer.encode(PRE_TOKENIZER_PROBE, true);
584    let listed = format!("[{}]", ids.iter().map(u32::to_string).collect::<Vec<_>>().join(", "));
585    let digest: String = sha256(listed.as_bytes()).iter().map(|b| format!("{b:02x}")).collect();
586    pre_tokenizer_named(&digest)
587        .map(str::to_string)
588        .ok_or_else(|| invalid(format!("{}: the byte-pair pre-tokenizer is not one the converter recognizes (probe digest {digest})", path.display())))
589}
590
591/// The pre-tokenizer the converter's table names for a probe digest: its last entry for the digest, as the converter's
592/// chain of `if`s leaves it when a digest is listed twice.
593pub fn pre_tokenizer_named(digest: &str) -> Option<&'static str> {
594    PRE_TOKENIZER_DIGESTS.iter().rev().find(|(d, _)| *d == digest).map(|(_, name)| *name)
595}
596
597/// A chat template as `tokenizer_config.json` or the template files state it: one, or several by name.
598enum Template {
599    One(String),
600    Named(Vec<(String, Option<String>)>),
601}
602
603/// `gguf.SpecialVocab`: the merges, special-token ids, added-token flags, chat template and normalizer flags read
604/// from `tokenizer.json`, `tokenizer_config.json`, `config.json` and `merges.txt`.
605struct SpecialVocab {
606    n_vocab: Option<usize>,
607    types: Vec<String>,
608    merges: Vec<Vec<u8>>,
609    ids: Vec<(String, u32)>,
610    add: Vec<(String, bool)>,
611    chat_template: Option<Template>,
612    lowercase: Option<bool>,
613    strip_accents: Option<bool>,
614}
615
616impl SpecialVocab {
617    fn load(dir: &Path, load_merges: bool, n_vocab: Option<usize>, types: &[&str]) -> io::Result<SpecialVocab> {
618        let mut sv = SpecialVocab {
619            n_vocab,
620            types: types.iter().map(|t| t.to_string()).collect(),
621            merges: Vec::new(),
622            ids: Vec::new(),
623            add: Vec::new(),
624            chat_template: None,
625            lowercase: None,
626            strip_accents: None,
627        };
628        sv.from_tokenizer_json(dir, load_merges)?;
629        sv.from_config_json(dir)?;
630        if load_merges && sv.merges.is_empty() {
631            sv.from_merges_txt(dir)?;
632        }
633        Ok(sv)
634    }
635
636    /// `_set_special_token`: the first id a kind is given stands; an id past the vocabulary is skipped.
637    fn set_id(&mut self, typ: &str, tid: &Value) -> io::Result<()> {
638        let tid = match tid {
639            Value::Bool(b) => i64::from(*b),
640            Value::Number(n) if n.is_i64() || n.is_u64() => n.as_i64().unwrap_or(i64::MAX),
641            _ => return Ok(()),
642        };
643        if tid < 0 {
644            return Err(invalid(format!("invalid value for special token type {typ}: {tid}")));
645        }
646        if self.n_vocab.is_none_or(|n| (tid as u64) < n as u64) && !self.ids.iter().any(|(t, _)| t == typ) {
647            self.ids.push((typ.to_string(), tid as u32));
648        }
649        Ok(())
650    }
651
652    fn set_add(&mut self, typ: &str, v: bool) {
653        match self.add.iter_mut().find(|(t, _)| t == typ) {
654            Some(e) => e.1 = v,
655            None => self.add.push((typ.to_string(), v)),
656        }
657    }
658
659    fn add_default(&mut self, typ: &str, v: bool) {
660        if !self.add.iter().any(|(t, _)| t == typ) {
661            self.add.push((typ.to_string(), v));
662        }
663    }
664
665    fn has_add(&self, typ: &str) -> bool {
666        self.add.iter().any(|(t, _)| t == typ)
667    }
668
669    fn parse_normalizer(&mut self, n: &Value) {
670        match n["type"].as_str() {
671            Some("Lowercase") => self.lowercase = Some(true),
672            Some("StripAccents") => self.strip_accents = Some(true),
673            Some("BertNormalizer") => {
674                if let Some(v) = n.get("lowercase") {
675                    self.lowercase = v.as_bool();
676                }
677                if let Some(v) = n.get("strip_accents") {
678                    self.strip_accents = v.as_bool();
679                }
680            }
681            Some("Sequence") => {
682                for m in n["normalizers"].as_array().into_iter().flatten() {
683                    self.parse_normalizer(m);
684                }
685            }
686            _ => {}
687        }
688    }
689
690    /// `_try_load_from_tokenizer_json`.
691    fn from_tokenizer_json(&mut self, dir: &Path, load_merges: bool) -> io::Result<()> {
692        let tj_path = dir.join("tokenizer.json");
693        let tokenizer = read_json(&tj_path)?;
694        let empty = Vec::new();
695        let mut added_tokens: &Vec<Value> = &empty;
696        if let Some(t) = &tokenizer {
697            if load_merges {
698                if let Some(merges) = t["model"]["merges"].as_array().filter(|m| !m.is_empty()) {
699                    if merges[0].is_string() {
700                        self.merges = merges.iter().map(|m| m.as_str().unwrap_or_default().as_bytes().to_vec()).collect();
701                    } else if merges[0].as_array().is_some_and(|p| p.len() == 2 && p[0].is_string()) {
702                        // A pair's own spaces are written as U+0120, as the converter writes them.
703                        self.merges = merges
704                            .iter()
705                            .map(|pair| {
706                                pair.as_array()
707                                    .into_iter()
708                                    .flatten()
709                                    .map(|part| part.as_str().unwrap_or_default().replace(' ', "\u{120}"))
710                                    .collect::<Vec<_>>()
711                                    .join(" ")
712                                    .into_bytes()
713                            })
714                            .collect();
715                    } else {
716                        return Err(invalid(format!("{}: an unknown merges format", tj_path.display())));
717                    }
718                }
719            }
720            if truthy(t.get("normalizer")) {
721                self.parse_normalizer(&t["normalizer"]);
722            }
723            if let Some(a) = t["added_tokens"].as_array() {
724                added_tokens = a;
725            }
726        }
727        let tc_path = dir.join("tokenizer_config.json");
728        let mut tc: Option<Map<String, Value>> = match read_json(&tc_path)? {
729            Some(Value::Object(o)) => Some(o),
730            Some(_) => return Err(invalid(format!("{}: not an object", tc_path.display()))),
731            None => None,
732        };
733        let tc_truthy = |tc: &Option<Map<String, Value>>| tc.as_ref().is_some_and(|c| !c.is_empty());
734        if truthy(tokenizer.as_ref()) {
735            let t = tokenizer.as_ref().unwrap();
736            let get = |tc: &Option<Map<String, Value>>, k: &str| tc.as_ref().and_then(|c| c.get(k)).cloned();
737            let mut special_bos = get(&tc, "bos_token");
738            let mut special_cls = get(&tc, "cls_token");
739            let mut special_eos = get(&tc, "eos_token");
740            let mut special_sep = get(&tc, "sep_token");
741            if !truthy(special_bos.as_ref()) && truthy(special_cls.as_ref()) && tc_truthy(&tc) {
742                special_bos = special_cls.clone();
743                tc.as_mut().unwrap().insert("bos_token".into(), special_bos.clone().unwrap());
744            }
745            if !truthy(special_eos.as_ref()) && truthy(special_sep.as_ref()) && tc_truthy(&tc) {
746                special_eos = special_sep.clone();
747                tc.as_mut().unwrap().insert("eos_token".into(), special_eos.clone().unwrap());
748            }
749            if truthy(t.get("post_processor")) {
750                let pp = &t["post_processor"];
751                let processors: Vec<Value> = match pp.get("processors") {
752                    Some(Value::Array(a)) => a.clone(),
753                    _ => vec![pp.clone()],
754                };
755                for processor in &processors {
756                    match processor["type"].as_str() {
757                        Some("ByteLevel") => {
758                            self.add_default("bos", false);
759                            self.add_default("eos", false);
760                        }
761                        Some("RobertaProcessing") => {
762                            self.set_add("bos", true);
763                            self.set_add("eos", true);
764                            self.set_add("sep", true);
765                            if !truthy(special_cls.as_ref()) && tc_truthy(&tc) {
766                                special_cls = Some(processor.get("cls").and_then(|c| c.get(0)).cloned().unwrap_or(special_bos.clone().unwrap_or(Value::Null)));
767                                tc.as_mut().unwrap().insert("cls_token".into(), special_cls.clone().unwrap());
768                            }
769                            if !truthy(special_sep.as_ref()) && tc_truthy(&tc) {
770                                special_sep = Some(processor.get("sep").and_then(|c| c.get(0)).cloned().unwrap_or(special_eos.clone().unwrap_or(Value::Null)));
771                                tc.as_mut().unwrap().insert("sep_token".into(), special_sep.clone().unwrap());
772                            }
773                        }
774                        Some("TemplateProcessing") => {
775                            let single = processor["single"].as_array().cloned().unwrap_or_default();
776                            let mut pair = processor["pair"].as_array().cloned().unwrap_or_default();
777                            let special_id = |v: &Value| v.get("SpecialToken").and_then(|s| s.get("id")).and_then(Value::as_str).filter(|s| !s.is_empty()).map(str::to_string);
778                            let mut special_first: Option<String> = None;
779                            let mut special_last: Option<String> = None;
780                            if single.len() > 1 {
781                                special_first = special_id(&single[0]);
782                                if let Some(first) = &special_first {
783                                    if !tc_truthy(&tc) {
784                                        special_bos = Some(Value::String(first.clone()));
785                                    } else if !is(&special_bos, first) && !is(&special_cls, first) {
786                                        if !truthy(special_bos.as_ref()) {
787                                            special_bos = Some(Value::String(first.clone()));
788                                            tc.as_mut().unwrap().insert("bos_token".into(), Value::String(first.clone()));
789                                        }
790                                        if !truthy(special_cls.as_ref()) {
791                                            special_cls = Some(Value::String(first.clone()));
792                                            tc.as_mut().unwrap().insert("cls_token".into(), Value::String(first.clone()));
793                                        }
794                                    }
795                                    let v = is(&special_bos, first) || is(&special_cls, first);
796                                    self.set_add("bos", v);
797                                }
798                                special_last = special_id(&single[single.len() - 1]);
799                                if let Some(last) = &special_last {
800                                    if !tc_truthy(&tc) {
801                                        special_eos = Some(Value::String(last.clone()));
802                                    } else if !is(&special_eos, last) {
803                                        if !self.types.iter().any(|t| t == "eot") {
804                                            self.types.push("eot".into());
805                                            tc.as_mut().unwrap().insert("eot_token".into(), special_eos.clone().unwrap_or(Value::Null));
806                                        } else if !self.types.iter().any(|t| t == "eom") {
807                                            self.types.push("eom".into());
808                                            tc.as_mut().unwrap().insert("eom_token".into(), special_eos.clone().unwrap_or(Value::Null));
809                                        }
810                                        special_eos = Some(Value::String(last.clone()));
811                                        tc.as_mut().unwrap().insert("eos_token".into(), Value::String(last.clone()));
812                                    }
813                                    let v = is(&special_eos, last);
814                                    self.set_add("eos", v);
815                                }
816                            }
817                            if !pair.is_empty() {
818                                let start = usize::from(special_first.is_some() && special_id(&pair[0]) == special_first);
819                                let stop = special_last.is_some() && special_id(&pair[pair.len() - 1]) == special_last;
820                                let end = if stop { pair.len().saturating_sub(1) } else { pair.len() };
821                                pair = if start <= end { pair[start..end].to_vec() } else { Vec::new() };
822                                if !pair.is_empty() {
823                                    let seq = |v: &Value| v.get("Sequence").and_then(|s| s.get("id")).and_then(Value::as_str).map(str::to_string);
824                                    let (a, b) = (seq(&pair[0]), seq(&pair[pair.len() - 1]));
825                                    if a.as_deref() == Some("A") && b.as_deref() == Some("B") && pair.len() > 2 {
826                                        let inner = &pair[1..pair.len() - 1];
827                                        let mut add_sep = false;
828                                        if let Some(entry) = special_id(&inner[0]) {
829                                            if (is(&special_sep, &entry) || is(&special_eos, &entry)) && special_last.is_none() {
830                                                add_sep = true;
831                                            }
832                                        }
833                                        if inner.len() == 2 {
834                                            if let Some(entry) = special_id(&inner[1]) {
835                                                if is(&special_sep, &entry) || is(&special_eos, &entry) {
836                                                    add_sep = true;
837                                                }
838                                            }
839                                        }
840                                        self.set_add("sep", add_sep);
841                                        if add_sep && !truthy(special_sep.as_ref()) && tc_truthy(&tc) {
842                                            tc.as_mut().unwrap().insert("sep_token".into(), special_eos.clone().unwrap_or(Value::Null));
843                                        }
844                                    }
845                                }
846                            }
847                        }
848                        _ => {}
849                    }
850                }
851            }
852        }
853        let Some(tc) = tc.filter(|c| !c.is_empty()) else { return Ok(()) };
854        let jinja = dir.join("chat_template.jinja");
855        let alt: Option<Template> = if jinja.is_file() {
856            let default = std::fs::read_to_string(&jinja)?;
857            let mut more: Vec<(String, Option<String>)> = Vec::new();
858            if let Ok(entries) = std::fs::read_dir(dir.join("additional_chat_templates")) {
859                for e in entries.flatten() {
860                    let p = e.path();
861                    if p.extension().is_some_and(|x| x == "jinja") {
862                        let name = p.file_stem().unwrap().to_string_lossy().to_string();
863                        more.push((name, Some(std::fs::read_to_string(&p)?)));
864                    }
865                }
866            }
867            if more.is_empty() {
868                Some(Template::One(default))
869            } else {
870                more.insert(0, ("default".into(), Some(default)));
871                Some(Template::Named(more))
872            }
873        } else {
874            match read_json(&dir.join("chat_template.json"))? {
875                Some(j) => template_of(j.get("chat_template")),
876                None => None,
877            }
878        };
879        self.chat_template = match tc.get("chat_template") {
880            Some(v) => template_of(Some(v)),
881            None => alt,
882        };
883        let types = self.types.clone();
884        for typ in &types {
885            if let Some(Value::Bool(b)) = tc.get(&format!("add_{typ}_token")) {
886                if !self.has_add(typ) {
887                    self.add.push((typ.clone(), *b));
888                }
889            }
890            let content = match tc.get(&format!("{typ}_token")) {
891                Some(Value::String(s)) => s.clone(),
892                Some(Value::Object(o)) => match o.get("content") {
893                    Some(Value::String(s)) => s.clone(),
894                    _ => continue,
895                },
896                _ => continue,
897            };
898            let id = added_tokens.iter().find(|a| a["content"].as_str() == Some(content.as_str())).and_then(|a| a.get("id")).cloned();
899            if let Some(id) = id {
900                self.set_id(typ, &id)?;
901            }
902        }
903        Ok(())
904    }
905
906    /// `_try_load_from_config_json`: `<kind>_token_id` of `config.json`, or of its `text_config`.
907    fn from_config_json(&mut self, dir: &Path) -> io::Result<()> {
908        let Some(config) = read_json(&dir.join("config.json"))? else { return Ok(()) };
909        for typ in self.types.clone() {
910            let key = format!("{typ}_token_id");
911            let mut id = config.get(&key).filter(|v| !v.is_null());
912            if id.is_none() {
913                id = config.get("text_config").and_then(|t| t.get(&key));
914            }
915            if let Some(id) = id {
916                self.set_id(&typ, &id.clone())?;
917            }
918        }
919        Ok(())
920    }
921
922    /// `_try_load_merges_txt`.
923    fn from_merges_txt(&mut self, dir: &Path) -> io::Result<()> {
924        let path = dir.join("merges.txt");
925        if !path.is_file() {
926            return Ok(());
927        }
928        let text = std::fs::read_to_string(&path)?;
929        let mut lines = text.lines().peekable();
930        if lines.peek().is_some_and(|l| l.trim().starts_with('#')) {
931            lines.next();
932        }
933        let mut merges = Vec::new();
934        for line in lines {
935            let parts: Vec<&str> = line.split_whitespace().collect();
936            if parts.len() == 2 {
937                merges.push(format!("{} {}", parts[0], parts[1]).into_bytes());
938            }
939        }
940        self.merges = merges;
941        Ok(())
942    }
943
944    /// `add_to_gguf`: every kind with a key of its own; a kind the writer has no key for is left out.
945    fn write(&self, w: &mut Keys) {
946        if !self.merges.is_empty() {
947            w.strings("tokenizer.ggml.merges", &self.merges);
948        }
949        for (typ, id) in &self.ids {
950            let key = match typ.as_str() {
951                "bos" => "tokenizer.ggml.bos_token_id",
952                "eos" => "tokenizer.ggml.eos_token_id",
953                "unk" => "tokenizer.ggml.unknown_token_id",
954                "sep" => "tokenizer.ggml.seperator_token_id",
955                "pad" => "tokenizer.ggml.padding_token_id",
956                "mask" => "tokenizer.ggml.mask_token_id",
957                "eot" => "tokenizer.ggml.eot_token_id",
958                "eom" => "tokenizer.ggml.eom_token_id",
959                _ => continue,
960            };
961            w.u32(key, *id);
962        }
963        for (typ, v) in &self.add {
964            let key = match typ.as_str() {
965                "bos" => "tokenizer.ggml.add_bos_token",
966                "eos" => "tokenizer.ggml.add_eos_token",
967                "sep" => "tokenizer.ggml.add_sep_token",
968                _ => continue,
969            };
970            w.bool(key, *v);
971        }
972        match &self.chat_template {
973            None => {}
974            Some(Template::One(t)) => w.str("tokenizer.chat_template", t),
975            Some(Template::Named(choices)) => {
976                let mut default = None;
977                let mut names: Vec<Vec<u8>> = Vec::new();
978                for (name, template) in choices {
979                    let name: String = name.chars().map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }).collect();
980                    let Some(template) = template else { continue };
981                    if name.is_empty() {
982                        continue;
983                    }
984                    if name == "default" {
985                        default = Some(template.clone());
986                    } else {
987                        if !names.contains(&name.as_bytes().to_vec()) {
988                            names.push(name.as_bytes().to_vec());
989                        }
990                        w.str(&format!("tokenizer.chat_template.{name}"), template);
991                    }
992                }
993                if !names.is_empty() {
994                    w.strings("tokenizer.chat_templates", &names);
995                }
996                if let Some(t) = default {
997                    w.str("tokenizer.chat_template", &t);
998                }
999            }
1000        }
1001        if let Some(b) = self.lowercase {
1002            w.bool("tokenizer.ggml.normalizer.lowercase", b);
1003        }
1004        if let Some(b) = self.strip_accents {
1005            w.bool("tokenizer.ggml.normalizer.strip_accents", b);
1006        }
1007    }
1008}
1009
1010/// A `chat_template` value: a string, or a list of `{name, template}` choices; anything else is ignored.
1011fn template_of(v: Option<&Value>) -> Option<Template> {
1012    match v {
1013        Some(Value::String(s)) => Some(Template::One(s.clone())),
1014        Some(Value::Array(a)) => Some(Template::Named(
1015            a.iter()
1016                .map(|c| (c["name"].as_str().unwrap_or_default().to_string(), c["template"].as_str().map(str::to_string)))
1017                .collect(),
1018        )),
1019        _ => None,
1020    }
1021}