ferrox_models/tokenizer.rs
1//! A real, reversible byte-level tokenizer: each UTF-8 byte maps to
2//! token id `byte as u32` (vocabulary 0..256). This is not a full
3//! BPE/tokenizer.json implementation -- GLM-5.2, DeepSeek V4 Pro, and
4//! Kimi K3 each ship their own trained BPE vocabulary alongside their
5//! weights, and none of those vocab files are guessable or available in
6//! this environment (see docs/MODELS.md) -- but unlike the
7//! previous placeholder (`byte % vocab_size`, which was lossy and could
8//! not decode back to the original text), this tokenizer is exact and
9//! round-trips perfectly. It is the honest "smallest real thing that
10//! works" rather than a fake stand-in.
11//!
12//! Loading a real BPE merge table from a GGUF file's
13//! `tokenizer.ggml.tokens` / `tokenizer.ggml.merges` metadata arrays
14//! (see `ferrox-gguf`'s `GgufValue::Array` support, already verified
15//! against a real downloaded llama.cpp vocab fixture) was the natural
16//! next step and now exists below (`GgufBpeTokenizer`,
17//! `GgufSpmTokenizer`, `GgufUnigramTokenizer`).
18//!
19//! The per-checkpoint pre-tokenization rules live next door in
20//! [`pretokenize`], which is a transcription of llama.cpp and is
21//! reviewed against it.
22//!
23//! Four of llama.cpp's six `tokenizer.ggml.model` values are covered:
24//! `gpt2`/`gemma4` by [`GgufBpeTokenizer`], `llama` by
25//! [`GgufSpmTokenizer`], `t5` by [`GgufUnigramTokenizer`], and `bert` by
26//! [`GgufWordPieceTokenizer`] in `wordpiece`, which brings its own
27//! normalizer and its own Unicode tables (`unicode`, `unicode_data`)
28//! because WordPiece does not use the pre-tokenizer regexes at all.
29//! Still missing: `rwkv`, which needs a trie tokenizer, and `none`.
30
31mod pretokenize;
32mod unicode;
33mod unicode_data;
34mod wordpiece;
35
36pub use wordpiece::{GgufWordPieceTokenizer, NormalizerOptions};
37
38/// The `tokenizer.ggml.pre` values whose llama.cpp arm sets
39/// `add_bos = true` for a BPE vocabulary.
40///
41/// Transcribed from `.scratch/llama.cpp/src/llama-vocab.cpp`: the
42/// `LLAMA_VOCAB_PRE_TYPE_LLAMA3` arm sets it for the whole llama3 group
43/// in one statement, and `tekken` and `chameleon` set it in arms of
44/// their own. Llama-3.x GGUFs ship no explicit
45/// `tokenizer.ggml.add_bos_token`, so leaving the group out made every
46/// raw completion prompt one `<|begin_of_text|>` short of llama.cpp's.
47const ADD_BOS_PRE: &[&str] = &[
48 // LLAMA_VOCAB_PRE_TYPE_LLAMA3
49 "llama3",
50 "llama-v3",
51 "llama-bpe",
52 "falcon3",
53 "falcon-h1",
54 "pixtral",
55 "midm-2.0",
56 "lfm2",
57 "jina-v5-nano",
58 // arms of their own, same flag
59 "tekken",
60 "chameleon",
61];
62
63/// Whether prompt encoding should prepend the GGUF BOS token.
64///
65/// Port of llama.cpp `llama_vocab` add_bos defaults
66/// (`.scratch/llama.cpp/src/llama-vocab.cpp`): explicit
67/// `tokenizer.ggml.add_bos_token` wins; else SPM → true, BPE → false
68/// unless the checkpoint's `pre` is one of [`ADD_BOS_PRE`]. Qwen2-MoE
69/// ships `bos_token_id=<|endoftext|>` but `add_bos=false` — always
70/// prepending that token poisons greedy decode.
71pub fn should_add_bos_token(file: &impl ferrox_gguf::TensorSource) -> bool {
72 if let Some(v) = file.metadata_bool("tokenizer.ggml.add_bos_token") {
73 return v;
74 }
75 let model = file.metadata_str("tokenizer.ggml.model").unwrap_or("");
76 let pre = file.metadata_str("tokenizer.ggml.pre").unwrap_or("");
77 // llama.cpp: SPM/WPM default add_bos=true; BPE defaults false unless
78 // its pre-tokenizer arm opts in. qwen2 leaves false.
79 // `bert` is WPM, whose upstream arm sets add_bos AND add_sep true.
80 // It was missing here, so every WordPiece prompt was one `[CLS]`
81 // short of llama.cpp's.
82 if matches!(model, "llama" | "spm" | "bert") || model.contains("sentencepiece") {
83 return true;
84 }
85 ADD_BOS_PRE.contains(&pre)
86}
87
88/// Prepends the checkpoint's BOS id to an already-encoded prompt, unless
89/// the prompt already starts with it.
90///
91/// # The rule, stated once
92///
93/// **The chat template owns BOS when it prints one; the loader owns it
94/// otherwise.** Which of the two happens is a property of the individual
95/// checkpoint, not of the family:
96///
97/// * Many upstream templates open with `{{ bos_token }}` — gemma-2/3
98/// (`<bos>`), Mistral-Instruct and TinyLlama (`<s>`), Llama-3
99/// (`<|begin_of_text|>`). Rendering one of those already puts BOS in
100/// the *text*, and both [`GgufBpeTokenizer::encode`] and
101/// [`GgufSpmTokenizer::encode`] split on special-token text first, so
102/// it comes back as the BOS *id* in position 0.
103/// * Unsloth deliberately **strips** `{{ bos_token }}` out of the
104/// templates it bakes into its GGUF exports, precisely so that a
105/// runtime which adds BOS itself does not double it. On those
106/// checkpoints the render carries no BOS and the loader must add it.
107///
108/// So neither "always add" nor "never add" is right, and a renderer
109/// cannot be sniffed for which case it is. This function implements the
110/// only rule that is correct for both: add the id, **idempotently**.
111/// `bos` is already the gated value — pass `None` when
112/// [`should_add_bos_token`] says this vocabulary does not take one
113/// (BPE/qwen2 ship a `bos_token_id` they never prepend).
114///
115/// Note this is *stricter* than llama.cpp, whose `add_special` path
116/// pushes BOS unconditionally and leaves the duplicate to a warning.
117/// Ferrox has no user-visible "you asked for two BOS tokens" surface, so
118/// it dedupes instead of warning.
119/// Generic over the id width because the CLI and server carry prompts as
120/// `Vec<usize>` and the tokenizers emit `Vec<u32>`.
121pub fn prepend_bos<T: Copy + PartialEq>(tokens: &mut Vec<T>, bos: Option<T>) {
122 let Some(bos) = bos else { return };
123 if tokens.first() != Some(&bos) {
124 tokens.insert(0, bos);
125 }
126}
127
128/// Token texts llama.cpp treats as end-of-generation regardless of what
129/// the metadata ids say (`llama-vocab.cpp`, the literal list right above
130/// its "sanity checks" block). Copied verbatim, including the comments
131/// naming which family each entry exists for, because the set is not
132/// derivable: it is a hand-maintained list of what real checkpoints ship.
133///
134/// Note `<|end|>` *is* here. The Unsloth study recorded in
135/// `docs/plans/llama-cpp-parity-push.md` claimed gpt-oss's `<|end|>` must
136/// not be EOG or every reply truncates; llama.cpp's own source says
137/// otherwise, and llama.cpp serves gpt-oss. Following the reference
138/// implementation, and flagging the claim as contradicted.
139const EOG_TOKEN_TEXTS: &[&str] = &[
140 "<|eot_id|>",
141 "<|im_end|>",
142 "<|end|>",
143 "<|return|>", // o200k_harmony
144 "<|call|>", // o200k_harmony
145 "<|flush|>", // solar-open
146 "<|calls|>", // solar-open
147 "<end_of_turn>",
148 "<|endoftext|>",
149 "</s>", // paddleocr
150 "<|eom_id|>",
151 "<EOT>",
152 "_<EOT>",
153 "[EOT]", // Kimi-K2
154 "[EOS]", // Kimi-K2
155 "<|end_of_text|>",
156 "<end_of_utterance>", // smoldocling
157 "<eos>", // gemma4
158 "<turn|>", // gemma4
159 "<|tool_response>", // gemma4
160 "<|end▁of▁sentence|>", // deepseek-ocr
161 "[e~[", // minimax-m2/m3
162];
163
164/// Every token id that ends generation, not just `eos_token_id`.
165///
166/// A single EOS id is wrong for most modern chat checkpoints: Llama-3
167/// ends turns with `<|eot_id|>` while its `eos_token_id` is
168/// `<|end_of_text|>`, and gemma-4 ends with `<turn|>`. Stopping only on
169/// the metadata EOS means the model keeps generating past the end of its
170/// turn and starts a new one — the "it answers, then interviews itself"
171/// failure.
172///
173/// Mirrors llama.cpp: the literal-name list above, plus the
174/// `eos`/`eot`/`eom` metadata ids, which it folds in with a warning when
175/// they were not already caught by name.
176pub fn eog_token_ids(file: &impl ferrox_gguf::TensorSource) -> std::collections::HashSet<u32> {
177 let mut out = std::collections::HashSet::new();
178 for key in [
179 "tokenizer.ggml.eos_token_id",
180 "tokenizer.ggml.eot_token_id",
181 "tokenizer.ggml.eom_token_id",
182 ] {
183 if let Some(id) = file.metadata_u64(key) {
184 out.insert(id as u32);
185 }
186 }
187 if let Some(ferrox_gguf::GgufValue::Array(items)) = file.metadata("tokenizer.ggml.tokens") {
188 for (id, v) in items.iter().enumerate() {
189 if let ferrox_gguf::GgufValue::String(text) = v {
190 if EOG_TOKEN_TEXTS.contains(&text.as_str()) {
191 out.insert(id as u32);
192 }
193 }
194 }
195 }
196 out
197}
198
199/// The set of token ids a decode loop must stop on, carried as one value
200/// so a caller cannot accidentally carry only half of it.
201///
202/// This type exists because `Option<usize>` was the shape of a real bug:
203/// every `ferrox-server` decode loop threaded a single `eos_id` from the
204/// loader to the sampler, so a Llama-3 or gemma checkpoint served over
205/// HTTP ran past `<|eot_id|>` / `<end_of_turn>` to `max_tokens` even
206/// after [`eog_token_ids`] landed for the CLI. Passing a `StopTokens`
207/// makes "I only have the metadata EOS" an explicit choice
208/// ([`StopTokens::from_eos`], for the synthetic-weights and Kimi paths
209/// that have no GGUF metadata to read) rather than the default.
210#[derive(Clone, Debug, Default)]
211pub struct StopTokens {
212 ids: std::collections::HashSet<u32>,
213}
214
215impl StopTokens {
216 /// Everything [`eog_token_ids`] finds in this checkpoint: the
217 /// `eos`/`eot`/`eom` metadata ids plus every vocabulary entry whose
218 /// text is on llama.cpp's literal EOG list.
219 pub fn from_gguf(file: &impl ferrox_gguf::TensorSource) -> Self {
220 Self {
221 ids: eog_token_ids(file),
222 }
223 }
224
225 /// Just the one id. For callers with no GGUF metadata behind them —
226 /// the synthetic random-weights demo model, and the Kimi checkpoint
227 /// directory whose tokenizer is a separate file format.
228 pub fn from_eos(eos: Option<usize>) -> Self {
229 Self {
230 ids: eos.map(|e| e as u32).into_iter().collect(),
231 }
232 }
233
234 /// For checkpoints whose vocabulary is not GGUF metadata — Kimi K3
235 /// ships a `tokenizer_config.json` with a name→id special-token map.
236 /// Folds in every entry whose *text* is on llama.cpp's EOG list, so
237 /// `[EOT]` stops a turn there exactly as it does in a GGUF.
238 pub fn from_special_tokens<'a>(specials: impl IntoIterator<Item = (&'a str, u32)>) -> Self {
239 Self {
240 ids: specials
241 .into_iter()
242 .filter(|(name, _)| EOG_TOKEN_TEXTS.contains(name))
243 .map(|(_, id)| id)
244 .collect(),
245 }
246 }
247
248 /// Folds one more id in — used to keep a metadata `eos_token_id` that
249 /// a vocabulary spells in a way the literal list does not know.
250 pub fn with_id(mut self, id: Option<usize>) -> Self {
251 if let Some(id) = id {
252 self.ids.insert(id as u32);
253 }
254 self
255 }
256
257 pub fn contains(&self, id: usize) -> bool {
258 u32::try_from(id).is_ok_and(|id| self.ids.contains(&id))
259 }
260
261 pub fn is_empty(&self) -> bool {
262 self.ids.is_empty()
263 }
264
265 pub fn len(&self) -> usize {
266 self.ids.len()
267 }
268}
269
270pub struct ByteTokenizer;
271
272impl ByteTokenizer {
273 pub fn encode(text: &str) -> Vec<u32> {
274 text.bytes().map(|b| b as u32).collect()
275 }
276
277 /// Decodes token ids back to a string. Ids outside 0..256 are
278 /// dropped rather than silently corrupting output; invalid UTF-8
279 /// byte sequences are replaced per Rust's standard lossy conversion.
280 pub fn decode(ids: &[u32]) -> String {
281 let bytes: Vec<u8> = ids.iter().filter_map(|&id| u8::try_from(id).ok()).collect();
282 String::from_utf8_lossy(&bytes).into_owned()
283 }
284
285 pub const VOCAB_SIZE: usize = 256;
286}
287
288/// How a GGUF BPE vocabulary remaps text before merge lookup.
289///
290/// GPT-2-style vocabs store merges in the OpenAI byte↔unicode remapped
291/// space; Gemma-4 (and similar SPM-flavoured BPE) stores merges over
292/// raw UTF-8 with spaces already escaped to U+2581 (`▁`).
293#[derive(Clone, Copy, Debug, PartialEq, Eq)]
294enum BpeEncodingStyle {
295 Gpt2,
296 /// llama.cpp `LLAMA_VOCAB_PRE_TYPE_GEMMA4`: escape `" "` → `▁`,
297 /// split only on newlines, merge on raw UTF-8 codepoints
298 /// (`byte_encode = false`).
299 SpmWhitespace,
300}
301
302/// Builds the GPT2 byte-to-unicode remap table: bytes in the "already
303/// printable, unambiguous" ranges (33..=126, 161..=172, 174..=255) map
304/// to themselves as Unicode codepoints; every other byte (control
305/// characters, space, and a few others that would be ambiguous or
306/// unprintable as raw codepoints) maps to a codepoint starting at 256.
307/// This is the exact algorithm from OpenAI's GPT-2 `encoder.py`
308/// `bytes_to_unicode()`, reimplemented independently in Rust: real BPE
309/// vocabularies (llama.cpp, mistral.rs via the `tokenizers` crate) list
310/// merge-table entries in *this* remapped space (e.g. "\u{0120}the",
311/// where the leading char is U+0120, the remapped space byte 0x20), not
312/// in raw byte or `char` space, so skipping this step -- which ferrox
313/// did before this function existed -- silently fails to match any real
314/// vocabulary's merge table on space- and control-byte-adjacent tokens.
315fn gpt2_byte_to_unicode() -> ([char; 256], std::collections::HashMap<char, u8>) {
316 let is_printable =
317 |b: u16| (33..=126).contains(&b) || (161..=172).contains(&b) || (174..=255).contains(&b);
318
319 let mut forward = ['\0'; 256];
320 let mut extra_offset = 0u32;
321 for b in 0..256u16 {
322 if is_printable(b) {
323 forward[b as usize] = char::from_u32(b as u32).unwrap();
324 } else {
325 forward[b as usize] = char::from_u32(256 + extra_offset).unwrap();
326 extra_offset += 1;
327 }
328 }
329
330 let mut reverse = std::collections::HashMap::with_capacity(256);
331 for (b, &c) in forward.iter().enumerate() {
332 reverse.insert(c, b as u8);
333 }
334 (forward, reverse)
335}
336
337/// The chunking that runs before BPE: raw text is cut into
338/// contractions, letter runs, digit runs, symbol runs and whitespace
339/// runs, and each chunk is merged separately. Without it `encode_word`
340/// would treat a whole sentence as one word and could merge across word
341/// boundaries in ways no real tokenizer does.
342///
343/// Which pattern a checkpoint gets, and what happens to the text
344/// between matches, is [`pretokenize`]'s job — it is a transcription of
345/// llama.cpp's `llama-vocab.cpp` and `unicode.cpp` and is reviewed
346/// against them.
347/// U+2581 FIGURE SPACE used by SentencePiece-style BPE merge tables.
348const SPM_SPACE: char = '\u{2581}';
349
350/// A real BPE tokenizer built from a GGUF file's own
351/// `tokenizer.ggml.tokens` / `tokenizer.ggml.merges` metadata arrays.
352/// Supports GPT-2 byte-remap BPE (`tokenizer.ggml.model == "gpt2"`) and
353/// Gemma-4 SPM-style BPE (`"gemma4"`: escape spaces to `▁`, merge on
354/// raw UTF-8, newline-only pre-split).
355///
356/// Verified against `tests/fixtures/llama-bpe-vocab.gguf` (GPT-2 path).
357/// See `crates/ferrox-models/tests/gguf_vocab.rs`.
358pub struct GgufBpeTokenizer {
359 token_to_id: std::collections::HashMap<String, u32>,
360 id_to_token: Vec<String>,
361 /// merge rank: lower = merges earlier (higher priority), matching
362 /// the standard BPE convention of applying the most-frequent
363 /// (lowest-rank) merge first.
364 merge_rank: std::collections::HashMap<(String, String), usize>,
365 byte_to_unicode: [char; 256],
366 unicode_to_byte: std::collections::HashMap<char, u8>,
367 /// Control/user-defined tokens (chat-template markers and similar
368 /// added special tokens) from `tokenizer.ggml.token_type`, matched
369 /// as atomic substrings before normal BPE runs -- see
370 /// `split_on_special_tokens`.
371 special_tokens: Vec<(String, u32)>,
372 /// Compiled pre-tokenization pattern (GPT-2 word regex, or
373 /// newline-only for Gemma-4 SPM-BPE).
374 pretokenize_pattern: fancy_regex::Regex,
375 style: BpeEncodingStyle,
376}
377
378#[derive(Debug, thiserror::Error)]
379pub enum TokenizerLoadError {
380 #[error("GGUF file has no 'tokenizer.ggml.tokens' metadata array")]
381 MissingTokens,
382 #[error("'tokenizer.ggml.tokens' is present but is not a string array")]
383 TokensNotStringArray,
384}
385
386/// GGUF's real `tokenizer.ggml.token_type` per-token integer tag,
387/// confirmed directly against llama.cpp's real `llama_token_type` enum
388/// (`include/llama.h`) and its real GGUF-loading code
389/// (`src/llama-vocab.cpp`'s `toktypes[i]` switch), not guessed: this is
390/// a plain sequential enum on disk (`1=NORMAL, 2=UNKNOWN, 3=CONTROL,
391/// 4=USER_DEFINED, 5=UNUSED, 6=BYTE`), a different and simpler
392/// representation than llama.cpp's own *internal* bit-flag
393/// `llama_token_attr` type, which is derived from this at load time,
394/// not what's actually stored in the file.
395const GGML_TOKEN_TYPE_CONTROL: i64 = 3;
396const GGML_TOKEN_TYPE_USER_DEFINED: i64 = 4;
397
398/// Reads `tokenizer.ggml.token_type` (if present) and returns the
399/// `(token_text, id)` pairs for every CONTROL or USER_DEFINED entry
400/// (chat-template markers like `<|user|>`/`<|assistant|>`, and similar
401/// added special tokens) -- these must be recognized as atomic
402/// vocabulary entries during encoding rather than shattered into
403/// ordinary BPE/SPM/Unigram pieces, matching real llama.cpp's
404/// `tokenizer_st_partition` behavior (`src/llama-vocab.cpp`): special
405/// tokens are located as literal substrings and carved out of the
406/// input *before* normal tokenization runs on what's left, not folded
407/// into the regular vocabulary-matching pass.
408fn load_special_tokens(
409 file: &impl ferrox_gguf::TensorSource,
410 id_to_token: &[String],
411) -> Vec<(String, u32)> {
412 let Some(ferrox_gguf::GgufValue::Array(items)) = file.metadata("tokenizer.ggml.token_type")
413 else {
414 return Vec::new();
415 };
416 items
417 .iter()
418 .zip(id_to_token.iter())
419 .enumerate()
420 .filter_map(|(id, (v, text))| {
421 let ty = match v {
422 ferrox_gguf::GgufValue::I32(t) => *t as i64,
423 ferrox_gguf::GgufValue::U32(t) => *t as i64,
424 _ => return None,
425 };
426 // A checkpoint that flags its markers is the easy case.
427 //
428 // Not every one does. Yi-1.5-6B-Chat lists `<|im_start|>`
429 // (id 6) and `<|im_end|>` (id 7) as NORMAL, so this filter
430 // dropped them, the splitter never saw them, and
431 // `<|im_end|>` tokenized as SIX ordinary BPE pieces instead
432 // of the single token the model was trained on. Its prompt
433 // was malformed and its turn marker could not be matched by
434 // id, so nothing stopped and the marker came back as text
435 // mid-answer.
436 //
437 // So a token that LOOKS like a marker and exists verbatim in
438 // the vocabulary is treated as one whatever its type says.
439 // The shape test is what keeps this narrow: ordinary words
440 // are in the vocabulary too, and promoting those would split
441 // real text.
442 (ty == GGML_TOKEN_TYPE_CONTROL
443 || ty == GGML_TOKEN_TYPE_USER_DEFINED
444 || looks_like_marker(text))
445 .then(|| (text.clone(), id as u32))
446 })
447 .collect()
448}
449
450/// Is this vocabulary entry shaped like a special marker?
451///
452/// Deliberately strict: `<|...|>` and `<...>` with no whitespace, at
453/// least three characters. That covers `<|im_end|>`, `<|endoftext|>`,
454/// `<end_of_turn>` and `</s>` while excluding ordinary text, which also
455/// lives in the vocabulary and must keep tokenizing normally.
456fn looks_like_marker(text: &str) -> bool {
457 let t = text.trim();
458 if t.len() < 3 || t.contains(char::is_whitespace) {
459 return false;
460 }
461 (t.starts_with("<|") && t.ends_with("|>")) || (t.starts_with('<') && t.ends_with('>'))
462}
463
464/// One chunk of `split_on_special_tokens`'s output: either a raw text
465/// run to tokenize normally, or an already-resolved special token id.
466enum TextOrSpecial<'a> {
467 Text(&'a str),
468 Special(u32),
469}
470
471/// Splits `text` around every literal occurrence of any of `specials`
472/// (longest-match-first on ties, matching real llama.cpp's
473/// `tokenizer_st_partition`), leaving the text runs between/around
474/// them untouched for the caller's normal tokenization pass. Returns
475/// the whole input as one `Text` chunk when `specials` is empty (the
476/// overwhelmingly common fast path, since most GGUF files carry no
477/// `tokenizer.ggml.token_type` metadata at all).
478fn split_on_special_tokens<'a>(
479 text: &'a str,
480 specials: &[(String, u32)],
481) -> Vec<TextOrSpecial<'a>> {
482 if specials.is_empty() {
483 return vec![TextOrSpecial::Text(text)];
484 }
485 let mut segments = Vec::new();
486 let mut pos = 0usize;
487 while pos < text.len() {
488 let mut best: Option<(usize, usize, u32)> = None; // (start, len, id)
489 for (s, id) in specials {
490 if s.is_empty() {
491 continue;
492 }
493 if let Some(rel) = text[pos..].find(s.as_str()) {
494 let start = pos + rel;
495 let len = s.len();
496 let better = match best {
497 None => true,
498 Some((bstart, blen, _)) => start < bstart || (start == bstart && len > blen),
499 };
500 if better {
501 best = Some((start, len, *id));
502 }
503 }
504 }
505 match best {
506 None => break,
507 Some((start, len, id)) => {
508 if start > pos {
509 segments.push(TextOrSpecial::Text(&text[pos..start]));
510 }
511 segments.push(TextOrSpecial::Special(id));
512 pos = start + len;
513 }
514 }
515 }
516 if pos < text.len() {
517 segments.push(TextOrSpecial::Text(&text[pos..]));
518 }
519 segments
520}
521
522#[cfg(test)]
523mod special_token_split_tests {
524 use super::*;
525
526 fn text_of<'a>(seg: &TextOrSpecial<'a>) -> Option<&'a str> {
527 match seg {
528 TextOrSpecial::Text(t) => Some(t),
529 TextOrSpecial::Special(_) => None,
530 }
531 }
532
533 #[test]
534 fn empty_specials_list_returns_the_whole_text_unsplit() {
535 let segs = split_on_special_tokens("hello world", &[]);
536 assert_eq!(segs.len(), 1);
537 assert_eq!(text_of(&segs[0]), Some("hello world"));
538 }
539
540 #[test]
541 fn splits_around_a_single_special_token_in_the_middle() {
542 let specials = vec![("<|user|>".to_string(), 42u32)];
543 let segs = split_on_special_tokens("before<|user|>after", &specials);
544 assert_eq!(segs.len(), 3);
545 assert_eq!(text_of(&segs[0]), Some("before"));
546 assert!(matches!(segs[1], TextOrSpecial::Special(42)));
547 assert_eq!(text_of(&segs[2]), Some("after"));
548 }
549
550 #[test]
551 fn multiple_occurrences_and_multiple_distinct_specials_all_split() {
552 let specials = vec![
553 ("<|user|>".to_string(), 1u32),
554 ("<|assistant|>".to_string(), 2u32),
555 ];
556 let segs = split_on_special_tokens("<|user|>hi<|assistant|>hello<|user|>bye", &specials);
557 let kinds: Vec<Option<&str>> = segs.iter().map(text_of).collect();
558 assert_eq!(
559 kinds,
560 vec![None, Some("hi"), None, Some("hello"), None, Some("bye")]
561 );
562 assert!(matches!(segs[0], TextOrSpecial::Special(1)));
563 assert!(matches!(segs[2], TextOrSpecial::Special(2)));
564 assert!(matches!(segs[4], TextOrSpecial::Special(1)));
565 }
566
567 #[test]
568 fn longest_match_wins_on_a_tied_start_position() {
569 // "<|user|>" and a hypothetical shorter overlapping prefix
570 // starting at the same position must prefer the longer match.
571 let specials = vec![("<|user|>".to_string(), 1u32), ("<|u".to_string(), 99u32)];
572 let segs = split_on_special_tokens("<|user|>x", &specials);
573 assert!(matches!(segs[0], TextOrSpecial::Special(1)));
574 }
575
576 #[test]
577 fn no_match_at_all_returns_the_whole_text_as_one_segment() {
578 let specials = vec![("<|user|>".to_string(), 1u32)];
579 let segs = split_on_special_tokens("plain text with no specials", &specials);
580 assert_eq!(segs.len(), 1);
581 assert_eq!(text_of(&segs[0]), Some("plain text with no specials"));
582 }
583}
584
585impl GgufBpeTokenizer {
586 /// Loads the vocabulary + merge table from a GGUF file's metadata.
587 /// Merges are optional (some tokenizer types, e.g. byte-level
588 /// unigram, don't use them); if absent, encoding falls back to
589 /// per-byte token lookup. `tokenizer.ggml.model == "gemma4"` selects
590 /// SPM-whitespace BPE; everything else with merges uses GPT-2 style.
591 pub fn from_gguf(file: &impl ferrox_gguf::TensorSource) -> Result<Self, TokenizerLoadError> {
592 let tokens_value = file
593 .metadata("tokenizer.ggml.tokens")
594 .ok_or(TokenizerLoadError::MissingTokens)?;
595 let id_to_token: Vec<String> = match tokens_value {
596 ferrox_gguf::GgufValue::Array(items) => items
597 .iter()
598 .map(|v| v.as_str().map(|s| s.to_string()))
599 .collect::<Option<Vec<_>>>()
600 .ok_or(TokenizerLoadError::TokensNotStringArray)?,
601 _ => return Err(TokenizerLoadError::TokensNotStringArray),
602 };
603
604 let token_to_id: std::collections::HashMap<String, u32> = id_to_token
605 .iter()
606 .enumerate()
607 .map(|(i, t)| (t.clone(), i as u32))
608 .collect();
609
610 let style = match file.metadata_str("tokenizer.ggml.model") {
611 Some("gemma4") => BpeEncodingStyle::SpmWhitespace,
612 _ => BpeEncodingStyle::Gpt2,
613 };
614
615 let mut merge_rank = std::collections::HashMap::new();
616 if let Some(ferrox_gguf::GgufValue::Array(items)) = file.metadata("tokenizer.ggml.merges") {
617 for (rank, item) in items.iter().enumerate() {
618 if let Some(s) = item.as_str() {
619 if let Some((a, b)) = split_bpe_merge_pair(s, style) {
620 merge_rank.insert((a, b), rank);
621 }
622 }
623 }
624 }
625
626 let (byte_to_unicode, unicode_to_byte) = gpt2_byte_to_unicode();
627 let pretokenize_pattern = match style {
628 // Keyed on the checkpoint's own `tokenizer.ggml.pre`, which
629 // was previously read only to decide BOS prepending.
630 BpeEncodingStyle::Gpt2 => {
631 pretokenize::regex_for(file.metadata_str("tokenizer.ggml.pre").unwrap_or(""))
632 }
633 BpeEncodingStyle::SpmWhitespace => pretokenize::newline_regex(),
634 };
635 let special_tokens = load_special_tokens(file, &id_to_token);
636
637 Ok(GgufBpeTokenizer {
638 token_to_id,
639 id_to_token,
640 merge_rank,
641 byte_to_unicode,
642 unicode_to_byte,
643 special_tokens,
644 pretokenize_pattern,
645 style,
646 })
647 }
648
649 pub fn vocab_size(&self) -> usize {
650 self.id_to_token.len()
651 }
652
653 pub fn has_merges(&self) -> bool {
654 !self.merge_rank.is_empty()
655 }
656
657 /// Greedy BPE merge over one pre-split chunk. GPT-2 style remaps
658 /// bytes through `byte_to_unicode`; Gemma-4 style merges raw UTF-8
659 /// codepoints (after `" "` → `▁` escaping in `encode`).
660 pub fn encode_word(&self, word: &str) -> Vec<u32> {
661 let mut pieces: Vec<String> = match self.style {
662 BpeEncodingStyle::Gpt2 => word
663 .bytes()
664 .map(|b| self.byte_to_unicode[b as usize].to_string())
665 .collect(),
666 BpeEncodingStyle::SpmWhitespace => word.chars().map(|c| c.to_string()).collect(),
667 };
668 if pieces.is_empty() {
669 return Vec::new();
670 }
671
672 loop {
673 let mut best: Option<(usize, usize)> = None; // (rank, index)
674 for i in 0..pieces.len().saturating_sub(1) {
675 if let Some(&rank) = self
676 .merge_rank
677 .get(&(pieces[i].clone(), pieces[i + 1].clone()))
678 {
679 if best.map(|(r, _)| rank < r).unwrap_or(true) {
680 best = Some((rank, i));
681 }
682 }
683 }
684 match best {
685 Some((_, i)) => {
686 let merged = format!("{}{}", pieces[i], pieces[i + 1]);
687 pieces.splice(i..=i + 1, [merged]);
688 }
689 None => break,
690 }
691 }
692
693 pieces.iter().flat_map(|p| self.piece_to_ids(p)).collect()
694 }
695
696 fn piece_to_ids(&self, piece: &str) -> Vec<u32> {
697 if let Some(&id) = self.token_to_id.get(piece) {
698 return vec![id];
699 }
700 match self.style {
701 BpeEncodingStyle::Gpt2 => {
702 // Fall back to first remapped-byte character (GPT-2 base).
703 piece
704 .chars()
705 .next()
706 .and_then(|c| self.token_to_id.get(&c.to_string()))
707 .copied()
708 .map(|id| vec![id])
709 .unwrap_or_else(|| vec![0])
710 }
711 BpeEncodingStyle::SpmWhitespace => {
712 // llama.cpp non-byte-encoded BPE: unknown pieces → `<0xXX>`.
713 piece
714 .bytes()
715 .filter_map(|b| {
716 let hex = format!("<0x{b:02X}>");
717 self.token_to_id.get(&hex).copied()
718 })
719 .collect()
720 }
721 }
722 }
723
724 /// Encodes text: specials first, then style-specific pretokenize +
725 /// `encode_word`. Gemma-4 escapes spaces to `▁` and splits only on
726 /// newlines; newline-only chunks look up the whole string in vocab
727 /// (multi-newline tokens) before BPE.
728 pub fn encode(&self, text: &str) -> Vec<u32> {
729 split_on_special_tokens(text, &self.special_tokens)
730 .into_iter()
731 .flat_map(|seg| -> Vec<u32> {
732 match seg {
733 TextOrSpecial::Special(id) => vec![id],
734 TextOrSpecial::Text(t) => self.encode_text_run(t),
735 }
736 })
737 .collect()
738 }
739
740 /// `pretokenize::split_with_gaps` rather than a bare `find_iter`
741 /// loop: the text BETWEEN matches is input too, and llama.cpp emits
742 /// it as its own chunk. Dropping it lost tabs, NBSPs, form feeds and
743 /// interior newlines out of the middle of every OLMo prompt.
744 fn encode_text_run(&self, text: &str) -> Vec<u32> {
745 match self.style {
746 BpeEncodingStyle::Gpt2 => pretokenize::split_with_gaps(&self.pretokenize_pattern, text)
747 .into_iter()
748 .flat_map(|chunk| self.encode_word(chunk))
749 .collect(),
750 BpeEncodingStyle::SpmWhitespace => {
751 let escaped: String = text
752 .chars()
753 .map(|c| if c == ' ' { SPM_SPACE } else { c })
754 .collect();
755 // Manual newline split (O(n)); avoids regex stack issues
756 // on long non-newline spans (llama.cpp PR #21587).
757 let mut out = Vec::new();
758 let bytes = escaped.as_bytes();
759 let mut i = 0usize;
760 while i < bytes.len() {
761 let is_nl = bytes[i] == b'\n';
762 let mut j = i + 1;
763 while j < bytes.len() && (bytes[j] == b'\n') == is_nl {
764 j += 1;
765 }
766 // Safe: we only split on ASCII `\n`, so `i..j` is UTF-8.
767 let word = std::str::from_utf8(&bytes[i..j]).expect("newline split keeps utf8");
768 if is_nl {
769 if let Some(&id) = self.token_to_id.get(word) {
770 out.push(id);
771 } else {
772 out.extend(self.encode_word(word));
773 }
774 } else {
775 out.extend(self.encode_word(word));
776 }
777 i = j;
778 }
779 out
780 }
781 }
782 }
783
784 /// GPT-2: remapped unicode → bytes. Gemma-4: unescape `▁` → space and
785 /// expand `<0xXX>` byte tokens (same shape as SPM decode).
786 pub fn decode(&self, ids: &[u32]) -> String {
787 match self.style {
788 BpeEncodingStyle::Gpt2 => {
789 let bytes: Vec<u8> = ids
790 .iter()
791 .filter_map(|&id| self.id_to_token.get(id as usize))
792 .flat_map(|token| token.chars())
793 .filter_map(|c| self.unicode_to_byte.get(&c).copied())
794 .collect();
795 String::from_utf8_lossy(&bytes).into_owned()
796 }
797 BpeEncodingStyle::SpmWhitespace => {
798 let mut bytes: Vec<u8> = Vec::new();
799 for &id in ids {
800 let Some(token) = self.id_to_token.get(id as usize) else {
801 continue;
802 };
803 if let Some(b) = spm_byte_fallback_value(token) {
804 bytes.push(b);
805 } else {
806 bytes.extend(token.replace(SPM_SPACE, " ").into_bytes());
807 }
808 }
809 String::from_utf8_lossy(&bytes).into_owned()
810 }
811 }
812 }
813}
814
815/// Split a GGUF merge line `"left right"` into pair. Gemma-4 / llama.cpp
816/// use `find(' ', 1)` on the raw byte string so a leading ASCII space in
817/// `left` is not the separator; search from byte 1 (not char 1) to match.
818fn split_bpe_merge_pair(s: &str, style: BpeEncodingStyle) -> Option<(String, String)> {
819 match style {
820 BpeEncodingStyle::Gpt2 => s
821 .split_once(' ')
822 .map(|(a, b)| (a.to_string(), b.to_string())),
823 BpeEncodingStyle::SpmWhitespace => {
824 let bytes = s.as_bytes();
825 if bytes.len() < 2 {
826 return None;
827 }
828 let pos = bytes[1..].iter().position(|&b| b == b' ')? + 1;
829 // ASCII space is always a UTF-8 char boundary.
830 Some((s[..pos].to_string(), s[pos + 1..].to_string()))
831 }
832 }
833}
834
835fn spm_byte_fallback_value(token: &str) -> Option<u8> {
836 let hex = token.strip_prefix("<0x")?.strip_suffix('>')?;
837 if hex.len() != 2 {
838 return None;
839 }
840 u8::from_str_radix(hex, 16).ok()
841}
842
843/// A real SentencePiece-BPE tokenizer, built from a GGUF file's
844/// `tokenizer.ggml.tokens` + `tokenizer.ggml.scores` metadata
845/// (`tokenizer.ggml.model == "llama"` in GGUF's convention -- this is
846/// SentencePiece's *BPE* model type, not its Unigram model type,
847/// despite both living under the umbrella term "SentencePiece"; the
848/// distinction matters because the encode algorithms are different).
849///
850/// # How this differs from `GgufBpeTokenizer`
851///
852/// `GgufBpeTokenizer` implements GPT2-style BPE: a fixed merge-rank
853/// table applied greedily left-to-right after GPT2's own
854/// byte-to-unicode remap and regex pre-tokenization. SentencePiece-BPE
855/// vocabularies (used by the original LLaMA, and generally any model
856/// whose GGUF reports `tokenizer.ggml.model = "llama"`) don't ship a
857/// merge-rank table at all -- instead every vocabulary entry carries a
858/// score, and encoding works by repeatedly merging whichever *currently
859/// adjacent* pair of symbols forms the highest-scoring known vocabulary
860/// piece, using a priority queue over merge candidates (this is the
861/// `llm_tokenizer_spm` algorithm from llama.cpp, reimplemented here
862/// independently against the public GGUF metadata, not from llama.cpp
863/// source). Preprocessing replaces spaces with `▁` (U+2581) and adds a
864/// leading `▁`, matching SentencePiece's own convention, rather than
865/// GPT2's byte-to-unicode remap.
866///
867/// # A real bug found and fixed while building this
868///
869/// The first implementation of this algorithm checked merge-candidate
870/// validity by adjacency alone (`is this pair still directly next to
871/// each other in the linked list?`). That's necessary but not
872/// sufficient: a symbol's *content* can change between when a
873/// candidate merge is queued and when it's popped, if that symbol was
874/// itself the survivor of a *different* merge in the meantime, while
875/// staying adjacency-valid at the same list position. The fix is to
876/// also store the exact left/right text expected at queue time and
877/// re-check it at pop time, discarding (not re-queuing) any candidate
878/// whose content has since changed. This was caught immediately by
879/// testing against real reference data (see below) rather than by
880/// code review -- the bug produced plausible-looking but wrong output
881/// ("Hello world" tokenized as 6 pieces instead of the correct 2)
882/// which would have been easy to miss without a real ground truth to
883/// check against.
884///
885/// # Verification
886///
887/// Tested against `tests/fixtures/llama-spm-vocab.gguf` (downloaded
888/// directly from `ggml-org/llama.cpp`'s own repository, the real
889/// LLaMA-1/2 tokenizer vocabulary) and its accompanying
890/// `.gguf.inp`/`.gguf.out` files -- llama.cpp's own CI test corpus of
891/// 45 input strings and their exact expected token ID sequences,
892/// covering ASCII, whitespace runs, control characters, CJK/Khmer/
893/// Vietnamese text, emoji, and byte-fallback. All 45 match exactly.
894pub struct GgufSpmTokenizer {
895 token_to_id: std::collections::HashMap<String, u32>,
896 id_to_token: Vec<String>,
897 scores: Vec<f32>,
898 /// Control/user-defined tokens from `tokenizer.ggml.token_type`,
899 /// matched as atomic substrings before normal merging -- see
900 /// `split_on_special_tokens`.
901 special_tokens: Vec<(String, u32)>,
902 /// `tokenizer.ggml.add_space_prefix` (llama.cpp default `true` for
903 /// SPM). When true, each normal-text run after a special (and the
904 /// start of the string) is prefixed with SentencePiece `▁`. Gemma
905 /// GGUFs set this to `false` so `<start_of_turn>user` encodes as
906 /// `[start_of_turn, user]` not `[start_of_turn, ▁user]`.
907 add_space_prefix: bool,
908}
909
910/// A merge candidate in the priority queue: pairs of currently-adjacent
911/// symbol positions, ordered by score (highest first), with ties
912/// broken in favor of the LEFTMOST candidate (smallest `left` symbol
913/// index) -- confirmed against llama.cpp's own real
914/// `llm_bigram_spm::comparator` (`src/llama-vocab.cpp`):
915/// `(l.score < r.score) || (l.score == r.score && l.left > r.left)`.
916/// This matters in practice: many real GGUF vocabularies carry an
917/// exact-zero score for every merge-derived (non-base) piece, so
918/// large stretches of a real tokenization are decided by this tie
919/// rule alone, not by score magnitude. `insertion_order` is kept only
920/// as a last-resort deterministic tiebreak for the (real, possible)
921/// case of two candidates tied on both score AND left index.
922struct SpmMergeCandidate {
923 score: f32,
924 left: usize,
925 right: usize,
926 insertion_order: u64,
927 expected_left_text: String,
928 expected_right_text: String,
929}
930
931impl PartialEq for SpmMergeCandidate {
932 fn eq(&self, other: &Self) -> bool {
933 self.score == other.score && self.insertion_order == other.insertion_order
934 }
935}
936impl Eq for SpmMergeCandidate {}
937impl PartialOrd for SpmMergeCandidate {
938 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
939 Some(self.cmp(other))
940 }
941}
942impl Ord for SpmMergeCandidate {
943 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
944 // BinaryHeap is a max-heap: higher score must compare Greater.
945 // On an exact score tie, the LEFTMOST candidate (smaller
946 // `left`) must compare Greater, so it pops first -- hence the
947 // reversed comparison on `left`. A final tie on `left` too
948 // (impossible for real distinct bigrams, kept for a total
949 // order) falls back to earliest-queued-first.
950 self.score
951 .partial_cmp(&other.score)
952 .unwrap_or(std::cmp::Ordering::Equal)
953 .then_with(|| other.left.cmp(&self.left))
954 .then_with(|| other.insertion_order.cmp(&self.insertion_order))
955 }
956}
957
958impl GgufSpmTokenizer {
959 pub fn from_gguf(file: &impl ferrox_gguf::TensorSource) -> Result<Self, TokenizerLoadError> {
960 let tokens_value = file
961 .metadata("tokenizer.ggml.tokens")
962 .ok_or(TokenizerLoadError::MissingTokens)?;
963 let id_to_token: Vec<String> = match tokens_value {
964 ferrox_gguf::GgufValue::Array(items) => items
965 .iter()
966 .map(|v| v.as_str().map(|s| s.to_string()))
967 .collect::<Option<Vec<_>>>()
968 .ok_or(TokenizerLoadError::TokensNotStringArray)?,
969 _ => return Err(TokenizerLoadError::TokensNotStringArray),
970 };
971 let token_to_id: std::collections::HashMap<String, u32> = id_to_token
972 .iter()
973 .enumerate()
974 .map(|(i, t)| (t.clone(), i as u32))
975 .collect();
976
977 let scores: Vec<f32> = match file.metadata("tokenizer.ggml.scores") {
978 Some(ferrox_gguf::GgufValue::Array(items)) => {
979 items.iter().map(|v| v.as_f32().unwrap_or(0.0)).collect()
980 }
981 _ => vec![0.0; id_to_token.len()],
982 };
983 let special_tokens = load_special_tokens(file, &id_to_token);
984 // llama.cpp defaults SPM `add_space_prefix` to true, then lets
985 // `tokenizer.ggml.add_space_prefix` override (Gemma sets false).
986 let add_space_prefix = match file.metadata("tokenizer.ggml.add_space_prefix") {
987 Some(ferrox_gguf::GgufValue::Bool(v)) => *v,
988 _ => true,
989 };
990
991 Ok(GgufSpmTokenizer {
992 token_to_id,
993 id_to_token,
994 scores,
995 special_tokens,
996 add_space_prefix,
997 })
998 }
999
1000 pub fn vocab_size(&self) -> usize {
1001 self.id_to_token.len()
1002 }
1003
1004 /// Encodes `text` using SentencePiece's space-replacement
1005 /// convention (`' '` -> `▁`, plus a leading `▁`) and the
1006 /// score-prioritized pairwise-merge algorithm described in this
1007 /// struct's doc comment. Characters with no direct vocabulary
1008 /// entry are expanded to UTF-8 byte-fallback tokens (`<0xXX>`,
1009 /// which every real SentencePiece-BPE vocabulary includes for
1010 /// exactly this purpose) before merging begins.
1011 ///
1012 /// Control/user-defined tokens (chat-template markers like
1013 /// `<|user|>`) are first carved out as atomic substrings via
1014 /// `split_on_special_tokens`, matching real llama.cpp's
1015 /// `tokenizer_st_partition` behavior, so they're never shattered
1016 /// into byte-fallback pieces; each remaining raw-text run between
1017 /// them is merged independently. A leading dummy `▁` is applied to
1018 /// a run only when [`Self::add_space_prefix`] is true (llama.cpp
1019 /// `add_space_prefix && is_prev_special` for each fragment).
1020 pub fn encode(&self, text: &str) -> Vec<u32> {
1021 split_on_special_tokens(text, &self.special_tokens)
1022 .into_iter()
1023 .flat_map(|seg| match seg {
1024 TextOrSpecial::Special(id) => vec![id],
1025 TextOrSpecial::Text(t) => self.encode_normal_run(t),
1026 })
1027 .collect()
1028 }
1029
1030 fn encode_normal_run(&self, text: &str) -> Vec<u32> {
1031 let replaced: String = text
1032 .chars()
1033 .map(|c| if c == ' ' { '\u{2581}' } else { c })
1034 .collect();
1035 let normalized = if self.add_space_prefix {
1036 format!("\u{2581}{replaced}")
1037 } else {
1038 replaced
1039 };
1040
1041 let mut symbols: Vec<String> = Vec::new();
1042 for ch in normalized.chars() {
1043 let s = ch.to_string();
1044 if self.token_to_id.contains_key(&s) {
1045 symbols.push(s);
1046 } else {
1047 for byte in s.as_bytes() {
1048 symbols.push(format!("<0x{byte:02X}>"));
1049 }
1050 }
1051 }
1052
1053 let n = symbols.len();
1054 if n == 0 {
1055 return Vec::new();
1056 }
1057 let mut nexts: Vec<Option<usize>> = (1..=n)
1058 .map(|i| if i < n { Some(i) } else { None })
1059 .collect();
1060 let mut prevs: Vec<Option<usize>> = (0..n)
1061 .map(|i| if i == 0 { None } else { Some(i - 1) })
1062 .collect();
1063 let mut alive = vec![true; n];
1064
1065 let mut heap: std::collections::BinaryHeap<SpmMergeCandidate> =
1066 std::collections::BinaryHeap::new();
1067 let mut insertion_order = 0u64;
1068
1069 let try_add_merge = |l: Option<usize>,
1070 r: Option<usize>,
1071 symbols: &[String],
1072 heap: &mut std::collections::BinaryHeap<SpmMergeCandidate>,
1073 insertion_order: &mut u64| {
1074 let (Some(l), Some(r)) = (l, r) else { return };
1075 let merged = format!("{}{}", symbols[l], symbols[r]);
1076 if let Some(&id) = self.token_to_id.get(&merged) {
1077 let score = self.scores.get(id as usize).copied().unwrap_or(0.0);
1078 *insertion_order += 1;
1079 heap.push(SpmMergeCandidate {
1080 score,
1081 left: l,
1082 right: r,
1083 insertion_order: *insertion_order,
1084 expected_left_text: symbols[l].clone(),
1085 expected_right_text: symbols[r].clone(),
1086 });
1087 }
1088 };
1089
1090 for i in 0..n.saturating_sub(1) {
1091 try_add_merge(
1092 Some(i),
1093 Some(i + 1),
1094 &symbols,
1095 &mut heap,
1096 &mut insertion_order,
1097 );
1098 }
1099
1100 while let Some(candidate) = heap.pop() {
1101 let (l, r) = (candidate.left, candidate.right);
1102 if !alive[l] || !alive[r] {
1103 continue;
1104 }
1105 if nexts[l] != Some(r) {
1106 continue;
1107 }
1108 if symbols[l] != candidate.expected_left_text
1109 || symbols[r] != candidate.expected_right_text
1110 {
1111 continue; // stale: content changed since this candidate was queued
1112 }
1113
1114 symbols[l] = format!("{}{}", symbols[l], symbols[r]);
1115 alive[r] = false;
1116 nexts[l] = nexts[r];
1117 if let Some(next_of_r) = nexts[r] {
1118 prevs[next_of_r] = Some(l);
1119 }
1120
1121 try_add_merge(prevs[l], Some(l), &symbols, &mut heap, &mut insertion_order);
1122 try_add_merge(Some(l), nexts[l], &symbols, &mut heap, &mut insertion_order);
1123 }
1124
1125 let mut result = Vec::new();
1126 let mut i = Some(0usize);
1127 while let Some(idx) = i {
1128 if alive[idx] {
1129 result.push(self.token_to_id.get(&symbols[idx]).copied().unwrap_or(0));
1130 }
1131 i = nexts[idx];
1132 }
1133 result
1134 }
1135
1136 /// Reverses a real SentencePiece byte-fallback token (`<0xXX>`,
1137 /// uppercase hex -- the exact format `encode` produces, see its doc
1138 /// comment) back to the raw byte it represents. `None` for any
1139 /// other (normal vocabulary) token.
1140 fn byte_fallback_value(token: &str) -> Option<u8> {
1141 let hex = token.strip_prefix("<0x")?.strip_suffix('>')?;
1142 if hex.len() != 2 {
1143 return None;
1144 }
1145 u8::from_str_radix(hex, 16).ok()
1146 }
1147
1148 pub fn decode(&self, ids: &[u32]) -> String {
1149 // Byte-fallback tokens must be collected as raw bytes (not
1150 // pushed as their 6-character literal token string) and
1151 // UTF-8-decoded together with the rest -- a single real
1152 // multi-byte UTF-8 character can be split across several
1153 // consecutive `<0xXX>` tokens, each individually invalid UTF-8
1154 // on its own. Found and fixed via real-world testing (a real
1155 // downloaded checkpoint's generated text was printing literal
1156 // "<0x0A>" instead of a newline).
1157 let mut bytes: Vec<u8> = Vec::new();
1158 for &id in ids {
1159 let Some(token) = self.id_to_token.get(id as usize) else {
1160 continue;
1161 };
1162 if let Some(b) = Self::byte_fallback_value(token) {
1163 bytes.push(b);
1164 } else {
1165 bytes.extend(token.replace('\u{2581}', " ").into_bytes());
1166 }
1167 }
1168 String::from_utf8_lossy(&bytes).into_owned()
1169 }
1170}
1171
1172/// A real SentencePiece Unigram (ULM) tokenizer, built from a GGUF
1173/// file's `tokenizer.ggml.tokens` + `tokenizer.ggml.scores` metadata
1174/// (`tokenizer.ggml.model == "t5"` in GGUF's convention -- confirmed
1175/// directly against llama.cpp's real vocab-type-loading source
1176/// (`src/llama-vocab.cpp`'s `tokenizer_model == "t5"` case), not
1177/// guessed; T5-family models are the real-world users of this tag).
1178///
1179/// # How this differs from `GgufSpmTokenizer`
1180///
1181/// Both are "SentencePiece" vocabularies, but with entirely different
1182/// encoding algorithms: `GgufSpmTokenizer` implements SentencePiece's
1183/// *BPE* model type (a merge-rank table, greedy pairwise merging).
1184/// Unigram has no merge table at all -- every vocabulary entry carries
1185/// a real log-probability score, and the *optimal* (highest total
1186/// log-probability) segmentation of the whole input is found by a
1187/// forward Viterbi dynamic-programming pass: `best[j]` is the highest-
1188/// scoring way to reach position `j`, computed as
1189/// `max over every vocabulary piece P that ends at j` of
1190/// `best[j - len(P)] + score(P)`. This is reimplemented independently
1191/// against real llama.cpp source read for this purpose
1192/// (`src/llama-vocab.cpp`'s `llm_tokenizer_ugm_session` class) -- not
1193/// copied, but the algorithm (including its unknown-token fallback
1194/// score and tie-breaking) is transcribed deliberately rather than
1195/// guessed, since a plausible-looking-but-wrong Viterbi variant would
1196/// silently produce different segmentations than the model was
1197/// actually trained to expect.
1198///
1199/// Preprocessing matches `GgufSpmTokenizer`'s exactly (`' '` -> `▁`
1200/// U+2581, plus a leading `▁`) -- both are real SentencePiece
1201/// conventions, this being the default `add_dummy_prefix=true` /
1202/// `treat_whitespace_as_suffix=false` behavior. Real SentencePiece
1203/// models can optionally ship a `precompiled_charsmap` (an auxiliary
1204/// normalization table, e.g. NFKC folding) via GGUF's
1205/// `tokenizer.ggml.precompiled_charsmap` key; this implementation does
1206/// not read or apply it (a real, disclosed scope decision, not an
1207/// oversight -- llama.cpp's own loader treats this key as optional
1208/// too, falling back to plain UTF-8 handling when absent).
1209///
1210/// Unlike `GgufSpmTokenizer`, Unigram has no byte-fallback token
1211/// convention in the real reference implementation: a character with
1212/// no matching vocabulary entry is scored via a fixed unknown-token
1213/// penalty (`min_score - 10.0`, matching the real
1214/// `unknown_token_score_penalty` constant) and mapped to the
1215/// vocabulary's real unknown-token id
1216/// (`tokenizer.ggml.unknown_token_id`, defaulting to `0` if absent)
1217/// rather than expanded into raw bytes.
1218///
1219/// Real user-defined/control tokens (GGUF's `tokenizer.ggml.token_type`
1220/// metadata) are not yet given longest-match priority over the
1221/// Viterbi pass the way the real reference implementation does --
1222/// deferred alongside `GgufSpmTokenizer`'s equivalent gap
1223/// (chat-template special-token handling), rather
1224/// than solved once per tokenizer independently.
1225///
1226/// # Verification
1227///
1228/// Cross-validated against a real Unigram model trained with the real
1229/// `sentencepiece` Python library (not a hand-built fixture) --
1230/// exact-match token-id-sequence comparison across ASCII text,
1231/// mixed-case, punctuation, digit runs, repeated whitespace, and
1232/// non-ASCII (accented Latin) text, plus text containing no matching
1233/// vocabulary substrings at all (exercising the unknown-token
1234/// fallback repeatedly).
1235pub struct GgufUnigramTokenizer {
1236 token_to_id: std::collections::HashMap<String, u32>,
1237 id_to_token: Vec<String>,
1238 scores: Vec<f32>,
1239 unk_id: u32,
1240 /// Longest vocabulary piece, in characters -- bounds the Viterbi
1241 /// pass's inner loop so it only ever tries substrings that could
1242 /// possibly be a real vocabulary entry, rather than every possible
1243 /// substring length.
1244 max_piece_chars: usize,
1245 /// `min_score - 10.0`, the real fixed penalty score assigned to the
1246 /// single-character "unknown token" fallback transition, matching
1247 /// the real `unknown_token_score_penalty` constant.
1248 unknown_token_score: f64,
1249 /// Control/user-defined tokens from `tokenizer.ggml.token_type`,
1250 /// matched as atomic substrings before the Viterbi pass -- see
1251 /// `split_on_special_tokens`.
1252 special_tokens: Vec<(String, u32)>,
1253}
1254
1255impl GgufUnigramTokenizer {
1256 pub fn from_gguf(file: &impl ferrox_gguf::TensorSource) -> Result<Self, TokenizerLoadError> {
1257 let tokens_value = file
1258 .metadata("tokenizer.ggml.tokens")
1259 .ok_or(TokenizerLoadError::MissingTokens)?;
1260 let id_to_token: Vec<String> = match tokens_value {
1261 ferrox_gguf::GgufValue::Array(items) => items
1262 .iter()
1263 .map(|v| v.as_str().map(|s| s.to_string()))
1264 .collect::<Option<Vec<_>>>()
1265 .ok_or(TokenizerLoadError::TokensNotStringArray)?,
1266 _ => return Err(TokenizerLoadError::TokensNotStringArray),
1267 };
1268 let token_to_id: std::collections::HashMap<String, u32> = id_to_token
1269 .iter()
1270 .enumerate()
1271 .map(|(i, t)| (t.clone(), i as u32))
1272 .collect();
1273
1274 let scores: Vec<f32> = match file.metadata("tokenizer.ggml.scores") {
1275 Some(ferrox_gguf::GgufValue::Array(items)) => {
1276 items.iter().map(|v| v.as_f32().unwrap_or(0.0)).collect()
1277 }
1278 _ => vec![0.0; id_to_token.len()],
1279 };
1280
1281 let unk_id = file
1282 .metadata("tokenizer.ggml.unknown_token_id")
1283 .and_then(|v| v.as_u64())
1284 .map(|v| v as u32)
1285 .unwrap_or(0);
1286
1287 let max_piece_chars = id_to_token
1288 .iter()
1289 .map(|t| t.chars().count())
1290 .max()
1291 .unwrap_or(1)
1292 .max(1);
1293 let min_score = scores.iter().copied().fold(f32::INFINITY, f32::min) as f64;
1294 let unknown_token_score = min_score - 10.0;
1295 let special_tokens = load_special_tokens(file, &id_to_token);
1296
1297 Ok(GgufUnigramTokenizer {
1298 token_to_id,
1299 id_to_token,
1300 scores,
1301 unk_id,
1302 max_piece_chars,
1303 unknown_token_score,
1304 special_tokens,
1305 })
1306 }
1307
1308 pub fn vocab_size(&self) -> usize {
1309 self.id_to_token.len()
1310 }
1311
1312 /// Encodes `text` via the real forward-Viterbi Unigram algorithm
1313 /// described in this struct's doc comment. Score accumulation uses
1314 /// `f64` (matching the real reference's `double score_sum`), since
1315 /// summing many `f32` log-probabilities over a long input can
1316 /// accumulate enough rounding error to flip which of two
1317 /// near-tied segmentations looks best.
1318 ///
1319 /// Control/user-defined tokens (chat-template markers) are first
1320 /// carved out as atomic substrings via `split_on_special_tokens`;
1321 /// each remaining raw-text run is Viterbi-segmented independently.
1322 pub fn encode(&self, text: &str) -> Vec<u32> {
1323 split_on_special_tokens(text, &self.special_tokens)
1324 .into_iter()
1325 .flat_map(|seg| match seg {
1326 TextOrSpecial::Special(id) => vec![id],
1327 TextOrSpecial::Text(t) => self.encode_normal_run(t),
1328 })
1329 .collect()
1330 }
1331
1332 fn encode_normal_run(&self, text: &str) -> Vec<u32> {
1333 // Real SentencePiece's default normalization rule ("nmt_nfkc",
1334 // used by the overwhelming majority of trained Unigram models
1335 // unless a model deliberately opts into the plain "identity"
1336 // rule) collapses any run of whitespace to a single space and
1337 // trims leading/trailing whitespace, before the dummy-prefix +
1338 // space->▁ substitution below -- confirmed empirically against
1339 // a real trained model, not assumed (a naive per-character
1340 // space->▁ substitution, `GgufSpmTokenizer`'s approach, gives
1341 // a different, wrong segmentation here: one `▁` per space
1342 // instead of one per whitespace *run*). A GGUF file does not
1343 // carry its normalization rule name as its own metadata key,
1344 // so this implements the common default rather than something
1345 // read from the file's own specific config.
1346 let collapsed = text.split_whitespace().collect::<Vec<_>>().join(" ");
1347 let replaced: String = collapsed
1348 .chars()
1349 .map(|c| if c == ' ' { '\u{2581}' } else { c })
1350 .collect();
1351 let normalized = format!("\u{2581}{replaced}");
1352 let chars: Vec<char> = normalized.chars().collect();
1353 let n = chars.len();
1354 if n == 0 {
1355 return Vec::new();
1356 }
1357
1358 struct Best {
1359 token_id: u32,
1360 from: usize,
1361 score: f64,
1362 }
1363 let mut dp: Vec<Best> = (0..=n)
1364 .map(|_| Best {
1365 token_id: 0,
1366 from: 0,
1367 score: f64::NEG_INFINITY,
1368 })
1369 .collect();
1370 dp[0].score = 0.0;
1371
1372 for i in 0..n {
1373 if dp[i].score == f64::NEG_INFINITY {
1374 continue; // unreachable position; never happens since the
1375 // unknown-token fallback below always advances by 1
1376 }
1377 let base = dp[i].score;
1378 let max_len = self.max_piece_chars.min(n - i);
1379 for len in 1..=max_len {
1380 let piece: String = chars[i..i + len].iter().collect();
1381 if let Some(&id) = self.token_to_id.get(&piece) {
1382 let candidate = base + self.scores[id as usize] as f64;
1383 let j = i + len;
1384 if candidate > dp[j].score {
1385 dp[j] = Best {
1386 token_id: id,
1387 from: i,
1388 score: candidate,
1389 };
1390 }
1391 }
1392 }
1393 let j = i + 1;
1394 let candidate = base + self.unknown_token_score;
1395 if candidate > dp[j].score {
1396 dp[j] = Best {
1397 token_id: self.unk_id,
1398 from: i,
1399 score: candidate,
1400 };
1401 }
1402 }
1403
1404 let mut result = Vec::new();
1405 let mut pos = n;
1406 while pos > 0 {
1407 result.push(dp[pos].token_id);
1408 pos = dp[pos].from;
1409 }
1410 result.reverse();
1411 result
1412 }
1413
1414 /// Reverses `encode`'s `' '` <-> `▁` convention. Unigram has no
1415 /// byte-fallback token convention (see this struct's doc comment),
1416 /// so every token here is decoded as plain text.
1417 pub fn decode(&self, ids: &[u32]) -> String {
1418 let mut out = String::new();
1419 for &id in ids {
1420 if let Some(token) = self.id_to_token.get(id as usize) {
1421 out.push_str(&token.replace('\u{2581}', " "));
1422 }
1423 }
1424 out
1425 }
1426}
1427
1428#[cfg(test)]
1429mod gguf_vocab_tests {
1430 use super::*;
1431
1432 fn load_real_fixture() -> GgufBpeTokenizer {
1433 let path = concat!(
1434 env!("CARGO_MANIFEST_DIR"),
1435 "/../../tests/fixtures/llama-bpe-vocab.gguf"
1436 );
1437 let file = ferrox_gguf::GgufFile::open(path).expect("real vocab fixture must open");
1438 GgufBpeTokenizer::from_gguf(&file).expect("real vocab fixture must parse as a tokenizer")
1439 }
1440
1441 #[test]
1442 fn loads_real_downloaded_llama_bpe_vocab() {
1443 let tok = load_real_fixture();
1444 // llama-bpe's real vocab is on the order of 128k tokens; assert
1445 // a loose lower bound so this test doesn't depend on an exact
1446 // upstream count.
1447 assert!(
1448 tok.vocab_size() > 100_000,
1449 "vocab_size={}",
1450 tok.vocab_size()
1451 );
1452 assert!(tok.has_merges(), "llama-bpe vocab ships a real merge table");
1453 }
1454
1455 #[test]
1456 fn decode_of_known_ids_is_stable() {
1457 let tok = load_real_fixture();
1458 // token id 0 exists in every llama-bpe vocab; decoding it must
1459 // not panic and must return the same string every call.
1460 let a = tok.decode(&[0]);
1461 let b = tok.decode(&[0]);
1462 assert_eq!(a, b);
1463 }
1464
1465 #[test]
1466 fn encode_word_never_panics_on_arbitrary_input() {
1467 let tok = load_real_fixture();
1468 for word in ["hello", "", "a", "the quick brown fox", "\u{1f980}"] {
1469 let ids = tok.encode_word(word);
1470 // round-trip through decode must not panic either
1471 let _ = tok.decode(&ids);
1472 }
1473 }
1474
1475 #[test]
1476 fn encode_sentence_round_trips_through_real_vocab() {
1477 let tok = load_real_fixture();
1478 for sentence in [
1479 "the quick brown fox jumps over the lazy dog",
1480 "Hello, World! 123",
1481 "ferrox is a pure-Rust inference engine.",
1482 ] {
1483 let ids = tok.encode(sentence);
1484 assert!(!ids.is_empty());
1485 let decoded = tok.decode(&ids);
1486 assert_eq!(
1487 decoded, sentence,
1488 "full sentence encode/decode through the pre-tokenizer must reproduce the input exactly"
1489 );
1490 }
1491 }
1492
1493 #[test]
1494 fn pretokenizer_splits_on_word_boundaries_not_mid_word() {
1495 let tok = load_real_fixture();
1496 // "cat dog" pre-tokenizes into ["cat", " dog"] (GPT2 convention:
1497 // leading space attaches to the following word). Encoding the
1498 // full sentence and encoding those two pieces separately with
1499 // encode_word must produce the exact same id sequence -- if
1500 // ferrox were still doing one giant merge over the whole
1501 // string (the pre-pretokenizer behavior), a cross-boundary
1502 // merge could produce a different sequence.
1503 let combined = tok.encode("cat dog");
1504 let mut separate = tok.encode_word("cat");
1505 separate.extend(tok.encode_word(" dog"));
1506 assert_eq!(
1507 combined, separate,
1508 "pre-tokenized sentence encoding must match word-by-word encoding at real word boundaries"
1509 );
1510 }
1511
1512 #[test]
1513 fn pretokenizer_keeps_contractions_as_gpt2_does() {
1514 let tok = load_real_fixture();
1515 // GPT2's pattern treats "'t" as its own pre-token (from the
1516 // 's|'t|'re|... alternatives), splitting "don't" into "don" +
1517 // "'t" pieces before BPE, not "do" + "n't" or a single
1518 // 6-character chunk. Confirm the pre-tokenizer actually
1519 // produces that split.
1520 let pieces: Vec<&str> = tok
1521 .pretokenize_pattern
1522 .find_iter("don't")
1523 .map(|m| m.expect("a fixed pattern cannot fail").as_str())
1524 .collect();
1525 assert_eq!(pieces, vec!["don", "'t"]);
1526 }
1527
1528 /// **Defect 3, at the tokenizer level.** The pre-tokenizer arms
1529 /// whose pattern has no catch-all leave text unmatched, and the
1530 /// encoder used to drop it: `find_iter(..).flat_map(..)` sees only
1531 /// the matches. That is silent data loss on the PROMPT, not a
1532 /// different segmentation, so the guard is a byte-for-byte
1533 /// round-trip rather than an id list.
1534 ///
1535 /// The fixture ships `pre = llama-bpe`, whose pattern ends in a
1536 /// catch-all `\s+` and so has no gaps to lose. The OLMo arm is
1537 /// swapped in to reproduce the checkpoint that actually broke —
1538 /// only the split rule changes, the vocabulary and merges stay real.
1539 #[test]
1540 fn every_byte_survives_encoding_on_an_arm_with_unmatched_gaps() {
1541 let mut tok = load_real_fixture();
1542 tok.pretokenize_pattern = super::pretokenize::regex_for("olmo");
1543
1544 // A tab, an NBSP, an interior newline, a form feed and a
1545 // trailing tab: every one of these was unmatched by the OLMo
1546 // pattern and vanished from the prompt.
1547 for text in [
1548 "a\tb\u{a0}c\nd\u{c}e",
1549 "\tif x:\n\t\treturn 1\n\t \treturn 2\n",
1550 "para one\n\npara two\n",
1551 "line one\r\nline two\r\n",
1552 ] {
1553 let ids = tok.encode(text);
1554 assert_eq!(
1555 tok.decode(&ids),
1556 text,
1557 "encoding {text:?} on the olmo arm lost input bytes"
1558 );
1559 }
1560
1561 // And the loss was real: the tab between `a` and `b` is its own
1562 // token, not absorbed into either neighbour.
1563 let ids = tok.encode("a\tb");
1564 assert_eq!(
1565 ids.len(),
1566 3,
1567 "a, the tab, b — the tab is a token of its own"
1568 );
1569 }
1570
1571 #[test]
1572 fn ascii_word_round_trips_through_real_vocab_encode_decode() {
1573 let tok = load_real_fixture();
1574 for word in ["hello", "ferrox", "test", "quick brown fox"] {
1575 let ids = tok.encode_word(word);
1576 assert!(!ids.is_empty(), "encoding {word:?} produced no tokens");
1577 let decoded = tok.decode(&ids);
1578 assert_eq!(
1579 decoded, word,
1580 "round-trip through the real vocab's encode/decode should reproduce ASCII text exactly"
1581 );
1582 }
1583 }
1584
1585 #[test]
1586 fn multibyte_utf8_round_trips_through_real_vocab_encode_decode() {
1587 let tok = load_real_fixture();
1588 for word in ["caf\u{e9}", "\u{1f980}", "\u{4e2d}\u{6587}"] {
1589 let ids = tok.encode_word(word);
1590 let decoded = tok.decode(&ids);
1591 assert_eq!(
1592 decoded, word,
1593 "byte-level BPE must round-trip arbitrary UTF-8, not just ASCII"
1594 );
1595 }
1596 }
1597
1598 #[test]
1599 fn gpt2_remap_matches_known_reference_points() {
1600 // These are well-known fixed points of the real GPT-2
1601 // byte-to-unicode table (verifiable against OpenAI's published
1602 // encoder.py): printable ASCII '!' (0x21) maps to itself, and
1603 // the space byte (0x20), which is NOT in the "already
1604 // printable" ranges, maps to U+0120 ("\u{120}", conventionally
1605 // rendered as "Ġ" in BPE merge tables).
1606 let (fwd, rev) = super::gpt2_byte_to_unicode();
1607 assert_eq!(fwd[0x21], '!');
1608 assert_eq!(fwd[0x20], '\u{120}');
1609 assert_eq!(rev[&'!'], 0x21);
1610 assert_eq!(rev[&'\u{120}'], 0x20);
1611 }
1612
1613 #[test]
1614 fn real_vocab_uses_gpt2_space_remap_in_its_own_tokens() {
1615 // If ferrox's remap table matches the real llama-bpe vocab's
1616 // own convention, at least one real vocabulary entry should
1617 // start with the remapped-space character (a leading-space
1618 // word piece, extremely common in any GPT2-style BPE vocab).
1619 let tok = load_real_fixture();
1620 let has_space_prefixed_token = tok.id_to_token.iter().any(|t| t.starts_with('\u{120}'));
1621 assert!(
1622 has_space_prefixed_token,
1623 "expected at least one real vocab token starting with the GPT2 remapped-space character"
1624 );
1625 }
1626}
1627
1628#[cfg(test)]
1629mod gguf_spm_tests {
1630 use super::*;
1631
1632 fn load_real_fixture() -> GgufSpmTokenizer {
1633 let path = concat!(
1634 env!("CARGO_MANIFEST_DIR"),
1635 "/../../tests/fixtures/llama-spm-vocab.gguf"
1636 );
1637 let file = ferrox_gguf::GgufFile::open(path).expect("real SPM vocab fixture must open");
1638 GgufSpmTokenizer::from_gguf(&file)
1639 .expect("real SPM vocab fixture must parse as a tokenizer")
1640 }
1641
1642 #[test]
1643 fn loads_real_downloaded_llama_spm_vocab() {
1644 let tok = load_real_fixture();
1645 assert_eq!(
1646 tok.vocab_size(),
1647 32000,
1648 "the real LLaMA-1/2 tokenizer vocab is exactly 32000 tokens"
1649 );
1650 }
1651
1652 #[test]
1653 fn matches_known_reference_encodings() {
1654 let tok = load_real_fixture();
1655 assert_eq!(tok.encode("Hello world"), vec![15043, 3186]);
1656 assert_eq!(tok.encode(" Hello world"), vec![29871, 15043, 3186]);
1657 assert_eq!(tok.encode("Hello World"), vec![15043, 2787]);
1658 }
1659
1660 /// Real regression test, found serving a real chat checkpoint:
1661 /// chat-template control tokens (`<|user|>`, `<|assistant|>`) must
1662 /// be recognized as atomic vocabulary entries, not shattered into
1663 /// byte-fallback pieces. Uses a real, hand-built GGUF fixture with
1664 /// genuine `tokenizer.ggml.token_type` CONTROL entries from the
1665 /// fixture generator, not the
1666 /// downloaded real-LLaMA fixture above (which carries no
1667 /// `token_type` array at all).
1668 #[test]
1669 fn chat_template_control_tokens_are_encoded_atomically_not_shattered() {
1670 let path = concat!(
1671 env!("CARGO_MANIFEST_DIR"),
1672 "/tests/fixtures/spm-special-tokens-test-vocab.gguf"
1673 );
1674 let file = ferrox_gguf::GgufFile::open(path).expect("fixture must open");
1675 let tok = GgufSpmTokenizer::from_gguf(&file).expect("fixture must parse");
1676
1677 let user_id = 269u32;
1678 let assistant_id = 270u32;
1679 let ids = tok.encode("<|user|>hello<|assistant|>");
1680
1681 assert_eq!(ids.first().copied(), Some(user_id), "ids={ids:?}");
1682 assert_eq!(ids.last().copied(), Some(assistant_id), "ids={ids:?}");
1683 // The control tokens' own byte-fallback expansions must NOT
1684 // appear anywhere in the output -- they'd show up as a long
1685 // run of ids >= the byte-fallback range if the old shattering
1686 // bug were still present.
1687 assert!(
1688 !ids[1..ids.len() - 1].contains(&user_id)
1689 && !ids[1..ids.len() - 1].contains(&assistant_id),
1690 "control tokens must appear exactly once each, at the boundaries: ids={ids:?}"
1691 );
1692 }
1693
1694 #[test]
1695 fn byte_fallback_handles_control_characters() {
1696 let tok = load_real_fixture();
1697 assert_eq!(
1698 tok.encode("\t"),
1699 vec![29871, 12],
1700 "tab must byte-fallback to <0x09> = token 12"
1701 );
1702 assert_eq!(
1703 tok.encode("\n"),
1704 vec![29871, 13],
1705 "newline must byte-fallback to <0x0A> = token 13"
1706 );
1707 }
1708
1709 /// The strongest test in this file: every one of llama.cpp's own
1710 /// 45 CI test cases for this exact vocabulary
1711 /// (`tests/fixtures/llama-spm-vocab.gguf.inp`/`.out`, downloaded
1712 /// directly from `ggml-org/llama.cpp`), covering ASCII, whitespace
1713 /// runs of every length, control characters, CJK/Khmer/Vietnamese
1714 /// text, emoji (including a ZWJ sequence), and mixed-script text,
1715 /// must produce EXACTLY the token IDs llama.cpp's own tokenizer
1716 /// produces for the same inputs. This is what caught the
1717 /// stale-merge-candidate bug described in `GgufSpmTokenizer`'s doc
1718 /// comment during development.
1719 #[test]
1720 fn matches_llama_cpp_full_reference_test_suite_exactly() {
1721 let tok = load_real_fixture();
1722
1723 let inp_path = concat!(
1724 env!("CARGO_MANIFEST_DIR"),
1725 "/../../tests/fixtures/llama-spm-vocab.gguf.inp"
1726 );
1727 let out_path = concat!(
1728 env!("CARGO_MANIFEST_DIR"),
1729 "/../../tests/fixtures/llama-spm-vocab.gguf.out"
1730 );
1731 let inp_raw = std::fs::read_to_string(inp_path).expect("reference .inp file must exist");
1732 let out_raw = std::fs::read_to_string(out_path).expect("reference .out file must exist");
1733
1734 let marker = "__ggml_vocab_test__\n";
1735 let mut inputs: Vec<&str> = inp_raw.split(marker).collect();
1736 // The split produces a leading/trailing artifact from the
1737 // marker boundaries; drop empty fragments and any trailing
1738 // newline each fragment carries from the format.
1739 inputs.retain(|s| !s.is_empty());
1740 let inputs: Vec<String> = inputs
1741 .iter()
1742 .map(|s| s.strip_suffix('\n').unwrap_or(s).to_string())
1743 .collect();
1744
1745 let outputs: Vec<&str> = out_raw.split('\n').collect();
1746
1747 assert!(
1748 inputs.len() >= 40,
1749 "expected the full ~45-case reference suite, got {}",
1750 inputs.len()
1751 );
1752
1753 let mut checked = 0;
1754 for (i, text) in inputs.iter().enumerate() {
1755 let Some(expected_line) = outputs.get(i) else {
1756 break;
1757 };
1758 let expected_line = expected_line.trim();
1759 if expected_line.is_empty() {
1760 continue;
1761 }
1762 let expected: Vec<u32> = expected_line
1763 .split_whitespace()
1764 .map(|s| s.parse().unwrap())
1765 .collect();
1766 let got = tok.encode(text);
1767 assert_eq!(got, expected, "case #{i}: text={text:?}");
1768 checked += 1;
1769 }
1770 assert!(
1771 checked >= 40,
1772 "expected to actually check at least 40 real cases, only checked {checked}"
1773 );
1774 }
1775
1776 #[test]
1777 fn decode_reverses_encode_for_ascii_text() {
1778 let tok = load_real_fixture();
1779 // SentencePiece's real convention (confirmed by the reference
1780 // suite above) always prepends a dummy leading space before
1781 // tokenizing, so decoding round-trips to " Hello world" (WITH
1782 // a leading space), not "Hello world" -- this is genuine
1783 // LLaMA-tokenizer behavior, not a bug in this test or the
1784 // encoder; downstream text-generation code conventionally
1785 // strips exactly one leading space from decoded output, but
1786 // the raw decode legitimately includes it.
1787 let text = "Hello world";
1788 let ids = tok.encode(text);
1789 assert_eq!(tok.decode(&ids), " Hello world");
1790 }
1791
1792 #[test]
1793 fn decode_reverses_byte_fallback_tokens_to_the_real_raw_bytes() {
1794 // Real bug found via real-world testing:
1795 // decode() used to emit the literal 6-character token string
1796 // "<0x0A>" instead of an actual newline byte.
1797 let tok = load_real_fixture();
1798 // `encode` always prepends a dummy leading space (SentencePiece
1799 // convention, see `decode_reverses_encode_for_ascii_text`
1800 // above), so the decoded round-trip carries it too.
1801 let newline_id = tok.encode("\n");
1802 assert_eq!(tok.decode(&newline_id), " \n");
1803
1804 // A multi-byte UTF-8 character split across several
1805 // consecutive byte-fallback tokens must still decode correctly
1806 // once reassembled -- not as mojibake or individually-invalid
1807 // UTF-8 fragments.
1808 let emoji = "🦀";
1809 let ids = tok.encode(emoji);
1810 assert_eq!(tok.decode(&ids), format!(" {emoji}"));
1811 }
1812}
1813
1814#[cfg(test)]
1815mod gguf_unigram_tests {
1816 use super::*;
1817
1818 /// Real trained SentencePiece Unigram model (100 pieces, trained
1819 /// with the real `sentencepiece` Python library on a small text
1820 /// corpus through a fixture generator), not a
1821 /// hand-guessed vocabulary.
1822 fn load_real_fixture() -> GgufUnigramTokenizer {
1823 let path = concat!(
1824 env!("CARGO_MANIFEST_DIR"),
1825 "/tests/fixtures/unigram-test-vocab.gguf"
1826 );
1827 let file = ferrox_gguf::GgufFile::open(path).expect("real Unigram vocab fixture must open");
1828 GgufUnigramTokenizer::from_gguf(&file)
1829 .expect("real Unigram vocab fixture must parse as a tokenizer")
1830 }
1831
1832 #[test]
1833 fn loads_real_trained_unigram_vocab() {
1834 let tok = load_real_fixture();
1835 assert_eq!(tok.vocab_size(), 100);
1836 }
1837
1838 /// Cross-validated against the exact same trained model's own
1839 /// `sentencepiece.SentencePieceProcessor.Encode` output -- not a
1840 /// hand-computed expectation. Covers ASCII, mixed case, digit runs,
1841 /// repeated whitespace, punctuation, and non-ASCII (accented Latin)
1842 /// text, plus a string with no real vocabulary substrings at all
1843 /// (exercising the unknown-token fallback repeatedly, including
1844 /// consecutive unknown tokens).
1845 #[test]
1846 fn matches_real_sentencepiece_reference_encodings() {
1847 let tok = load_real_fixture();
1848 let cases: &[(&str, &[u32])] = &[
1849 ("hello world", &[3, 63, 4, 95, 8, 3, 36, 14, 11]),
1850 (
1851 "The quick brown fox",
1852 &[34, 3, 89, 10, 65, 70, 57, 49, 73, 12, 54, 8, 30],
1853 ),
1854 (
1855 "Testing unicode: café",
1856 &[74, 44, 20, 35, 47, 4, 83, 3, 62, 13, 25, 18],
1857 ),
1858 (
1859 "Numbers 12345",
1860 &[3, 86, 50, 15, 53, 5, 3, 75, 76, 77, 81, 82],
1861 ),
1862 ("a", &[58]),
1863 (
1864 " multiple spaces ",
1865 &[55, 10, 14, 64, 16, 99, 22, 3, 5, 99, 13, 27, 5],
1866 ),
1867 (
1868 "Zurich naive resume",
1869 &[3, 88, 10, 7, 16, 51, 38, 16, 33, 60, 4, 5, 50, 4],
1870 ),
1871 (
1872 "punctuation! test? yes.",
1873 &[24, 10, 72, 43, 29, 80, 3, 64, 44, 84, 3, 28, 4, 5, 6],
1874 ),
1875 (
1876 "unknown_gibberish_xyz_qqq_zzz",
1877 &[
1878 3, 10, 12, 70, 12, 8, 73, 12, 0, 17, 16, 15, 15, 53, 56, 63, 0, 30, 28, 90, 0,
1879 89, 89, 89, 0, 90, 90, 90,
1880 ],
1881 ),
1882 ];
1883 for (text, expected) in cases {
1884 let got = tok.encode(text);
1885 assert_eq!(&got, expected, "text={text:?}");
1886 }
1887 }
1888
1889 #[test]
1890 fn decode_reverses_encode_for_ascii_text() {
1891 let tok = load_real_fixture();
1892 let ids = tok.encode("hello world");
1893 // encode's leading dummy `▁` decodes back to a leading space,
1894 // same SentencePiece convention as GgufSpmTokenizer.
1895 assert_eq!(tok.decode(&ids), " hello world");
1896 }
1897}
1898
1899#[cfg(test)]
1900mod tests {
1901 use super::*;
1902
1903 #[test]
1904 fn ascii_round_trips_exactly() {
1905 let text = "hello ferrox";
1906 let ids = ByteTokenizer::encode(text);
1907 assert_eq!(ids.len(), text.len());
1908 assert_eq!(ByteTokenizer::decode(&ids), text);
1909 }
1910
1911 #[test]
1912 fn utf8_multibyte_round_trips_exactly() {
1913 let text = "caffe\u{300} \u{1f980}"; // combining accent + emoji, multi-byte UTF-8
1914 let ids = ByteTokenizer::encode(text);
1915 assert_eq!(ByteTokenizer::decode(&ids), text);
1916 }
1917
1918 #[test]
1919 fn all_ids_are_within_byte_vocab_range() {
1920 let ids = ByteTokenizer::encode("mixed ASCII and \u{00e9}\u{00e8} text");
1921 assert!(ids
1922 .iter()
1923 .all(|&id| (id as usize) < ByteTokenizer::VOCAB_SIZE));
1924 }
1925
1926 #[test]
1927 fn empty_string_round_trips() {
1928 assert_eq!(ByteTokenizer::encode(""), Vec::<u32>::new());
1929 assert_eq!(ByteTokenizer::decode(&[]), "");
1930 }
1931
1932 #[test]
1933 fn out_of_range_ids_are_dropped_not_corrupting() {
1934 // 300 is outside the byte vocab; decode should simply skip it
1935 // rather than panicking or wrapping into a wrong byte.
1936 let decoded = ByteTokenizer::decode(&[104, 105, 300, 33]); // "hi" + garbage + "!"
1937 assert_eq!(decoded, "hi!");
1938 }
1939}
1940
1941#[cfg(test)]
1942mod eog_tests {
1943 use super::*;
1944 use ferrox_gguf::{GgufValue, TensorInfo, TensorSource};
1945 use std::collections::HashMap;
1946
1947 struct MetaOnly(HashMap<String, GgufValue>);
1948
1949 impl TensorSource for MetaOnly {
1950 fn metadata(&self, key: &str) -> Option<&GgufValue> {
1951 self.0.get(key)
1952 }
1953 fn find_tensor(&self, _name: &str) -> Option<&TensorInfo> {
1954 None
1955 }
1956 fn tensor_bytes(&self, name: &str) -> Result<&[u8], ferrox_gguf::GgufError> {
1957 Err(ferrox_gguf::GgufError::TensorNotFound(name.to_string()))
1958 }
1959 fn tensor_mapped_range(
1960 &self,
1961 name: &str,
1962 ) -> Result<
1963 (
1964 std::sync::Arc<ferrox_gguf::MmapHandle>,
1965 std::ops::Range<usize>,
1966 ),
1967 ferrox_gguf::GgufError,
1968 > {
1969 Err(ferrox_gguf::GgufError::TensorNotFound(name.to_string()))
1970 }
1971 }
1972
1973 fn source(tokens: &[&str], kv: &[(&str, u64)]) -> MetaOnly {
1974 let mut m = HashMap::new();
1975 m.insert(
1976 "tokenizer.ggml.tokens".to_string(),
1977 GgufValue::Array(
1978 tokens
1979 .iter()
1980 .map(|t| GgufValue::String((*t).to_string()))
1981 .collect(),
1982 ),
1983 );
1984 for (k, v) in kv {
1985 m.insert((*k).to_string(), GgufValue::U32(*v as u32));
1986 }
1987 MetaOnly(m)
1988 }
1989
1990 /// A vocabulary described only by the three metadata keys
1991 /// `should_add_bos_token` reads.
1992 fn vocab_meta(model: &str, pre: Option<&str>, add_bos: Option<bool>) -> MetaOnly {
1993 let mut m = HashMap::new();
1994 m.insert(
1995 "tokenizer.ggml.model".to_string(),
1996 GgufValue::String(model.to_string()),
1997 );
1998 if let Some(pre) = pre {
1999 m.insert(
2000 "tokenizer.ggml.pre".to_string(),
2001 GgufValue::String(pre.to_string()),
2002 );
2003 }
2004 if let Some(v) = add_bos {
2005 m.insert(
2006 "tokenizer.ggml.add_bos_token".to_string(),
2007 GgufValue::Bool(v),
2008 );
2009 }
2010 MetaOnly(m)
2011 }
2012
2013 /// **Defect 4.** llama.cpp sets `add_bos = true` for the whole
2014 /// `LLAMA_VOCAB_PRE_TYPE_LLAMA3` group (`llama-vocab.cpp`, the
2015 /// `tokenizer_pre == "llama-bpe"` arm), and Llama-3.x GGUFs ship no
2016 /// explicit `tokenizer.ggml.add_bos_token`, so a missing group
2017 /// member is a prompt one `<|begin_of_text|>` short of llama.cpp's
2018 /// on every raw completion.
2019 #[test]
2020 fn the_llama_bpe_group_takes_bos_even_with_no_metadata_flag() {
2021 for pre in [
2022 "llama3",
2023 "llama-v3",
2024 "llama-bpe",
2025 "falcon3",
2026 "falcon-h1",
2027 "pixtral",
2028 "midm-2.0",
2029 "lfm2",
2030 "jina-v5-nano",
2031 "tekken",
2032 "chameleon",
2033 ] {
2034 assert!(
2035 should_add_bos_token(&vocab_meta("gpt2", Some(pre), None)),
2036 "llama.cpp sets add_bos for pre={pre}"
2037 );
2038 }
2039 }
2040
2041 /// The other half of the same rule, so the fix cannot be "return
2042 /// true": BPE arms outside that group leave `add_bos` false, and
2043 /// Qwen2's `bos_token_id` is `<|endoftext|>` — prepending it poisons
2044 /// greedy decode.
2045 #[test]
2046 fn other_bpe_pretokenizers_still_do_not_take_bos() {
2047 for pre in ["qwen2", "deepseek-r1-qwen", "gpt-4o", "olmo", "gpt-2", ""] {
2048 assert!(
2049 !should_add_bos_token(&vocab_meta("gpt2", Some(pre), None)),
2050 "llama.cpp leaves add_bos false for pre={pre}"
2051 );
2052 }
2053 }
2054
2055 /// An explicit `tokenizer.ggml.add_bos_token` still wins in both
2056 /// directions — the group default only applies when the key is
2057 /// absent, which is what makes this a *default* rather than an
2058 /// override.
2059 #[test]
2060 fn an_explicit_add_bos_flag_beats_the_pretokenizer_default() {
2061 assert!(!should_add_bos_token(&vocab_meta(
2062 "gpt2",
2063 Some("llama-bpe"),
2064 Some(false)
2065 )));
2066 assert!(should_add_bos_token(&vocab_meta(
2067 "gpt2",
2068 Some("qwen2"),
2069 Some(true)
2070 )));
2071 // SPM still defaults to true with no flag and no pre.
2072 assert!(should_add_bos_token(&vocab_meta("llama", None, None)));
2073 }
2074
2075 /// The failure this exists to stop: a Llama-3 chat checkpoint whose
2076 /// `eos_token_id` is `<|end_of_text|>` while turns actually end with
2077 /// `<|eot_id|>`. Stopping only on the metadata id runs the model past
2078 /// its own turn and it starts interviewing itself.
2079 #[test]
2080 fn turn_enders_count_even_when_they_are_not_the_metadata_eos() {
2081 let src = source(
2082 &["hello", "<|end_of_text|>", "<|eot_id|>", "world"],
2083 &[("tokenizer.ggml.eos_token_id", 1)],
2084 );
2085 let eog = eog_token_ids(&src);
2086 assert!(eog.contains(&1), "metadata eos");
2087 assert!(eog.contains(&2), "<|eot_id|> ends the turn");
2088 assert!(
2089 !eog.contains(&0) && !eog.contains(&3),
2090 "ordinary tokens are not EOG"
2091 );
2092 }
2093
2094 /// gemma-4 ends on `<turn|>`; both it and `<eos>` are in llama.cpp's
2095 /// list, so a gemma checkpoint must stop on either.
2096 #[test]
2097 fn gemma_style_turn_and_eos_are_both_end_of_generation() {
2098 let src = source(&["<eos>", "<turn|>", "x"], &[]);
2099 let eog = eog_token_ids(&src);
2100 assert!(eog.contains(&0) && eog.contains(&1));
2101 assert!(!eog.contains(&2));
2102 }
2103
2104 /// `eot`/`eom` ids are folded in even when the vocabulary spells them
2105 /// something llama.cpp's literal list does not know.
2106 #[test]
2107 fn eot_and_eom_metadata_ids_are_included() {
2108 let src = source(
2109 &["a", "b", "c"],
2110 &[
2111 ("tokenizer.ggml.eot_token_id", 1),
2112 ("tokenizer.ggml.eom_token_id", 2),
2113 ],
2114 );
2115 let eog = eog_token_ids(&src);
2116 assert!(eog.contains(&1) && eog.contains(&2));
2117 }
2118
2119 /// A file with neither the ids nor any known name yields an empty
2120 /// set, so callers keep their previous `eos_id`-only behaviour rather
2121 /// than stopping on something arbitrary.
2122 #[test]
2123 fn a_file_with_nothing_to_go_on_yields_no_stop_tokens() {
2124 let src = source(&["a", "b"], &[]);
2125 assert!(eog_token_ids(&src).is_empty());
2126 }
2127
2128 /// The case a template-evaluating loader creates: gemma-3, Mistral,
2129 /// Phi-3 and DeepSeek-R1-Distill all open their real
2130 /// `tokenizer.chat_template` with `{{ bos_token }}`, so the encoded
2131 /// prompt already starts with the BOS id before the loader gets a
2132 /// look. Measured on the local corpus by
2133 /// `tests/bos_policy.rs::sweep_local_gguf_bos_policy`: 6 of 26
2134 /// checkpoints double their BOS if this is an unconditional insert.
2135 #[test]
2136 fn a_template_that_already_emitted_bos_is_not_given_a_second_one() {
2137 let mut ids = vec![2u32, 105, 2364];
2138 prepend_bos(&mut ids, Some(2));
2139 assert_eq!(ids, vec![2, 105, 2364]);
2140 }
2141
2142 /// The other half of the same rule: Unsloth strips `{{ bos_token }}`
2143 /// out of the templates it exports (TinyLlama's checked-in template
2144 /// is the local example), so on those checkpoints nobody adds BOS
2145 /// unless the loader does.
2146 #[test]
2147 fn a_template_that_stripped_bos_gets_one_from_the_loader() {
2148 let mut ids = vec![529u32, 29989];
2149 prepend_bos(&mut ids, Some(1));
2150 assert_eq!(ids, vec![1, 529, 29989]);
2151 }
2152
2153 /// `None` is the `should_add_bos_token` gate having said no — BPE
2154 /// vocabularies ship a `bos_token_id` they never prepend, and
2155 /// Qwen2-MoE's is `<|endoftext|>`.
2156 #[test]
2157 fn a_vocabulary_that_does_not_take_bos_gets_nothing() {
2158 let mut ids = vec![151644u32, 872];
2159 prepend_bos(&mut ids, None);
2160 assert_eq!(ids, vec![151644, 872]);
2161 }
2162}