Skip to main content

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}