Skip to main content

cortiq_engine/
tokenizer.rs

1//! Byte-level BPE tokenizer — HF tokenizer.json parity.
2//!
3//! Faithful pipeline (matches `tokenizers` for Qwen-style files):
4//!   added-token split (raw text) → NFC → pre-tokenizer regex
5//!   (GPT-2 style, needs lookahead) → byte-level mapping → ranked BPE
6//!   merges → vocab ids. Decode reverses through the byte-level map,
7//!   assembling UTF-8 across token boundaries.
8//!
9//! No silent corruption: a symbol that cannot be encoded is reported
10//! (tracing::error), never dropped without a trace.
11
12use serde::Deserialize;
13use std::collections::{HashMap, HashSet};
14use std::path::Path;
15use std::sync::OnceLock;
16use unicode_normalization::UnicodeNormalization;
17
18/// GPT-2 pre-tokenizer pattern — used when tokenizer.json carries no
19/// explicit Split regex (Qwen files carry their own; see `from_json`).
20const DEFAULT_SPLIT: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+";
21
22/// Tool-call grammar tokens that decode as TEXT even when the vocabulary
23/// marks them special.
24///
25/// `decode` drops special tokens, which is right for control tokens
26/// (`<|im_end|>`, `</s>` …) and wrong for these: they ARE the tool call.
27/// MiniCPM5's tokenizer marks `<function`, `</function>`, `<param`,
28/// `</param>` (and `<tool_call>` …) special, so a call came out as
29/// ` name="get_weather">…Paris` — nothing a parser can find. Qwen's
30/// tokenizer marks its `<tool_call>` non-special for the same reason;
31/// this list gives every vocabulary that behaviour. Only exact tool
32/// markup is listed, so no chat control token can leak into content.
33pub const TOOL_MARKUP_TOKENS: &[&str] = &[
34    "<tool_call>",
35    "</tool_call>",
36    "<function",
37    "</function>",
38    "<param",
39    "</param>",
40];
41
42/// A loaded BPE tokenizer.
43pub struct Tokenizer {
44    /// Token string → ID
45    vocab: HashMap<String, u32>,
46    /// ID → Token string
47    id_to_token: Vec<String>,
48    /// BPE merge ranks: (left, right) → rank (lower merges first)
49    ranks: HashMap<(String, String), u32>,
50    /// The SentencePiece merges over token ids (built on first use; `None`
51    /// when a merge names a string outside the vocabulary — the string
52    /// path then runs)
53    sp_ids: OnceLock<Option<SpMerges>>,
54    /// All added tokens (split during encode; emitted raw at decode)
55    added: Vec<(String, u32)>,
56    /// IDs of added tokens (decode: emit content raw, no byte-map)
57    added_ids: HashSet<u32>,
58    /// IDs of special tokens (skipped by `decode`)
59    special_ids: HashSet<u32>,
60    /// Pre-tokenizer split pattern (None = whitespace fallback for the
61    /// synthetic `byte_level()` tokenizer)
62    /// Applied in order; each subdivides the previous stage's pieces.
63    split_res: Vec<fancy_regex::Regex>,
64    /// Per-stage gap policy, parallel to `split_res`: `true` when the
65    /// stage's Split declares `behavior: Removed, invert: true`, so the
66    /// NON-matching spans (whitespace) are dropped instead of kept as
67    /// pieces. The CLIP family isolates words that way — keeping the
68    /// gaps would hand BPE a stray `Ġ` piece per space.
69    split_drop_gap: Vec<bool>,
70    /// `model.end_of_word_suffix` (CLIP: `</w>`), appended to a piece's
71    /// LAST symbol before the merges run — the merges themselves carry
72    /// the suffix (`["th", "e</w>"]`), so plain words resolve to their
73    /// `…</w>` ids, never to the mid-word spelling.
74    eow_suffix: Option<String>,
75    /// Normalizer `Lowercase` (CLIP): lower-case the segment before the
76    /// pre-tokenizer runs.
77    lowercase: bool,
78    /// Normalizer `Replace(\s+ → " ")` (CLIP): collapse whitespace runs
79    /// to a single space, before lowercasing.
80    collapse_ws: bool,
81    /// SentencePiece Prepend("▁") normalizer present (llama family).
82    /// Gemma replaces spaces with ▁ but does NOT prepend one.
83    sp_prepend: bool,
84    /// Metaspace `prepend_scheme: "first"` in the PRE-tokenizer (no
85    /// Prepend normalizer): the ▁ goes on the very first section of the
86    /// input only, so text after an added token gets none. Nanbeige 4.2
87    /// is this shape — reading only the normalizer left every raw prompt
88    /// short one leading ▁ ("Hello" instead of "▁Hello").
89    sp_prepend_first: bool,
90    /// SentencePiece family (TinyLlama/Llama-2/Mistral): metaspace ▁
91    /// normalization + byte_fallback, no byte-level alphabet.
92    metaspace: bool,
93    /// NFC only when the file's normalizer declares it (Qwen does,
94    /// TinyLlama does not — forcing it broke combining-accent parity).
95    nfc: bool,
96    /// byte → byte-level char (GPT-2 visible-alphabet mapping)
97    byte_to_char: [char; 256],
98    /// byte-level char → byte
99    char_to_byte: HashMap<char, u8>,
100    /// Special tokens
101    pub bos_token_id: Option<u32>,
102    pub eos_token_id: Option<u32>,
103    pub pad_token_id: Option<u32>,
104    /// Chat template special tokens
105    pub im_start_id: Option<u32>,
106    pub im_end_id: Option<u32>,
107    /// Jinja chat template carried by the container (spec §6.1);
108    /// None → hardcoded ChatML fallback.
109    pub chat_template: Option<String>,
110    /// Extra stop ids from the container's generation config.
111    pub extra_eos: HashSet<u32>,
112    /// Generation prepends BOS (llama post_processor semantics).
113    pub add_bos: bool,
114}
115
116impl std::fmt::Debug for Tokenizer {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("Tokenizer")
119            .field("vocab", &self.vocab.len())
120            .field("merges", &self.ranks.len())
121            .field("added", &self.added.len())
122            .finish()
123    }
124}
125
126/// GPT-2 byte↔unicode bijection: printable bytes map to themselves,
127/// the rest get consecutive codepoints from U+0100 up.
128fn bytes_to_unicode() -> ([char; 256], HashMap<char, u8>) {
129    let mut b2c = ['\0'; 256];
130    let mut c2b = HashMap::with_capacity(256);
131    let mut n = 0u32;
132    for b in 0..=255u16 {
133        let printable =
134            (0x21..=0x7E).contains(&b) || (0xA1..=0xAC).contains(&b) || (0xAE..=0xFF).contains(&b);
135        let c = if printable {
136            char::from_u32(b as u32).unwrap()
137        } else {
138            let c = char::from_u32(256 + n).unwrap();
139            n += 1;
140            c
141        };
142        b2c[b as usize] = c;
143        c2b.insert(c, b as u8);
144    }
145    (b2c, c2b)
146}
147
148/// HuggingFace tokenizer.json schema (the parts we execute).
149#[derive(Deserialize)]
150struct HfTokenizerJson {
151    model: HfModel,
152    #[serde(default)]
153    added_tokens: Vec<HfAddedToken>,
154    #[serde(default)]
155    pre_tokenizer: Option<serde_json::Value>,
156    #[serde(default)]
157    normalizer: Option<serde_json::Value>,
158    #[serde(default)]
159    post_processor: Option<serde_json::Value>,
160}
161
162#[derive(Deserialize)]
163struct HfModel {
164    vocab: HashMap<String, u32>,
165    #[serde(default)]
166    merges: Vec<HfMerge>,
167    #[serde(default)]
168    byte_fallback: bool,
169    /// CLIP's `</w>`: the BPE model appends it to a word's last symbol.
170    #[serde(default)]
171    end_of_word_suffix: Option<String>,
172}
173
174/// Merge rules come in two HF flavours: legacy `"a b"` strings and
175/// modern `["a", "b"]` pairs (Qwen3.5 tokenizer.json uses pairs).
176#[derive(Deserialize)]
177#[serde(untagged)]
178enum HfMerge {
179    Pair([String; 2]),
180    Text(String),
181}
182
183#[derive(Deserialize)]
184struct HfAddedToken {
185    id: u32,
186    content: String,
187    special: bool,
188}
189
190/// Collect EVERY Split regex from a pre_tokenizer subtree, in order.
191///
192/// A Sequence applies its splits one after another, each subdividing what
193/// the previous one produced — taking only the first is not an
194/// approximation, it is a different tokenizer. DeepSeek-V4 puts the digit
195/// rule first and the word rule third, so reading one pattern sent whole
196/// sentences into BPE as a single piece and produced ids the model has
197/// never seen.
198fn collect_split_patterns(pt: &serde_json::Value, out: &mut Vec<String>) {
199    if pt.get("type").and_then(|t| t.as_str()) == Some("Split") {
200        if let Some(r) = pt
201            .get("pattern")
202            .and_then(|p| p.get("Regex"))
203            .and_then(|r| r.as_str())
204        {
205            out.push(r.to_string());
206        }
207        return;
208    }
209    if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
210        for p in list {
211            collect_split_patterns(p, out);
212        }
213    }
214}
215
216/// The gap policy of every Split stage, in the SAME traversal order as
217/// [`collect_split_patterns`]: `true` = `behavior: Removed` with
218/// `invert: true` (non-matching spans dropped — CLIP's word isolation).
219fn collect_split_drop_gaps(pt: &serde_json::Value, out: &mut Vec<bool>) {
220    if pt.get("type").and_then(|t| t.as_str()) == Some("Split") {
221        let behavior = pt.get("behavior").and_then(|b| b.as_str()).unwrap_or("");
222        let invert = pt.get("invert").and_then(|i| i.as_bool()).unwrap_or(false);
223        out.push(behavior == "Removed" && invert);
224        return;
225    }
226    if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
227        for p in list {
228            collect_split_drop_gaps(p, out);
229        }
230    }
231}
232
233/// `prepend_scheme` of a Metaspace pre-tokenizer ("always" | "first" |
234/// "never"), searched through a Sequence.
235fn find_prepend_scheme(pt: &serde_json::Value) -> Option<String> {
236    if pt.get("type").and_then(|t| t.as_str()) == Some("Metaspace") {
237        return pt
238            .get("prepend_scheme")
239            .and_then(|p| p.as_str())
240            .map(String::from);
241    }
242    if let Some(list) = pt.get("pretokenizers").and_then(|l| l.as_array()) {
243        return list.iter().find_map(find_prepend_scheme);
244    }
245    None
246}
247
248/// Neutralise `{% generation %}` / `{% endgeneration %}`.
249///
250/// Transformers uses that pair to mark which span of the render is the
251/// assistant's own tokens, so a trainer can build a loss mask. For
252/// INFERENCE it is transparent — the body renders either way — but
253/// minijinja does not know the statement and fails the whole template,
254/// after which the caller quietly serves a ChatML approximation and the
255/// model answers a differently-shaped prompt than the one it was tuned
256/// on. LiquidAI's LFM2.5 templates use it.
257///
258/// Deleting the tag is not enough: `{%- generation -%}` also carries
259/// whitespace control, and dropping it would leave the newline and the
260/// indentation around it in the output. Each tag becomes an assignment
261/// that does nothing, carrying the SAME dashes, so minijinja trims
262/// exactly what the original would have.
263pub(crate) fn strip_generation_tags(tpl: &str) -> std::borrow::Cow<'_, str> {
264    if !tpl.contains("generation") {
265        return std::borrow::Cow::Borrowed(tpl);
266    }
267    let mut out = String::with_capacity(tpl.len());
268    let mut rest = tpl;
269    let mut touched = false;
270    while let Some(open) = rest.find("{%") {
271        let Some(close_rel) = rest[open..].find("%}") else {
272            break;
273        };
274        let close = open + close_rel + 2;
275        let tag = &rest[open..close];
276        let inner = tag[2..tag.len() - 2].trim();
277        let lead = inner.starts_with('-');
278        let trail = inner.ends_with('-');
279        let name = inner.trim_matches('-').trim();
280        out.push_str(&rest[..open]);
281        if name == "generation" || name == "endgeneration" {
282            out.push_str(if lead { "{%-" } else { "{%" });
283            out.push_str(" set _generation_span = true ");
284            out.push_str(if trail { "-%}" } else { "%}" });
285            touched = true;
286        } else {
287            out.push_str(tag);
288        }
289        rest = &rest[close..];
290    }
291    if !touched {
292        return std::borrow::Cow::Borrowed(tpl);
293    }
294    out.push_str(rest);
295    std::borrow::Cow::Owned(out)
296}
297
298/// Close a template's generation-time `<think>` marker for direct answers.
299///
300/// GLM-5's template ends in `<|assistant|><think>` and controls the reasoning
301/// budget with `reasoning_effort`; unlike Qwen templates it does not inspect
302/// `enable_thinking`.  The marker must be closed after rendering without
303/// splitting the complete `<|assistant|>` special token.  Keep this helper
304/// shared by pair- and JSON-shaped chat rendering so both APIs use identical
305/// token sequences.
306fn close_direct_think(rendered: String, enable_thinking: Option<bool>) -> String {
307    if enable_thinking != Some(false) || rendered.contains("</think>") {
308        return rendered;
309    }
310    if let Some(pos) = rendered.rfind("<|assistant|>") {
311        let marker_end = pos + "<|assistant|>".len();
312        let after = &rendered[marker_end..];
313        let mut out = String::with_capacity(rendered.len() + 24);
314        out.push_str(&rendered[..marker_end]);
315        if let Some(rest) = after.strip_prefix("<think>") {
316            out.push_str("<think></think>\n\n");
317            out.push_str(rest);
318        } else {
319            out.push_str(after);
320            if !out.ends_with('\n') {
321                out.push('\n');
322            }
323            out.push_str("<think>\n\n</think>\n\n");
324        }
325        return out;
326    }
327    // Templates without the angle-bracket assistant marker use a
328    // line-oriented role prefix.  Keep that marker intact and append the
329    // empty think block after its newline.
330    if let Some(pos) = rendered.rfind("\nassistant") {
331        let mut insert_at = pos + 1 + "assistant".len();
332        if let Some(idx) = rendered[insert_at..].find('\n') {
333            insert_at += idx + 1;
334        }
335        let mut out = String::with_capacity(rendered.len() + 24);
336        out.push_str(&rendered[..insert_at]);
337        if !out.ends_with('\n') {
338            out.push('\n');
339        }
340        out.push_str("<think>\n\n</think>\n\n");
341        out.push_str(&rendered[insert_at..]);
342        return out;
343    }
344    rendered
345}
346
347impl Tokenizer {
348    /// Load tokenizer from HuggingFace tokenizer.json file.
349    pub fn from_file(path: impl AsRef<Path>) -> Result<Self, TokenizerError> {
350        let data = std::fs::read_to_string(path.as_ref())
351            .map_err(|e| TokenizerError::Io(e.to_string()))?;
352        Self::from_json(&data)
353    }
354
355    /// Load tokenizer from raw tokenizer.json bytes (CMF VOCAB section).
356    pub fn from_bytes(bytes: &[u8]) -> Result<Self, TokenizerError> {
357        let s = std::str::from_utf8(bytes)
358            .map_err(|e| TokenizerError::Parse(format!("vocab is not UTF-8: {e}")))?;
359        Self::from_json(s)
360    }
361
362    /// Load tokenizer from JSON string.
363    pub fn from_json(json: &str) -> Result<Self, TokenizerError> {
364        let hf: HfTokenizerJson =
365            serde_json::from_str(json).map_err(|e| TokenizerError::Parse(e.to_string()))?;
366
367        let mut vocab = hf.model.vocab;
368        let mut ranks = HashMap::new();
369        for (rank, m) in hf.model.merges.into_iter().enumerate() {
370            let (a, b) = match m {
371                HfMerge::Pair([a, b]) => (a, b),
372                HfMerge::Text(s) => {
373                    let mut it = s.splitn(2, ' ');
374                    match (it.next(), it.next()) {
375                        (Some(a), Some(b)) => (a.to_string(), b.to_string()),
376                        _ => continue,
377                    }
378                }
379            };
380            ranks.insert((a, b), rank as u32);
381        }
382
383        // Family detection: SentencePiece carries byte_fallback and/or a
384        // Prepend("▁") normalizer; byte-level BPE carries a Split regex.
385        // Llama-family post_processor prepends <s> at add_special_tokens
386        // time; generation must honor it (word salad without BOS).
387        let mut saw_gemma_bos = false;
388        let add_bos_detected = hf
389            .post_processor
390            .as_ref()
391            .map(|p| {
392                let pp = p.to_string();
393                pp.contains("\"<s>\"") || pp.contains("\"<bos>\"")
394            })
395            .unwrap_or(false);
396        let nfc = hf
397            .normalizer
398            .as_ref()
399            .map(|n| n.to_string().contains("NFC"))
400            .unwrap_or(false);
401        let metaspace = hf.model.byte_fallback
402            || hf
403                .normalizer
404                .as_ref()
405                .map(|n| n.to_string().contains("\u{2581}") || n.to_string().contains("▁"))
406                .unwrap_or(false);
407        let sp_prepend = hf
408            .normalizer
409            .as_ref()
410            .map(|n| n.to_string().contains("Prepend"))
411            .unwrap_or(false);
412        // Metaspace can also live in the pre-tokenizer, carrying its own
413        // prepend_scheme: "always" behaves like the llama normalizer,
414        // "first" only marks the head of the input (see `sp_prepend_first`).
415        let (sp_prepend, sp_prepend_first) = if sp_prepend {
416            (true, false)
417        } else {
418            match hf.pre_tokenizer.as_ref().and_then(find_prepend_scheme) {
419                Some(s) if s == "always" => (true, false),
420                Some(s) if s == "first" => (false, true),
421                _ => (false, false),
422            }
423        };
424        let (split_res, split_drop_gap) = if metaspace {
425            (Vec::new(), Vec::new())
426        } else {
427            let mut pats = Vec::new();
428            if let Some(pt) = hf.pre_tokenizer.as_ref() {
429                collect_split_patterns(pt, &mut pats);
430            }
431            if pats.is_empty() {
432                pats.push(DEFAULT_SPLIT.to_string());
433            }
434            let mut split_drop_gap = Vec::new();
435            if let Some(pt) = hf.pre_tokenizer.as_ref() {
436                collect_split_drop_gaps(pt, &mut split_drop_gap);
437            }
438            while split_drop_gap.len() < pats.len() {
439                split_drop_gap.push(false);
440            }
441            split_drop_gap.truncate(pats.len());
442            (
443                pats.iter()
444                    .map(|p| {
445                        fancy_regex::Regex::new(p)
446                            .map_err(|e| TokenizerError::Parse(format!("pre-tokenizer regex: {e}")))
447                    })
448                    .collect::<Result<Vec<_>, _>>()?,
449                split_drop_gap,
450            )
451        };
452        // CLIP's normalizer: NFC → Replace(\s+ → " ") → Lowercase.
453        let norm_s = hf
454            .normalizer
455            .as_ref()
456            .map(|n| n.to_string())
457            .unwrap_or_default();
458        let collapse_ws = norm_s.contains("\"Replace\"") && norm_s.contains("\\\\s+");
459        let lowercase = norm_s.contains("\"Lowercase\"");
460
461        // Added tokens: longest-first so overlapping contents match right.
462        let mut bos_token_id = None;
463        let mut eos_token_id = None;
464        let mut pad_token_id = None;
465        let mut im_start_id = None;
466        let mut im_end_id = None;
467        let mut special_ids = HashSet::new();
468        let mut added_ids = HashSet::new();
469        let mut added = Vec::new();
470
471        for at in &hf.added_tokens {
472            vocab.insert(at.content.clone(), at.id);
473            added.push((at.content.clone(), at.id));
474            added_ids.insert(at.id);
475            if at.special && !TOOL_MARKUP_TOKENS.contains(&at.content.as_str()) {
476                special_ids.insert(at.id);
477            }
478            match at.content.as_str() {
479                "<|endoftext|>" | "</s>" | "[EOS]" => eos_token_id = Some(at.id),
480                "<|im_start|>" => im_start_id = Some(at.id),
481                "<|im_end|>" => im_end_id = Some(at.id),
482                "<s>" | "[BOS]" | "<|startoftext|>" => bos_token_id = Some(at.id),
483                // Gemma spells BOS as literal "<bos>" — the family
484                // REQUIRES it on every sequence, and newer tokenizers
485                // (gemma-4) no longer say so in a post_processor.
486                "<bos>" => {
487                    bos_token_id = Some(at.id);
488                    saw_gemma_bos = true;
489                }
490                "<pad>" => pad_token_id = Some(at.id),
491                _ => {}
492            }
493        }
494
495        // DeepSeek-V4.1 keeps harmony/DSML markers in the ordinary vocabulary
496        // on some tokenizer revisions and in `added_tokens` on others. Treat
497        // either spelling as atomic so image placeholders cannot split into
498        // byte-BPE pieces, and make the configured BOS/EOS ids discoverable.
499        const DSV41_SPECIALS: &[&str] = &[
500            "<|begin▁of▁sentence|>",
501            "<|end▁of▁sentence|>",
502            "<|User|>",
503            "<|Assistant|>",
504            "<|System|>",
505            "<|latest_reminder|>",
506            "<|deepseek_image|>",
507            "<|action|>",
508            "<|query|>",
509            "<|authority|>",
510            "<|domain|>",
511            "<|title|>",
512            "<|read_url|>",
513            "<think>",
514            "</think>",
515            "|DSML|",
516        ];
517        for token in DSV41_SPECIALS {
518            if let Some(&id) = vocab.get(*token) {
519                if !added.iter().any(|(content, _)| content.as_str() == *token) {
520                    added.push(((*token).to_string(), id));
521                }
522                added_ids.insert(id);
523                match *token {
524                    "<|begin▁of▁sentence|>" => bos_token_id = Some(id),
525                    "<|end▁of▁sentence|>" => eos_token_id = Some(id),
526                    _ => {}
527                }
528            }
529        }
530        added.sort_by_key(|(c, _)| std::cmp::Reverse(c.len()));
531
532        // Gemma REQUIRES a leading <bos> on every sequence, but newer
533        // tokenizers (gemma-4, 262k vocab) no longer spell it in a
534        // post_processor template — the family marker <start_of_turn>
535        // is the reliable tell. Without this, raw-text scoring runs
536        // unanchored and the first ~30 positions read worse than
537        // uniform (the chat path masked it: the template carries <bos>).
538        let gemma_family = saw_gemma_bos
539            || vocab.contains_key("<start_of_turn>")
540            || added.iter().any(|(c, _)| c == "<start_of_turn>");
541
542        // The post_processor template names the exact BOS content
543        // (llama "<s>", gemma "<bos>" — gemma's vocab carries BOTH, so
544        // added-token scan order must not decide).
545        if let Some(pp) = hf.post_processor.as_ref() {
546            let pp = pp.to_string();
547            for name in ["<bos>", "<s>"] {
548                if pp.contains(&format!("\"{name}\"")) {
549                    if let Some(&id) = vocab.get(name) {
550                        bos_token_id = Some(id);
551                    }
552                    break;
553                }
554            }
555        }
556
557        // Build reverse map
558        let max_id = vocab.values().copied().max().unwrap_or(0) as usize;
559        let mut id_to_token = vec![String::new(); max_id + 1];
560        for (token, &id) in &vocab {
561            if (id as usize) < id_to_token.len() {
562                id_to_token[id as usize] = token.clone();
563            }
564        }
565
566        let (byte_to_char, char_to_byte) = bytes_to_unicode();
567
568        tracing::info!(
569            "Tokenizer loaded: {} vocab, {} merges, {} added, eos={:?}",
570            vocab.len(),
571            ranks.len(),
572            added.len(),
573            eos_token_id
574        );
575
576        Ok(Self {
577            vocab,
578            id_to_token,
579            ranks,
580            sp_ids: OnceLock::new(),
581            added,
582            added_ids,
583            special_ids,
584            split_res,
585            split_drop_gap,
586            eow_suffix: hf.model.end_of_word_suffix,
587            lowercase,
588            collapse_ws,
589            metaspace,
590            sp_prepend,
591            sp_prepend_first,
592            nfc,
593            byte_to_char,
594            char_to_byte,
595            bos_token_id,
596            eos_token_id,
597            pad_token_id,
598            im_start_id,
599            im_end_id,
600            chat_template: None,
601            extra_eos: HashSet::new(),
602            add_bos: add_bos_detected || gemma_family,
603        })
604    }
605
606    /// Create a minimal tokenizer for testing (byte tokens, no merges).
607    pub fn byte_level() -> Self {
608        let mut vocab = HashMap::new();
609        let mut id_to_token = Vec::with_capacity(256);
610        for i in 0..256u32 {
611            let tok = format!("<0x{:02X}>", i);
612            vocab.insert(tok.clone(), i);
613            id_to_token.push(tok);
614        }
615        let (byte_to_char, char_to_byte) = bytes_to_unicode();
616        Self {
617            vocab,
618            id_to_token,
619            ranks: HashMap::new(),
620            sp_ids: OnceLock::new(),
621            added: Vec::new(),
622            added_ids: HashSet::new(),
623            special_ids: HashSet::new(),
624            split_res: Vec::new(),
625            split_drop_gap: Vec::new(),
626            eow_suffix: None,
627            lowercase: false,
628            collapse_ws: false,
629            metaspace: false,
630            sp_prepend: false,
631            sp_prepend_first: false,
632            nfc: false,
633            byte_to_char,
634            char_to_byte,
635            bos_token_id: None,
636            eos_token_id: None,
637            pad_token_id: None,
638            im_start_id: None,
639            im_end_id: None,
640            chat_template: None,
641            extra_eos: HashSet::new(),
642            add_bos: false,
643        }
644    }
645
646    /// Encode text to token IDs.
647    pub fn encode(&self, text: &str) -> Vec<u32> {
648        let mut ids = Vec::new();
649        // Added tokens match on raw text (normalized: false), longest first.
650        let mut rest = text;
651        // `prepend_scheme: "first"` marks only the section that starts at
652        // offset 0 — HF drops the ▁ for everything after an added token,
653        // which is why a chat prompt opening with <|im_start|> tokenizes
654        // the same either way and only raw prompts were wrong.
655        let mut head = true;
656        'outer: while !rest.is_empty() {
657            let mut best: Option<(usize, usize, u32)> = None; // (pos, len, id)
658            for (content, id) in &self.added {
659                if let Some(pos) = rest.find(content.as_str()) {
660                    let better = match best {
661                        None => true,
662                        Some((bp, bl, _)) => pos < bp || (pos == bp && content.len() > bl),
663                    };
664                    if better {
665                        best = Some((pos, content.len(), *id));
666                    }
667                    if pos == 0 {
668                        break; // earliest possible; added is longest-first
669                    }
670                }
671            }
672            match best {
673                Some((pos, len, id)) => {
674                    self.encode_segment_at(&rest[..pos], head, &mut ids);
675                    ids.push(id);
676                    rest = &rest[pos + len..];
677                    head = false;
678                }
679                None => {
680                    self.encode_segment_at(rest, head, &mut ids);
681                    break 'outer;
682                }
683            }
684        }
685        ids
686    }
687
688    /// Encode text as PLAIN text: the whole input is one added-token-free
689    /// segment (NFC → split → byte-map → BPE), so a literal `<|im_end|>`
690    /// in it stays bytes instead of becoming the special id. This is the
691    /// trainer's `Bpe::encode` (user text, never the template frame) —
692    /// router v2 tokenizes the user message with it (spec §9.4).
693    pub fn encode_plain(&self, text: &str) -> Vec<u32> {
694        let mut ids = Vec::new();
695        self.encode_segment_at(text, true, &mut ids);
696        ids
697    }
698
699    /// Encode one added-token-free segment: NFC → whitespace collapse →
700    /// lowercase (the CLIP normalizer chain) → split → byte-map → BPE.
701    /// `head` says whether this section starts at offset 0 of the
702    /// input — only that one takes a `prepend_scheme: "first"` ▁.
703    fn encode_segment_at(&self, segment: &str, head: bool, out: &mut Vec<u32>) {
704        if segment.is_empty() {
705            return;
706        }
707        let mut norm: String = if self.nfc {
708            segment.nfc().collect()
709        } else {
710            segment.to_string()
711        };
712        if self.collapse_ws {
713            // Replace(\s+ → " "): every whitespace RUN becomes one space.
714            let mut collapsed = String::with_capacity(norm.len());
715            let mut in_ws = false;
716            for ch in norm.chars() {
717                if ch.is_whitespace() {
718                    if !in_ws {
719                        collapsed.push(' ');
720                        in_ws = true;
721                    }
722                } else {
723                    collapsed.push(ch);
724                    in_ws = false;
725                }
726            }
727            norm = collapsed;
728        }
729        if self.lowercase {
730            norm = norm.to_lowercase();
731        }
732        if self.metaspace {
733            // SentencePiece: [Prepend("▁") +] Replace(" "→"▁"), BPE over
734            // chars of the whole span (no pre-tokenizer, no byte map).
735            // Gemma's normalizer replaces only — no dummy prefix.
736            let sp = if self.sp_prepend {
737                // Normalizer order: Prepend THEN Replace, unguarded — so
738                // " hello" really does become ▁▁hello on llama.
739                format!("\u{2581}{}", norm).replace(' ', "\u{2581}")
740            } else {
741                // Pre-tokenizer Metaspace: Replace, then prepend only if
742                // the span does not already start with ▁ (HF's guard, so
743                // " hello" stays one ▁, not two).
744                let replaced = norm.replace(' ', "\u{2581}");
745                if self.sp_prepend_first && head && !replaced.starts_with('\u{2581}') {
746                    format!("\u{2581}{replaced}")
747                } else {
748                    replaced
749                }
750            };
751            self.bpe_piece_sp(&sp, out);
752            return;
753        }
754        if !self.split_res.is_empty() {
755            // Each stage subdivides the pieces the previous one left, with
756            // Isolated behaviour: both the matches and the gaps survive.
757            // A stage declaring `behavior: Removed, invert: true` keeps
758            // only the matches — its gaps (whitespace) are dropped.
759            let mut pieces: Vec<(usize, usize)> = vec![(0, norm.len())];
760            for (stage, re) in self.split_res.iter().enumerate() {
761                let drop = self.split_drop_gap.get(stage).copied().unwrap_or(false);
762                let mut next: Vec<(usize, usize)> = Vec::with_capacity(pieces.len() * 2);
763                for (ps, pe) in pieces {
764                    let seg = &norm[ps..pe];
765                    let mut last = 0usize;
766                    for m in re.find_iter(seg) {
767                        let m = match m {
768                            Ok(m) => m,
769                            Err(e) => {
770                                tracing::error!("pre-tokenizer regex failed: {e}");
771                                break;
772                            }
773                        };
774                        if m.start() > last && !drop {
775                            next.push((ps + last, ps + m.start()));
776                        }
777                        if m.end() > m.start() {
778                            next.push((ps + m.start(), ps + m.end()));
779                        }
780                        last = m.end();
781                    }
782                    if last < seg.len() && !drop {
783                        next.push((ps + last, pe));
784                    }
785                }
786                pieces = next;
787            }
788            for (ps, pe) in pieces {
789                self.bpe_piece(&norm[ps..pe], out);
790            }
791        } else {
792            {
793                // Synthetic byte_level() tokenizer: raw byte tokens.
794                for b in norm.bytes() {
795                    let tok = format!("<0x{:02X}>", b);
796                    if let Some(&id) = self.vocab.get(&tok) {
797                        out.push(id);
798                    }
799                }
800            }
801        }
802    }
803
804    /// The id-level merge table of a SentencePiece vocabulary.
805    fn sp_merges(&self) -> Option<&SpMerges> {
806        self.sp_ids
807            .get_or_init(|| SpMerges::build(&self.vocab, &self.ranks))
808            .as_ref()
809    }
810
811    /// SentencePiece BPE: symbols are chars (no byte-level alphabet);
812    /// unknown symbols fall back to <0xNN> tokens per UTF-8 byte.
813    ///
814    /// The rule is "merge every occurrence of the lowest-ranked adjacent
815    /// pair, left to right, until none is left". [`SpMerges::encode`] runs
816    /// it over ids with a heap (O(n log n)); this string form (O(n²) with a
817    /// clone per pair) remains for vocabularies whose merges leave the
818    /// vocabulary and as the reference the fast form is tested against.
819    fn bpe_piece_sp(&self, piece: &str, out: &mut Vec<u32>) {
820        if piece.is_empty() {
821            return;
822        }
823        if !SP_STRINGS.with(|f| f.get()) {
824            if let Some(m) = self.sp_merges() {
825                m.encode(&self.vocab, piece, out);
826                return;
827            }
828        }
829        self.bpe_piece_sp_strings(piece, out);
830    }
831
832    /// [`encode`](Self::encode) with the SentencePiece string form forced:
833    /// the reference the id-level BPE is checked against.
834    #[doc(hidden)]
835    pub fn encode_reference(&self, text: &str) -> Vec<u32> {
836        SP_STRINGS.with(|f| f.set(true));
837        let ids = self.encode(text);
838        SP_STRINGS.with(|f| f.set(false));
839        ids
840    }
841
842    fn bpe_piece_sp_strings(&self, piece: &str, out: &mut Vec<u32>) {
843        let mut sym: Vec<String> = piece.chars().map(|c| c.to_string()).collect();
844        loop {
845            let mut best: Option<(u32, usize)> = None;
846            for i in 0..sym.len().saturating_sub(1) {
847                if let Some(&r) = self.ranks.get(&(sym[i].clone(), sym[i + 1].clone())) {
848                    if best.map(|(br, _)| r < br).unwrap_or(true) {
849                        best = Some((r, i));
850                    }
851                }
852            }
853            let Some((_, i)) = best else { break };
854            let merged = format!("{}{}", sym[i], sym[i + 1]);
855            let (left, right) = (sym[i].clone(), sym[i + 1].clone());
856            let mut j = 0;
857            while j + 1 < sym.len() {
858                if sym[j] == left && sym[j + 1] == right {
859                    sym[j] = merged.clone();
860                    sym.remove(j + 1);
861                }
862                j += 1;
863            }
864        }
865        for t in &sym {
866            if let Some(&id) = self.vocab.get(t) {
867                out.push(id);
868            } else {
869                let mut ok = true;
870                for byte in t.bytes() {
871                    let tok = format!("<0x{:02X}>", byte);
872                    match self.vocab.get(&tok) {
873                        Some(&id) => out.push(id),
874                        None => {
875                            ok = false;
876                            break;
877                        }
878                    }
879                }
880                if !ok {
881                    tracing::error!("tokenizer: no id for SP symbol {t:?} — dropped");
882                }
883            }
884        }
885    }
886
887    /// Byte-level map one pre-token piece, then ranked BPE merges.
888    fn bpe_piece(&self, piece: &str, out: &mut Vec<u32>) {
889        if piece.is_empty() {
890            return;
891        }
892        let mapped: Vec<String> = piece
893            .bytes()
894            .map(|b| self.byte_to_char[b as usize].to_string())
895            .collect();
896        let mut sym = mapped;
897        // CLIP's `end_of_word_suffix`: the last symbol of the word wears
898        // it before any merge runs, so the suffixed merges ("th",
899        // "e</w>") can form and the word resolves to its `…</w>` id.
900        if let Some(sfx) = &self.eow_suffix {
901            if let Some(last) = sym.last_mut() {
902                last.push_str(sfx);
903            }
904        }
905
906        // Classic BPE: repeatedly merge the lowest-rank adjacent pair.
907        loop {
908            let mut best: Option<(u32, usize)> = None;
909            for i in 0..sym.len().saturating_sub(1) {
910                if let Some(&r) = self.ranks.get(&(sym[i].clone(), sym[i + 1].clone())) {
911                    if best.map(|(br, _)| r < br).unwrap_or(true) {
912                        best = Some((r, i));
913                    }
914                }
915            }
916            let Some((_, i)) = best else { break };
917            let merged = format!("{}{}", sym[i], sym[i + 1]);
918            // Merge ALL occurrences of this exact pair, left to right.
919            let (left, right) = (sym[i].clone(), sym[i + 1].clone());
920            let mut j = 0;
921            while j + 1 < sym.len() {
922                if sym[j] == left && sym[j + 1] == right {
923                    sym[j] = merged.clone();
924                    sym.remove(j + 1);
925                }
926                j += 1;
927            }
928        }
929
930        for s in &sym {
931            if let Some(&id) = self.vocab.get(s) {
932                out.push(id);
933            } else {
934                // Byte-fallback (synthetic vocabs); never drop silently.
935                let mut ok = true;
936                for ch in s.chars() {
937                    let Some(&b) = self.char_to_byte.get(&ch) else {
938                        ok = false;
939                        break;
940                    };
941                    let tok = format!("<0x{:02X}>", b);
942                    if let Some(&id) = self.vocab.get(&tok) {
943                        out.push(id);
944                    } else {
945                        ok = false;
946                        break;
947                    }
948                }
949                if !ok {
950                    tracing::error!("tokenizer: no id for symbol {s:?} — dropped");
951                }
952            }
953        }
954    }
955
956    /// Decode token IDs back to text. Special tokens are skipped; added
957    /// tokens are raw text; everything else reverses the byte-level map.
958    pub fn decode(&self, ids: &[u32]) -> String {
959        let mut bytes: Vec<u8> = Vec::new();
960        for &id in ids {
961            if self.special_ids.contains(&id) {
962                continue;
963            }
964            let idx = id as usize;
965            if idx >= self.id_to_token.len() {
966                continue;
967            }
968            let tok = &self.id_to_token[idx];
969            if self.added_ids.contains(&id) {
970                // Gemma-3n declares its multi-space ▁-runs as ADDED
971                // tokens — verbatim passthrough leaked ▁ into output.
972                if self.metaspace && tok.contains('\u{2581}') {
973                    bytes.extend_from_slice(tok.replace('\u{2581}', " ").as_bytes());
974                } else {
975                    bytes.extend_from_slice(tok.as_bytes());
976                }
977                continue;
978            }
979            // Byte-fallback / legacy byte tokens
980            if tok.starts_with("<0x") && tok.ends_with('>') && tok.len() == 6 {
981                if let Ok(b) = u8::from_str_radix(&tok[3..5], 16) {
982                    bytes.push(b);
983                    continue;
984                }
985            }
986            if self.metaspace {
987                // SP decoder: Replace(▁→" "); UTF-8 chars pass through.
988                for ch in tok.chars() {
989                    if ch == '\u{2581}' {
990                        bytes.push(b' ');
991                    } else {
992                        let mut buf = [0u8; 4];
993                        bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
994                    }
995                }
996                continue;
997            }
998            for ch in tok.chars() {
999                match self.char_to_byte.get(&ch) {
1000                    Some(&b) => bytes.push(b),
1001                    // Not a byte-level char (shouldn't happen for real
1002                    // vocabs) — pass the char through as UTF-8.
1003                    None => {
1004                        let mut buf = [0u8; 4];
1005                        bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
1006                    }
1007                }
1008            }
1009        }
1010        let text = String::from_utf8_lossy(&bytes).into_owned();
1011        if self.metaspace && (self.sp_prepend || self.sp_prepend_first) {
1012            // SP decoder Strip(start=1): one leading space from Prepend.
1013            if let Some(stripped) = text.strip_prefix(' ') {
1014                return stripped.to_string();
1015            }
1016        }
1017        text
1018    }
1019
1020    /// Streaming decode of ONE token: no sequence-level Strip — a
1021    /// per-token strip would eat the ▁-spaces of every SP word.
1022    pub fn decode_token(&self, id: u32) -> String {
1023        if self.special_ids.contains(&id) {
1024            return String::new();
1025        }
1026        let idx = id as usize;
1027        if idx >= self.id_to_token.len() {
1028            return String::new();
1029        }
1030        let tok = &self.id_to_token[idx];
1031        if self.added_ids.contains(&id) {
1032            if self.metaspace && tok.contains('\u{2581}') {
1033                return tok.replace('\u{2581}', " ");
1034            }
1035            return tok.clone();
1036        }
1037        if tok.starts_with("<0x") && tok.ends_with('>') && tok.len() == 6 {
1038            if let Ok(b) = u8::from_str_radix(&tok[3..5], 16) {
1039                return String::from_utf8_lossy(&[b]).into_owned();
1040            }
1041        }
1042        if self.metaspace {
1043            return tok.replace('\u{2581}', " ");
1044        }
1045        let mut bytes = Vec::new();
1046        for ch in tok.chars() {
1047            match self.char_to_byte.get(&ch) {
1048                Some(&b) => bytes.push(b),
1049                None => {
1050                    let mut buf = [0u8; 4];
1051                    bytes.extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
1052                }
1053            }
1054        }
1055        String::from_utf8_lossy(&bytes).into_owned()
1056    }
1057
1058    /// Decode one vocabulary entry for Engram's compressed token map while
1059    /// retaining special tokens.
1060    pub fn decode_token_for_hash(&self, id: u32) -> String {
1061        let idx = id as usize;
1062        if idx >= self.id_to_token.len() {
1063            return String::new();
1064        }
1065        if self.special_ids.contains(&id) {
1066            return self.id_to_token[idx].clone();
1067        }
1068        self.decode_token(id)
1069    }
1070
1071    /// Decode generated protocol text while retaining special markers. The
1072    /// V4.1 harmony parser needs `<think>`, EOS, and spaced DSML tags.
1073    pub fn decode_for_protocol(&self, ids: &[u32]) -> String {
1074        let mut out = String::new();
1075        for &id in ids {
1076            let idx = id as usize;
1077            if self.special_ids.contains(&id) {
1078                if let Some(token) = self.id_to_token.get(idx) {
1079                    out.push_str(token);
1080                }
1081            } else {
1082                out.push_str(&self.decode_token(id));
1083            }
1084        }
1085        out
1086    }
1087
1088    /// Return the backend vocabulary spelling for an Engram map entry.
1089    pub fn raw_token_for_hash(&self, id: u32) -> String {
1090        self.id_to_token
1091            .get(id as usize)
1092            .cloned()
1093            .unwrap_or_default()
1094    }
1095
1096    /// Render the container's Jinja chat template (HF semantics:
1097    /// trim_blocks + lstrip_blocks + loop controls) and encode it.
1098    /// Falls back to hardcoded ChatML when the file carries none.
1099    pub fn apply_chat_template(&self, messages: &[(String, String)]) -> Vec<u32> {
1100        self.apply_chat_template_opts(messages, None)
1101    }
1102
1103    /// Like `apply_chat_template`, with an explicit `enable_thinking` value for
1104    /// reasoning-model templates (Qwen3/3.5 emit an empty <think> block when it
1105    /// is false, so the model answers directly). `None` leaves the variable
1106    /// undefined — the template's own default applies.
1107
1108    /// Chat template with the FULL message shape and a tool list.
1109    ///
1110    /// The pair-based API below flattens every message to (role, text),
1111    /// which silently drops exactly what agentic use needs: the `tools`
1112    /// array, `role: "tool"` results, and `tool_calls` on assistant
1113    /// turns. The templates this format embeds — Qwen-family, Nanbeige —
1114    /// have carried a `{%- if tools %}` branch all along; this is the
1115    /// call that finally feeds it. Messages arrive as JSON objects in
1116    /// the OpenAI shape and pass through to minijinja unflattened, so a
1117    /// template sees the same fields a Python `apply_chat_template`
1118    /// would.
1119    pub fn apply_chat_template_json(
1120        &self,
1121        messages: &[serde_json::Value],
1122        tools: Option<&[serde_json::Value]>,
1123        enable_thinking: Option<bool>,
1124    ) -> Vec<u32> {
1125        self.apply_chat_template_json_with_reasoning_effort(messages, tools, enable_thinking, None)
1126    }
1127
1128    /// JSON chat template rendering with an optional template-native reasoning
1129    /// effort. An explicit value is passed through to Jinja unchanged.
1130    pub fn apply_chat_template_json_with_reasoning_effort(
1131        &self,
1132        messages: &[serde_json::Value],
1133        tools: Option<&[serde_json::Value]>,
1134        enable_thinking: Option<bool>,
1135        reasoning_effort: Option<&str>,
1136    ) -> Vec<u32> {
1137        match self.try_apply_chat_template_json_with_reasoning_effort(
1138            messages,
1139            tools,
1140            enable_thinking,
1141            reasoning_effort,
1142        ) {
1143            Ok(ids) => ids,
1144            Err(e) => {
1145                tracing::error!("chat template render failed ({e}); ChatML fallback");
1146                self.chatml_json_fallback(messages, enable_thinking)
1147            }
1148        }
1149    }
1150
1151    /// Like [`Self::apply_chat_template_json`], but a template that fails
1152    /// to render is an ERROR instead of a quiet ChatML approximation.
1153    ///
1154    /// The fallback flattens every message to (role, text): it has no
1155    /// place for `tools`, `tool_calls` or `role: "tool"`. For a plain chat
1156    /// that is a tolerable degradation; for a request with tools it means
1157    /// the model never sees the functions and answers as if none were
1158    /// offered — a failure no client can detect. The server calls this
1159    /// variant when tools are present and reports the error instead.
1160    /// Files without a template still take the ChatML path (Ok).
1161    pub fn try_apply_chat_template_json(
1162        &self,
1163        messages: &[serde_json::Value],
1164        tools: Option<&[serde_json::Value]>,
1165        enable_thinking: Option<bool>,
1166    ) -> Result<Vec<u32>, String> {
1167        self.try_apply_chat_template_json_with_reasoning_effort(
1168            messages,
1169            tools,
1170            enable_thinking,
1171            None,
1172        )
1173    }
1174
1175    /// Fallible JSON template rendering with a native reasoning-effort
1176    /// variable. MiniMax M3's upstream template is adaptive and consumes
1177    /// `<effort>` from this variable instead of a generic `<think>` injection.
1178    pub fn try_apply_chat_template_json_with_reasoning_effort(
1179        &self,
1180        messages: &[serde_json::Value],
1181        tools: Option<&[serde_json::Value]>,
1182        enable_thinking: Option<bool>,
1183        reasoning_effort: Option<&str>,
1184    ) -> Result<Vec<u32>, String> {
1185        if let Some(tpl) = &self.chat_template {
1186            return self
1187                .render_template_json(tpl, messages, tools, enable_thinking, reasoning_effort)
1188                .map(|text| self.with_bos(self.encode(&text)))
1189                .map_err(|e| format!("{e:#}"));
1190        }
1191        Ok(self.chatml_json_fallback(messages, enable_thinking))
1192    }
1193
1194    fn chatml_json_fallback(
1195        &self,
1196        messages: &[serde_json::Value],
1197        enable_thinking: Option<bool>,
1198    ) -> Vec<u32> {
1199        let pairs: Vec<(String, String)> = messages
1200            .iter()
1201            .map(|m| {
1202                (
1203                    m.get("role")
1204                        .and_then(|v| v.as_str())
1205                        .unwrap_or("user")
1206                        .to_string(),
1207                    m.get("content")
1208                        .and_then(|v| v.as_str())
1209                        .unwrap_or("")
1210                        .to_string(),
1211                )
1212            })
1213            .collect();
1214        self.with_bos(self.chatml_fallback_opts(&pairs, enable_thinking))
1215    }
1216
1217    /// Render the template against JSON-shaped messages (parity surface).
1218    pub fn render_chat_json(
1219        &self,
1220        messages: &[serde_json::Value],
1221        tools: Option<&[serde_json::Value]>,
1222        enable_thinking: Option<bool>,
1223    ) -> Option<String> {
1224        self.render_chat_json_with_reasoning_effort(messages, tools, enable_thinking, None)
1225    }
1226
1227    /// Render a JSON chat template with an optional native reasoning effort.
1228    /// Kept public for exact-template parity tests and protocol adapters.
1229    pub fn render_chat_json_with_reasoning_effort(
1230        &self,
1231        messages: &[serde_json::Value],
1232        tools: Option<&[serde_json::Value]>,
1233        enable_thinking: Option<bool>,
1234        reasoning_effort: Option<&str>,
1235    ) -> Option<String> {
1236        let tpl = self.chat_template.as_ref()?;
1237        match self.render_template_json(tpl, messages, tools, enable_thinking, reasoning_effort) {
1238            Ok(t) => Some(t),
1239            Err(e) => {
1240                tracing::error!("chat template render (json): {e:#}");
1241                eprintln!("chat template render (json): {e:#}");
1242                None
1243            }
1244        }
1245    }
1246
1247    fn render_template_json(
1248        &self,
1249        tpl: &str,
1250        messages: &[serde_json::Value],
1251        tools: Option<&[serde_json::Value]>,
1252        enable_thinking: Option<bool>,
1253        reasoning_effort: Option<&str>,
1254    ) -> Result<String, minijinja::Error> {
1255        let mut env = crate::chat_template::environment();
1256        let tpl_src = strip_generation_tags(tpl);
1257        env.add_template("chat", &tpl_src)?;
1258        let msgs: Vec<minijinja::Value> = messages
1259            .iter()
1260            .map(minijinja::Value::from_serialize)
1261            .collect();
1262        let tools_v: Option<Vec<minijinja::Value>> =
1263            tools.map(|ts| ts.iter().map(minijinja::Value::from_serialize).collect());
1264        let tpl = env.get_template("chat")?;
1265        // Three axes, each present only when meaningful: templates guard
1266        // with `is defined`, and an explicit null flips those guards.
1267        // `tool_call_format` picks the grammar in two-mode templates
1268        // (Nanbeige): undefined falls into their XML branch. The JSON
1269        // grammar is the one every parser downstream speaks, so choose
1270        // it explicitly; templates without the knob never read it.
1271        let rendered = match (tools_v, enable_thinking, reasoning_effort) {
1272            (Some(ts), Some(v), Some(effort)) => tpl.render(minijinja::context! {
1273                messages => msgs, tools => ts, add_generation_prompt => true, enable_thinking => v,
1274                reasoning_effort => effort,
1275                tool_call_format => "json",
1276            })?,
1277            (Some(ts), Some(v), None) => tpl.render(minijinja::context! {
1278                messages => msgs, tools => ts, add_generation_prompt => true, enable_thinking => v,
1279                // GLM-5's template controls its reasoning budget with
1280                // `reasoning_effort` rather than `enable_thinking`.  A direct
1281                // answer therefore needs the low budget when thinking is
1282                // disabled; otherwise the template defaults to `max`.
1283                reasoning_effort => if !v { Some("low") } else { None::<&str> },
1284                tool_call_format => "json",
1285            })?,
1286            (Some(ts), None, Some(effort)) => tpl.render(minijinja::context! {
1287                messages => msgs, tools => ts, add_generation_prompt => true,
1288                reasoning_effort => effort,
1289                tool_call_format => "json",
1290            })?,
1291            (Some(ts), None, None) => tpl.render(minijinja::context! {
1292                messages => msgs, tools => ts, add_generation_prompt => true,
1293                tool_call_format => "json",
1294            })?,
1295            (None, Some(v), Some(effort)) => tpl.render(minijinja::context! {
1296                messages => msgs, add_generation_prompt => true, enable_thinking => v,
1297                reasoning_effort => effort,
1298                tool_call_format => "json",
1299            })?,
1300            (None, Some(v), None) => tpl.render(minijinja::context! {
1301                messages => msgs, add_generation_prompt => true, enable_thinking => v,
1302                // See the tool-bearing branch above.
1303                reasoning_effort => if !v { Some("low") } else { None::<&str> },
1304                tool_call_format => "json",
1305            })?,
1306            (None, None, Some(effort)) => tpl.render(minijinja::context! {
1307                messages => msgs, add_generation_prompt => true,
1308                reasoning_effort => effort,
1309                tool_call_format => "json",
1310            })?,
1311            (None, None, None) => tpl.render(minijinja::context! {
1312                messages => msgs, add_generation_prompt => true,
1313                tool_call_format => "json",
1314            })?,
1315        };
1316        Ok(close_direct_think(rendered, enable_thinking))
1317    }
1318
1319    pub fn apply_chat_template_opts(
1320        &self,
1321        messages: &[(String, String)],
1322        enable_thinking: Option<bool>,
1323    ) -> Vec<u32> {
1324        if let Some(tpl) = &self.chat_template {
1325            match self.render_template(tpl, messages, enable_thinking) {
1326                Ok(text) => return self.with_bos(self.encode(&text)),
1327                Err(e) => {
1328                    tracing::error!("chat template render failed ({e}); ChatML fallback");
1329                }
1330            }
1331        }
1332        self.with_bos(self.chatml_fallback_opts(messages, enable_thinking))
1333    }
1334
1335    /// Prepend BOS when the tokenizer declares it (llama family).
1336    pub fn with_bos(&self, mut ids: Vec<u32>) -> Vec<u32> {
1337        if self.add_bos {
1338            if let Some(b) = self.bos_token_id {
1339                if ids.first() != Some(&b) {
1340                    ids.insert(0, b);
1341                }
1342            }
1343        }
1344        ids
1345    }
1346
1347    /// Render the carried template to text (parity-testable surface).
1348    pub fn render_chat(&self, messages: &[(String, String)]) -> Option<String> {
1349        self.render_chat_opts(messages, None)
1350    }
1351
1352    /// Render the carried template to text with explicit thinking mode.
1353    pub fn render_chat_opts(
1354        &self,
1355        messages: &[(String, String)],
1356        enable_thinking: Option<bool>,
1357    ) -> Option<String> {
1358        let tpl = self.chat_template.as_ref()?;
1359        match self.render_template(tpl, messages, enable_thinking) {
1360            Ok(t) => Some(t),
1361            Err(e) => {
1362                tracing::error!("chat template render: {e:#}");
1363                None
1364            }
1365        }
1366    }
1367
1368    fn render_template(
1369        &self,
1370        tpl: &str,
1371        messages: &[(String, String)],
1372        enable_thinking: Option<bool>,
1373    ) -> Result<String, minijinja::Error> {
1374        let mut env = crate::chat_template::environment();
1375        let tpl_src = strip_generation_tags(tpl);
1376        env.add_template("chat", &tpl_src)?;
1377        let msgs: Vec<minijinja::Value> = messages
1378            .iter()
1379            .map(|(role, content)| {
1380                minijinja::context! { role => role, content => content }
1381            })
1382            .collect();
1383        // `enable_thinking` stays UNDEFINED when None — reasoning templates
1384        // check `enable_thinking is defined` and fall back to their default.
1385        let rendered = match enable_thinking {
1386            Some(v) => env.get_template("chat")?.render(minijinja::context! {
1387                messages => msgs,
1388                add_generation_prompt => true,
1389                enable_thinking => v,
1390                // GLM-5 uses this variable (not `enable_thinking`) to select
1391                // its reasoning budget and defaults to max when omitted.
1392                reasoning_effort => if !v { Some("low") } else { None::<&str> },
1393            })?,
1394            None => env.get_template("chat")?.render(minijinja::context! {
1395                messages => msgs,
1396                add_generation_prompt => true,
1397            })?,
1398        };
1399        Ok(close_direct_think(rendered, enable_thinking))
1400    }
1401
1402    /// Hardcoded Qwen ChatML (pre-§6.1 files).
1403    fn chatml_fallback(&self, messages: &[(String, String)]) -> Vec<u32> {
1404        self.chatml_fallback_opts(messages, None)
1405    }
1406
1407    /// Hardcoded Qwen ChatML (pre-§6.1 files) with optional thinking suppression.
1408    fn chatml_fallback_opts(
1409        &self,
1410        messages: &[(String, String)],
1411        enable_thinking: Option<bool>,
1412    ) -> Vec<u32> {
1413        let mut tokens = Vec::new();
1414
1415        for (role, content) in messages {
1416            // <|im_start|>role\ncontent<|im_end|>\n
1417            if let Some(start_id) = self.im_start_id {
1418                tokens.push(start_id);
1419            }
1420            tokens.extend(self.encode(&format!("{}\n{}", role, content)));
1421            if let Some(end_id) = self.im_end_id {
1422                tokens.push(end_id);
1423            }
1424            tokens.extend(self.encode("\n"));
1425        }
1426
1427        // Add assistant prefix
1428        if let Some(start_id) = self.im_start_id {
1429            tokens.push(start_id);
1430        }
1431        tokens.extend(self.encode("assistant\n"));
1432        if enable_thinking == Some(false) {
1433            tokens.extend(self.encode("<think>\n\n</think>\n\n"));
1434        }
1435
1436        tokens
1437    }
1438
1439    /// Vocabulary size.
1440    pub fn vocab_size(&self) -> usize {
1441        self.id_to_token.len()
1442    }
1443
1444    /// Return the ID for an exact token spelling, including added/special
1445    /// tokens. Multimodal prompt preparation uses this to validate the image
1446    /// placeholder against the model configuration.
1447    pub fn token_to_id(&self, token: &str) -> Option<u32> {
1448        self.vocab.get(token).copied()
1449    }
1450
1451    /// Alias matching the HuggingFace tokenizer API used by the official
1452    /// DeepSeek image processor.
1453    pub fn convert_tokens_to_ids(&self, token: &str) -> Option<u32> {
1454        self.token_to_id(token)
1455    }
1456
1457    /// Check if token ID is EOS.
1458    pub fn is_eos(&self, id: u32) -> bool {
1459        self.eos_token_id == Some(id) || self.im_end_id == Some(id) || self.extra_eos.contains(&id)
1460    }
1461}
1462
1463#[derive(Debug, thiserror::Error)]
1464pub enum TokenizerError {
1465    #[error("IO error: {0}")]
1466    Io(String),
1467    #[error("Parse error: {0}")]
1468    Parse(String),
1469}
1470
1471/// SentencePiece merges over token ids: `(left, right)` → `(rank, merged)`.
1472struct SpMerges {
1473    pairs: HashMap<u64, (u32, u32)>,
1474    /// `<0xNN>` ids for the byte fallback (`u32::MAX` = absent)
1475    bytes: [u32; 256],
1476}
1477
1478const SP_NONE: u32 = u32::MAX;
1479
1480thread_local! {
1481    /// [`Tokenizer::encode_reference`] in progress on this thread.
1482    static SP_STRINGS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1483}
1484
1485impl SpMerges {
1486    fn build(vocab: &HashMap<String, u32>, ranks: &HashMap<(String, String), u32>) -> Option<Self> {
1487        if ranks.is_empty() {
1488            return None;
1489        }
1490        let mut pairs = HashMap::with_capacity(ranks.len());
1491        let mut cat = String::new();
1492        for ((a, b), &r) in ranks {
1493            let (&ia, &ib) = (vocab.get(a)?, vocab.get(b)?);
1494            cat.clear();
1495            cat.push_str(a);
1496            cat.push_str(b);
1497            let &m = vocab.get(&cat)?;
1498            pairs.insert(((ia as u64) << 32) | ib as u64, (r, m));
1499        }
1500        let mut bytes = [SP_NONE; 256];
1501        for (b, slot) in bytes.iter_mut().enumerate() {
1502            if let Some(&id) = vocab.get(&format!("<0x{b:02X}>")) {
1503                *slot = id;
1504            }
1505        }
1506        Some(SpMerges { pairs, bytes })
1507    }
1508
1509    #[inline]
1510    fn rank(&self, a: u32, b: u32) -> Option<(u32, u32)> {
1511        if a == SP_NONE || b == SP_NONE {
1512            return None;
1513        }
1514        self.pairs.get(&(((a as u64) << 32) | b as u64)).copied()
1515    }
1516
1517    /// The string form's rule over ids: take the lowest rank among the
1518    /// adjacent pairs, merge every live occurrence of that pair left to
1519    /// right (non-overlapping), and only then admit the pairs the merges
1520    /// formed — exactly the order of `bpe_piece_sp_strings`.
1521    fn encode(&self, vocab: &HashMap<String, u32>, piece: &str, out: &mut Vec<u32>) {
1522        use std::cmp::Reverse;
1523        use std::collections::BinaryHeap;
1524        let spans: Vec<(usize, usize)> = piece
1525            .char_indices()
1526            .map(|(i, c)| (i, i + c.len_utf8()))
1527            .collect();
1528        let n = spans.len();
1529        let mut buf = [0u8; 4];
1530        let mut id: Vec<u32> = piece
1531            .chars()
1532            .map(|c| *vocab.get(c.encode_utf8(&mut buf) as &str).unwrap_or(&SP_NONE))
1533            .collect();
1534        let mut next: Vec<usize> = (1..=n).collect();
1535        let mut prev: Vec<usize> = (0..n).map(|i| i.wrapping_sub(1)).collect();
1536        let mut alive = vec![true; n];
1537        let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(n);
1538        for i in 0..n.saturating_sub(1) {
1539            if let Some((r, _)) = self.rank(id[i], id[i + 1]) {
1540                heap.push(Reverse((r, i)));
1541            }
1542        }
1543        let mut batch: Vec<usize> = Vec::new();
1544        let mut merged: Vec<usize> = Vec::new();
1545        while let Some(&Reverse((r, _))) = heap.peek() {
1546            batch.clear();
1547            while let Some(&Reverse((r2, pos))) = heap.peek() {
1548                if r2 != r {
1549                    break;
1550                }
1551                heap.pop();
1552                batch.push(pos);
1553            }
1554            batch.sort_unstable();
1555            batch.dedup();
1556            merged.clear();
1557            for &pos in &batch {
1558                if !alive[pos] || next[pos] >= n {
1559                    continue;
1560                }
1561                let j = next[pos];
1562                match self.rank(id[pos], id[j]) {
1563                    Some((rr, m)) if rr == r => {
1564                        id[pos] = m;
1565                        alive[j] = false;
1566                        next[pos] = next[j];
1567                        if next[j] < n {
1568                            prev[next[j]] = pos;
1569                        }
1570                        merged.push(pos);
1571                    }
1572                    _ => {}
1573                }
1574            }
1575            for &pos in &merged {
1576                if !alive[pos] {
1577                    continue;
1578                }
1579                let pv = prev[pos];
1580                if pv < n && alive[pv] {
1581                    if let Some((r2, _)) = self.rank(id[pv], id[pos]) {
1582                        heap.push(Reverse((r2, pv)));
1583                    }
1584                }
1585                let nx = next[pos];
1586                if nx < n {
1587                    if let Some((r2, _)) = self.rank(id[pos], id[nx]) {
1588                        heap.push(Reverse((r2, pos)));
1589                    }
1590                }
1591            }
1592        }
1593        let mut i = 0usize;
1594        while i < n {
1595            if id[i] != SP_NONE {
1596                out.push(id[i]);
1597            } else {
1598                // a character outside the vocabulary: its UTF-8 bytes
1599                let (a, b) = spans[i];
1600                let mut ok = true;
1601                for &byte in piece[a..b].as_bytes() {
1602                    match self.bytes[byte as usize] {
1603                        SP_NONE => {
1604                            ok = false;
1605                            break;
1606                        }
1607                        t => out.push(t),
1608                    }
1609                }
1610                if !ok {
1611                    tracing::error!("tokenizer: no id for SP symbol {:?} — dropped", &piece[a..b]);
1612                }
1613            }
1614            i = next[i];
1615        }
1616    }
1617}
1618
1619#[cfg(test)]
1620mod tests {
1621    use super::*;
1622
1623    #[test]
1624    fn byte_unicode_bijection() {
1625        let (b2c, c2b) = bytes_to_unicode();
1626        for b in 0..=255u8 {
1627            assert_eq!(c2b[&b2c[b as usize]], b);
1628        }
1629        // GPT-2 well-known mappings: space → Ġ, newline → Ċ
1630        assert_eq!(b2c[b' ' as usize], 'Ġ');
1631        assert_eq!(b2c[b'\n' as usize], 'Ċ');
1632    }
1633
1634    #[test]
1635    fn byte_level_roundtrip_utf8() {
1636        let tok = Tokenizer::byte_level();
1637        let text = "hello 🌍 hi\n";
1638        let ids = tok.encode(text);
1639        assert_eq!(ids.len(), text.len()); // one id per byte
1640        assert_eq!(tok.decode(&ids), text);
1641    }
1642
1643    /// A tiny real-format tokenizer.json exercising the full pipeline:
1644    /// GPT-2 regex, byte-level alphabet, one merge, an added token.
1645    fn mini_json() -> String {
1646        // vocab: byte-level chars for h,e,l,o,Ġ,w,r,d + merged "he"
1647        let vocab: Vec<(&str, u32)> = vec![
1648            ("h", 0),
1649            ("e", 1),
1650            ("l", 2),
1651            ("o", 3),
1652            ("Ġ", 4),
1653            ("w", 5),
1654            ("r", 6),
1655            ("d", 7),
1656            ("he", 8),
1657            ("Ġw", 9),
1658        ];
1659        let vocab_json: String = vocab
1660            .iter()
1661            .map(|(t, i)| format!("\"{t}\": {i}"))
1662            .collect::<Vec<_>>()
1663            .join(", ");
1664        format!(
1665            r#"{{
1666              "model": {{
1667                "type": "BPE",
1668                "vocab": {{ {vocab_json} }},
1669                "merges": [["h", "e"], ["Ġ", "w"]]
1670              }},
1671              "added_tokens": [
1672                {{"id": 10, "content": "<|eot|>", "special": true}}
1673              ]
1674            }}"#
1675        )
1676    }
1677
1678    /// The id-level SentencePiece BPE gives the string form's ids on every
1679    /// input, including merges whose products rank below the merge that
1680    /// made them, overlapping runs ("aaa"), two merges producing one
1681    /// string, and characters outside the vocabulary (byte fallback, with
1682    /// one byte missing).
1683    #[test]
1684    fn sp_heap_bpe_matches_the_string_form() {
1685        let toks = [
1686            "a", "b", "c", "▁", "bc", "abc", "ab", "aa", "aab", "ca", "▁a", "aabc",
1687            "aaa", "<0x64>", "<0xC3>", "▁▁",
1688        ];
1689        let vocab: Vec<String> = toks
1690            .iter()
1691            .enumerate()
1692            .map(|(i, t)| format!("\"{t}\": {i}"))
1693            .collect();
1694        let merges = [
1695            ("aab", "c"),
1696            ("b", "c"),
1697            ("a", "bc"),
1698            ("a", "b"),
1699            ("ab", "c"),
1700            ("a", "a"),
1701            ("aa", "b"),
1702            ("c", "a"),
1703            ("▁", "a"),
1704            ("aa", "a"),
1705            ("▁", "▁"),
1706        ];
1707        let merges: Vec<String> = merges
1708            .iter()
1709            .map(|(a, b)| format!("[\"{a}\", \"{b}\"]"))
1710            .collect();
1711        let json = format!(
1712            r#"{{"model": {{"type": "BPE", "byte_fallback": true, "vocab": {{ {} }}, "merges": [{}]}},
1713               "normalizer": {{"type": "Replace", "pattern": {{"String": " "}}, "content": "\u2581"}},
1714               "added_tokens": []}}"#,
1715            vocab.join(", "),
1716            merges.join(", ")
1717        );
1718        let t = Tokenizer::from_json(&json).unwrap();
1719        assert!(t.metaspace, "a byte_fallback vocabulary takes the SentencePiece path");
1720        assert!(t.sp_merges().is_some());
1721        let alphabet = ['a', 'b', 'c', '\u{2581}', 'd', 'é', 'a', 'a', 'b'];
1722        let mut state = 0x2545_f491_4f6c_dd1du64;
1723        for len in 0..40 {
1724            for _ in 0..60 {
1725                let piece: String = (0..len)
1726                    .map(|_| {
1727                        state ^= state << 13;
1728                        state ^= state >> 7;
1729                        state ^= state << 17;
1730                        alphabet[(state % alphabet.len() as u64) as usize]
1731                    })
1732                    .collect();
1733                let (mut fast, mut slow) = (Vec::new(), Vec::new());
1734                t.sp_merges().unwrap().encode(&t.vocab, &piece, &mut fast);
1735                t.bpe_piece_sp_strings(&piece, &mut slow);
1736                assert_eq!(fast, slow, "{piece:?}");
1737            }
1738        }
1739    }
1740
1741    /// MiniCPM5 marks its tool grammar special. Decoding must keep the
1742    /// markup (the call IS that text) and still drop the chat control
1743    /// tokens around it.
1744    #[test]
1745    fn tool_markup_decodes_even_when_special() {
1746        let json = r#"{
1747          "model": {"type": "BPE", "vocab": {"h": 0, "e": 1, "l": 2, "o": 3}, "merges": []},
1748          "added_tokens": [
1749            {"id": 10, "content": "<|im_end|>", "special": true},
1750            {"id": 11, "content": "<function", "special": true},
1751            {"id": 12, "content": "</function>", "special": true},
1752            {"id": 13, "content": "<param", "special": true},
1753            {"id": 14, "content": "</param>", "special": true},
1754            {"id": 15, "content": "<tool_call>", "special": true}
1755          ]
1756        }"#;
1757        let t = Tokenizer::from_json(json).unwrap();
1758        let ids = [11, 0, 1, 13, 2, 14, 12, 15, 10];
1759        assert_eq!(
1760            t.decode(&ids),
1761            "<functionhe<paraml</param></function><tool_call>"
1762        );
1763        let streamed: String = ids.iter().map(|&i| t.decode_token(i)).collect();
1764        assert_eq!(streamed, t.decode(&ids), "streaming must agree with decode");
1765        assert!(
1766            !t.decode(&[10]).contains("im_end"),
1767            "control tokens stay hidden"
1768        );
1769    }
1770
1771    /// Parity against HuggingFace on a real tokenizer, run only when the
1772    /// file is present (CMF_TOK_PARITY=/path/to/tokenizer.json).
1773    #[test]
1774    fn real_tokenizer_parity_when_available() {
1775        let Ok(path) = std::env::var("CMF_TOK_PARITY") else {
1776            return;
1777        };
1778        let t = Tokenizer::from_file(&path).expect("load");
1779        for (text, want) in [
1780            (
1781                "The capital of France is",
1782                vec![671u32, 6102, 294, 8760, 344],
1783            ),
1784            ("2 + 2 =", vec![20, 940, 223, 20, 438]),
1785        ] {
1786            let got = t.encode(text);
1787            assert_eq!(got, want, "«{text}»");
1788        }
1789    }
1790
1791    /// Granite 4.2 ships its prompt grammar as a sidecar
1792    /// `chat_template.jinja` (not tokenizer_config.chat_template).  Exercise
1793    /// the exact upstream file when supplied so macro/namespace support and
1794    /// both reasoning prefixes cannot silently fall back to generic ChatML.
1795    #[test]
1796    fn granite_42_chat_template_when_available() {
1797        let Ok(path) = std::env::var("CMF_GRANITE_CHAT_TEMPLATE") else {
1798            return;
1799        };
1800        let mut tok = Tokenizer::byte_level();
1801        tok.chat_template = Some(std::fs::read_to_string(path).expect("read Granite template"));
1802        let messages = vec![("user".to_string(), "Hello".to_string())];
1803
1804        let thinking = tok
1805            .render_chat_opts(&messages, Some(true))
1806            .expect("render Granite thinking prompt");
1807        assert_eq!(
1808            thinking,
1809            "<|im_start|>system\n<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think>\n"
1810        );
1811
1812        let direct = tok
1813            .render_chat_opts(&messages, Some(false))
1814            .expect("render Granite direct prompt");
1815        assert_eq!(
1816            direct,
1817            "<|im_start|>system\n<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n<think></think>"
1818        );
1819    }
1820
1821    #[test]
1822    fn json_template_passes_explicit_reasoning_effort_without_overriding_it() {
1823        let mut tok = Tokenizer::byte_level();
1824        tok.chat_template = Some(
1825            "{% if reasoning_effort is defined %}<effort>{{ reasoning_effort }}</effort>{% else %}<default-effort>{% endif %}|{{ enable_thinking }}|{{ messages[0].reasoning_content }}"
1826                .to_string(),
1827        );
1828        let messages = vec![serde_json::json!({
1829            "role": "assistant",
1830            "content": "answer",
1831            "reasoning_content": "prior chain"
1832        })];
1833        let explicit = tok
1834            .render_chat_json_with_reasoning_effort(&messages, None, Some(false), Some("high"))
1835            .expect("render explicit effort");
1836        assert_eq!(explicit, "<effort>high</effort>|False|prior chain");
1837
1838        let native_default = tok
1839            .render_chat_json_with_reasoning_effort(&messages, None, None, None)
1840            .expect("render native default");
1841        assert_eq!(native_default, "<default-effort>||prior chain");
1842    }
1843
1844    #[test]
1845    fn fallible_json_template_render_does_not_use_chatml_on_jinja_error() {
1846        let mut tok = Tokenizer::byte_level();
1847        tok.chat_template = Some("{% for message in messages %}".to_string());
1848        let messages = vec![serde_json::json!({"role": "user", "content": "hello"})];
1849        let error = tok
1850            .try_apply_chat_template_json_with_reasoning_effort(
1851                &messages,
1852                None,
1853                None,
1854                Some("high"),
1855            )
1856            .expect_err("unclosed Jinja must fail instead of producing ChatML");
1857        assert!(!error.is_empty());
1858    }
1859
1860    /// MiniMax M3.1 carries an unusually featureful sidecar template: macros,
1861    /// `namespace`, Python-style mapping access, custom `tojson` keywords,
1862    /// a namespaced XML tool grammar and M3-specific reasoning markers.  Run
1863    /// its *actual* upstream template when supplied rather than accepting a
1864    /// synthetic toy template that would miss a minijinja compatibility gap.
1865    #[test]
1866    fn minimax_m3_chat_template_when_available() {
1867        let Ok(path) = std::env::var("CMF_MINIMAX_M3_CHAT_TEMPLATE") else {
1868            return;
1869        };
1870        let mut tok = Tokenizer::byte_level();
1871        tok.chat_template = Some(std::fs::read_to_string(path).expect("read MiniMax M3 template"));
1872        let messages = vec![
1873            serde_json::json!({"role": "system", "content": "Keep answers concise."}),
1874            serde_json::json!({"role": "user", "content": "What is 2 + 2?"}),
1875        ];
1876        let tools = vec![serde_json::json!({
1877            "type": "function",
1878            "function": {
1879                "name": "get_weather",
1880                "description": "Return the weather for a city.",
1881                "parameters": {
1882                    "type": "object",
1883                    "properties": {"city": {"type": "string"}},
1884                    "required": ["city"]
1885                }
1886            }
1887        })];
1888        let rendered = tok
1889            .render_chat_json(&messages, Some(&tools), Some(false))
1890            .expect("render MiniMax M3 native prompt");
1891        assert!(rendered.starts_with("]~!b[]~b]system\n"), "{rendered:?}");
1892        assert!(rendered.contains("<effort>low</effort>"), "{rendered:?}");
1893        assert!(
1894            rendered.contains("]<]minimax[>[<tool_call>"),
1895            "{rendered:?}"
1896        );
1897        assert!(rendered.contains("get_weather"), "{rendered:?}");
1898        assert!(rendered.ends_with("]~b]ai\n"), "{rendered:?}");
1899    }
1900
1901    /// GLM-5's generation suffix is `<|assistant|><think>`.  The direct-answer
1902    /// path must close that marker without splitting the `<|assistant|>` token
1903    /// (a bare `rfind("assistant")` inserts bytes in the middle of the marker).
1904    #[test]
1905    fn glm_direct_template_keeps_assistant_special_token_intact() {
1906        let mut tok = Tokenizer::byte_level();
1907        tok.chat_template = Some(
1908            "[gMASK]<sop>{%- set effort = reasoning_effort if reasoning_effort is defined and reasoning_effort in ['low', 'high'] else 'max' -%}<|system|>Reasoning Effort: {{ effort | capitalize }}{%- for m in messages -%}<|user|>{{ m.content }}{%- endfor -%}<|assistant|><think>"
1909                .to_string(),
1910        );
1911        let messages = vec![("user".to_string(), "2+2?".to_string())];
1912        let rendered = tok
1913            .render_chat_opts(&messages, Some(false))
1914            .expect("render GLM direct prompt");
1915        assert!(rendered.contains("Reasoning Effort: Low"));
1916        assert!(rendered.contains("<|assistant|><think></think>"));
1917        assert!(!rendered.contains("<|assistant\n"));
1918
1919        // Leaving thinking enabled must preserve GLM's template default
1920        // (`max`) rather than serializing our direct-answer hint.
1921        let thinking = tok
1922            .render_chat_opts(&messages, Some(true))
1923            .expect("render GLM thinking prompt");
1924        assert!(thinking.contains("Reasoning Effort: Max"));
1925
1926        let json_messages = vec![serde_json::json!({
1927            "role": "user",
1928            "content": "2+2?"
1929        })];
1930        let rendered_json = tok
1931            .render_chat_json(&json_messages, None, Some(false))
1932            .expect("render GLM JSON direct prompt");
1933        assert!(rendered_json.contains("Reasoning Effort: Low"));
1934        assert!(rendered_json.contains("<|assistant|><think></think>"));
1935        assert!(!rendered_json.contains("<|assistant\n"));
1936    }
1937
1938    /// A Sequence of Splits applies ALL of them, in order. Reading only the
1939    /// first is not a near-miss: DeepSeek-V4 puts a digit rule first and the
1940    /// word rule third, so one-pattern behaviour hands BPE a whole sentence
1941    /// as a single piece and the ids that come back are ones the model was
1942    /// never trained on.
1943    #[test]
1944    fn every_split_in_a_sequence_is_applied() {
1945        let pt = serde_json::json!({
1946            "type": "Sequence",
1947            "pretokenizers": [
1948                {"type": "Split", "behavior": "Isolated",
1949                 "pattern": {"Regex": r"\p{N}{1,3}"}},
1950                {"type": "Split", "behavior": "Isolated",
1951                 "pattern": {"Regex": r" ?[\p{L}]+"}},
1952                {"type": "ByteLevel", "add_prefix_space": false, "use_regex": false}
1953            ]
1954        });
1955        let mut pats = Vec::new();
1956        collect_split_patterns(&pt, &mut pats);
1957        assert_eq!(
1958            pats.len(),
1959            2,
1960            "both Split stages must be collected: {pats:?}"
1961        );
1962        assert!(pats[0].contains("p{N}"), "digit rule first");
1963        assert!(pats[1].contains("p{L}"), "word rule second");
1964
1965        // And the staged subdivision reaches the word boundaries. Building a
1966        // byte-level tokenizer over this pre_tokenizer, "ab cd" has to become
1967        // two pieces rather than one.
1968        let re: Vec<fancy_regex::Regex> = pats
1969            .iter()
1970            .map(|p| fancy_regex::Regex::new(p).unwrap())
1971            .collect();
1972        let norm = "ab cd12";
1973        let mut pieces: Vec<(usize, usize)> = vec![(0, norm.len())];
1974        for r in &re {
1975            let mut next = Vec::new();
1976            for (ps, pe) in pieces {
1977                let seg = &norm[ps..pe];
1978                let mut last = 0;
1979                for m in r.find_iter(seg).flatten() {
1980                    if m.start() > last {
1981                        next.push((ps + last, ps + m.start()));
1982                    }
1983                    if m.end() > m.start() {
1984                        next.push((ps + m.start(), ps + m.end()));
1985                    }
1986                    last = m.end();
1987                }
1988                if last < seg.len() {
1989                    next.push((ps + last, pe));
1990                }
1991            }
1992            pieces = next;
1993        }
1994        let got: Vec<&str> = pieces.iter().map(|(a, b)| &norm[*a..*b]).collect();
1995        assert_eq!(
1996            got,
1997            vec!["ab", " cd", "12"],
1998            "staged split produced {got:?}"
1999        );
2000    }
2001
2002    #[test]
2003    fn full_pipeline_merges_and_added_tokens() {
2004        let tok = Tokenizer::from_json(&mini_json()).unwrap();
2005        // "hello world" → [he,l,l,o, Ġw,o,r,l,d]
2006        let ids = tok.encode("hello world");
2007        assert_eq!(ids, vec![8, 2, 2, 3, 9, 3, 6, 2, 7]);
2008        assert_eq!(tok.decode(&ids), "hello world");
2009        // Added token splits and is skipped at decode (special).
2010        let ids2 = tok.encode("he<|eot|>he");
2011        assert_eq!(ids2, vec![8, 10, 8]);
2012        assert_eq!(tok.decode(&ids2), "hehe");
2013    }
2014
2015    #[test]
2016    fn non_ascii_is_never_silently_dropped() {
2017        let tok = Tokenizer::from_json(&mini_json()).unwrap();
2018        // A non-ASCII char is not encodable by the mini vocab (no byte tokens either):
2019        // the id list may be empty, but ASCII around it must survive.
2020        let ids = tok.encode("hello");
2021        assert!(!ids.is_empty());
2022    }
2023}
2024
2025#[cfg(test)]
2026mod generation_tag_tests {
2027    use super::strip_generation_tags;
2028
2029    /// LFM2.5's template wraps the assistant branch in `{%- generation -%}`.
2030    /// minijinja does not know the statement, the whole template failed,
2031    /// and the caller served a ChatML approximation instead — the model
2032    /// then answers a differently-shaped prompt than it was tuned on.
2033    #[test]
2034    fn a_generation_block_becomes_a_no_op_keeping_its_whitespace_control() {
2035        let tpl = "a{%- generation -%}b{%- endgeneration -%}c";
2036        let out = strip_generation_tags(tpl);
2037        assert!(!out.contains("{%- generation"));
2038        assert!(!out.contains("endgeneration"));
2039        // Both dashes survive on both tags: the trimming must not change.
2040        assert_eq!(out.matches("{%-").count(), 2);
2041        assert_eq!(out.matches("-%}").count(), 2);
2042        assert!(out.starts_with('a') && out.ends_with('c'));
2043    }
2044
2045    /// Whitespace control is per-side, and a tag without dashes must not
2046    /// grow any.
2047    #[test]
2048    fn each_side_keeps_its_own_dash() {
2049        let out = strip_generation_tags("{% generation %}x{%- endgeneration %}");
2050        assert!(out.starts_with("{% set"), "no dash added on the left");
2051        assert!(out.contains("{%- set"), "the right tag keeps its dash");
2052        assert!(!out.contains("-%}"), "no trailing dash invented");
2053    }
2054
2055    /// Templates that never use it are returned untouched, and other
2056    /// statements are never rewritten.
2057    #[test]
2058    fn everything_else_is_left_alone() {
2059        let plain = "{%- if x -%}{{ y }}{%- endif -%}";
2060        assert_eq!(strip_generation_tags(plain), plain);
2061        // The word appearing in TEXT is not a statement.
2062        let prose = "{{ 'the generation of tokens' }}";
2063        assert_eq!(strip_generation_tags(prose), prose);
2064    }
2065}