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;
65mod chat_render;
66mod deepseek_v4;
67mod jinja_helpers;
68mod message_preprocess;
69
70pub(crate) use message_preprocess::{
71    autoclose_assistant_think, remap_developer_role, resolve_think_control,
72};
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75enum ChatEncoding {
76    Jinja,
77    DeepseekV4,
78}
79
80pub struct ChatTokenizer {
81    tokenizer: Tokenizer,
82    eos_token_id: u32,
83    supports_thinking: bool,
84    chat_encoding: ChatEncoding,
85    /// Compiled Jinja chat template (from tokenizer_config.json).
86    #[allow(dead_code)]
87    chat_template: String,
88    /// Precompiled minijinja environment (avoids re-creating + re-compiling each call).
89    jinja_env: minijinja::Environment<'static>,
90    /// OpenAI-variant template: gates historical `<think>` wrappers on enable_thinking.
91    /// Falls back to jinja_env if no openai/ variant exists.
92    openai_jinja_env: Option<minijinja::Environment<'static>>,
93}
94
95/// Wrapper around tokenizers::DecodeStream that hides the generic parameters.
96/// O(1) per step vs O(n) for full re-decode.
97pub struct StreamingDecoder<'a> {
98    inner: tokenizers::DecodeStream<
99        'a,
100        tokenizers::models::ModelWrapper,
101        tokenizers::normalizers::NormalizerWrapper,
102        tokenizers::pre_tokenizers::PreTokenizerWrapper,
103        tokenizers::processors::PostProcessorWrapper,
104        tokenizers::decoders::DecoderWrapper,
105    >,
106}
107
108impl StreamingDecoder<'_> {
109    /// Feed one token. Returns Some(text) when valid UTF-8 is ready.
110    pub fn step(&mut self, id: u32) -> Result<Option<String>> {
111        self.inner
112            .step(id)
113            .map_err(|e| anyhow::anyhow!("Streaming decode error: {e}"))
114    }
115}
116
117#[cfg(test)]
118mod tests;
spark_server/tokenizer.rs

spark_server/
tokenizer.rs
