metrale_server/
tokenizer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ChatTokenizer`: a Hugging Face tokenizer, the chat template prompts are
4//! rendered with (minijinja, or the DeepSeek-V4 message encoder), and a streaming decoder.
5//!
6//! Owner: server (tokenizer).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use tokenizers::Tokenizer;
11
12/// 2026-09-26: Return a copy of `messages` in which every string
13/// `tool_calls[*].function.arguments` that parses as JSON is replaced by the parsed value;
14/// anything else is left as it is. Templates iterate the arguments as a map, e.g.
15/// `_args.items()` in jinja-templates/openai/minimax_m2.jinja.
16fn normalize_tool_call_arguments(messages: &[serde_json::Value]) -> Vec<serde_json::Value> {
17    let mut total_parsed = 0usize;
18    let mut total_seen = 0usize;
19    let out: Vec<_> = messages
20        .iter()
21        .map(|msg| {
22            let mut msg = msg.clone();
23            let Some(tool_calls) = msg.get_mut("tool_calls").and_then(|v| v.as_array_mut()) else {
24                return msg;
25            };
26            for tc in tool_calls.iter_mut() {
27                let Some(function) = tc.get_mut("function") else {
28                    continue;
29                };
30                let Some(args) = function.get_mut("arguments") else {
31                    continue;
32                };
33                total_seen += 1;
34                let parsed_owned = if let Some(s) = args.as_str() {
35                    serde_json::from_str::<serde_json::Value>(s).ok()
36                } else {
37                    None
38                };
39                if let Some(parsed) = parsed_owned {
40                    *args = parsed;
41                    total_parsed += 1;
42                }
43            }
44            msg
45        })
46        .collect();
47    if total_seen > 0 {
48        tracing::debug!(
49            "F76 normalize: {}/{} tool_call arguments parsed string→dict",
50            total_parsed,
51            total_seen,
52        );
53    }
54    out
55}
56
57mod chat_impl;
58pub(crate) mod chat_render;
59mod deepseek_v4;
60pub(crate) mod jinja_helpers;
61mod kimi_k3;
62mod message_preprocess;
63
64pub(crate) use message_preprocess::{
65    autoclose_assistant_think, remap_developer_role, resolve_think_control,
66};
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69enum ChatEncoding {
70    Jinja,
71    DeepseekV4,
72    KimiK3XtmlUnsupported,
73}
74
75pub struct ChatTokenizer {
76    tokenizer: Tokenizer,
77    eos_token_id: u32,
78    supports_thinking: bool,
79    chat_encoding: ChatEncoding,
80    /// 2026-09-26: The checkpoint's own Qwen template renders the tool schema, so the parser's
81    /// tool prompt is not added (`checkpoint_owns_qwen_tool_prompt` in chat_impl.rs).
82    native_qwen_tool_template: bool,
83    /// 2026-09-26: Template source `jinja_env` was built from: the override, the model's own
84    /// template, or the ChatML default; empty for an official Kimi K3 checkpoint.
85    #[allow(dead_code)]
86    chat_template: String,
87    /// 2026-09-26: Environment holding `chat_template`, compiled once in `from_model_dir`.
88    jinja_env: minijinja::Environment<'static>,
89    /// 2026-09-26: Environment for `jinja-templates/openai/{model_type}.jinja` when that file
90    /// exists and compiles; without it the OpenAI apply paths use `jinja_env`.
91    openai_jinja_env: Option<minijinja::Environment<'static>>,
92}
93
94/// 2026-09-26: `tokenizers::DecodeStream` with its generic parameters fixed.
95pub struct StreamingDecoder<'a> {
96    inner: tokenizers::DecodeStream<
97        'a,
98        tokenizers::models::ModelWrapper,
99        tokenizers::normalizers::NormalizerWrapper,
100        tokenizers::pre_tokenizers::PreTokenizerWrapper,
101        tokenizers::processors::PostProcessorWrapper,
102        tokenizers::decoders::DecoderWrapper,
103    >,
104}
105
106impl StreamingDecoder<'_> {
107    /// 2026-09-26: Feed one token. `Ok(Some(text))` returns the new text once the decode has
108    /// grown and does not end in U+FFFD; `Ok(None)` otherwise.
109    pub fn step(&mut self, id: u32) -> Result<Option<String>> {
110        self.inner
111            .step(id)
112            .map_err(|e| anyhow::anyhow!("Streaming decode error: {e}"))
113    }
114}
115
116#[cfg(test)]
117mod tests;
metrale_server/tokenizer.rs

metrale_server/
tokenizer.rs
