spark_server/
tokenizer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Tokenizer wrapper using HuggingFace tokenizers + minijinja chat template.
4//!
5//! Loads the model's official Jinja template from `tokenizer_config.json` and
6//! renders it with minijinja for byte-exact alignment with the model's training
7//! format. No fallback — if there's no Jinja template, the model is misconfigured.
8
9use anyhow::Result;
10use tokenizers::Tokenizer;
11
12/// F76 (2026-04-29): pre-parse `tool_calls[*].function.arguments` from
13/// OpenAI's wire format (JSON-encoded string) into the JSON value the
14/// model's chat template expects. MiniMax M2.7's template iterates
15/// `tool_call.function.arguments.items()` which crashes on a string.
16/// We rebuild the message list with parsed arguments where present,
17/// leaving every other field untouched. Returns a fresh Vec rather
18/// than mutating the caller's slice.
19fn normalize_tool_call_arguments(messages: &[serde_json::Value]) -> Vec<serde_json::Value> {
20    let mut total_parsed = 0usize;
21    let mut total_seen = 0usize;
22    let out: Vec<_> = messages
23        .iter()
24        .map(|msg| {
25            let mut msg = msg.clone();
26            let Some(tool_calls) = msg.get_mut("tool_calls").and_then(|v| v.as_array_mut()) else {
27                return msg;
28            };
29            for tc in tool_calls.iter_mut() {
30                let Some(function) = tc.get_mut("function") else {
31                    continue;
32                };
33                let Some(args) = function.get_mut("arguments") else {
34                    continue;
35                };
36                total_seen += 1;
37                let parsed_owned = if let Some(s) = args.as_str() {
38                    serde_json::from_str::<serde_json::Value>(s).ok()
39                } else {
40                    None
41                };
42                if let Some(parsed) = parsed_owned {
43                    *args = parsed;
44                    total_parsed += 1;
45                }
46                // If parse fails or args wasn't a string, leave as-is —
47                // template may handle via tojson, or surface the
48                // original error for the operator.
49            }
50            msg
51        })
52        .collect();
53    if total_seen > 0 {
54        tracing::debug!(
55            "F76 normalize: {}/{} tool_call arguments parsed string→dict",
56            total_parsed,
57            total_seen,
58        );
59    }
60    out
61}
62
63/// Wraps a HuggingFace tokenizer with Jinja chat template support.
64mod chat_impl;
65pub(crate) mod chat_render;
66mod deepseek_v4;
67pub(crate) mod jinja_helpers;
68mod kimi_k3;
69mod message_preprocess;
70
71pub(crate) use message_preprocess::{
72    autoclose_assistant_think, remap_developer_role, resolve_think_control,
73};
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76enum ChatEncoding {
77    Jinja,
78    DeepseekV4,
79    KimiK3XtmlUnsupported,
80}
81
82pub struct ChatTokenizer {
83    tokenizer: Tokenizer,
84    eos_token_id: u32,
85    supports_thinking: bool,
86    chat_encoding: ChatEncoding,
87    /// Checkpoint Qwen tool instructions own schema rendering (no override selected).
88    native_qwen_tool_template: bool,
89    /// Compiled Jinja chat template (from tokenizer_config.json).
90    #[allow(dead_code)]
91    chat_template: String,
92    /// Precompiled minijinja environment (avoids re-creating + re-compiling each call).
93    jinja_env: minijinja::Environment<'static>,
94    /// OpenAI-variant template: gates historical `<think>` wrappers on enable_thinking.
95    /// Falls back to jinja_env if no openai/ variant exists.
96    openai_jinja_env: Option<minijinja::Environment<'static>>,
97}
98
99/// Wrapper around tokenizers::DecodeStream that hides the generic parameters.
100/// O(1) per step vs O(n) for full re-decode.
101pub struct StreamingDecoder<'a> {
102    inner: tokenizers::DecodeStream<
103        'a,
104        tokenizers::models::ModelWrapper,
105        tokenizers::normalizers::NormalizerWrapper,
106        tokenizers::pre_tokenizers::PreTokenizerWrapper,
107        tokenizers::processors::PostProcessorWrapper,
108        tokenizers::decoders::DecoderWrapper,
109    >,
110}
111
112impl StreamingDecoder<'_> {
113    /// Feed one token. Returns Some(text) when valid UTF-8 is ready.
114    pub fn step(&mut self, id: u32) -> Result<Option<String>> {
115        self.inner
116            .step(id)
117            .map_err(|e| anyhow::anyhow!("Streaming decode error: {e}"))
118    }
119}
120
121#[cfg(test)]
122mod tests;
spark_server/tokenizer.rs

spark_server/
tokenizer.rs
