metrale_server/tokenizer/
chat_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ChatTokenizer` construction (template selection), encode and decode, and the
4//! chat-template apply paths.
5//!
6//! Owner: server (tokenizer).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use std::path::Path;
11use tokenizers::Tokenizer;
12
13#[path = "chat_impl/vision_pads.rs"]
14mod vision_pads;
15
16use super::{
17    ChatEncoding, ChatTokenizer, StreamingDecoder, autoclose_assistant_think,
18    normalize_tool_call_arguments, remap_developer_role, resolve_think_control,
19};
20
21/// 2026-09-26: Message rewrites applied before every Jinja render (`render_chat`), in order:
22/// parse string tool-call arguments (`normalize_tool_call_arguments`), turn `developer`
23/// messages into `system` (`remap_developer_role`), close an open `<think>` before a
24/// `<tool_call>` in assistant history (`autoclose_assistant_think`), and strip inline
25/// `<|think_on|>`/`<|think_off|>` tokens (`resolve_think_control`).
26///
27/// Returns the rewritten messages and the thinking flag to render with: the last inline
28/// control token when there is one, else `enable_thinking`.
29pub(crate) fn preprocess_for_render(
30    messages: &[serde_json::Value],
31    enable_thinking: bool,
32) -> (Vec<serde_json::Value>, bool) {
33    let prepared = normalize_tool_call_arguments(messages);
34    let mut prepared = remap_developer_role(prepared);
35    autoclose_assistant_think(&mut prepared);
36    let (prepared, control_override) = resolve_think_control(&prepared);
37    let effective_thinking = control_override.unwrap_or(enable_thinking);
38    (prepared, effective_thinking)
39}
40
41impl ChatTokenizer {
42    pub fn from_model_dir(
43        model_dir: &Path,
44        eos_token_id: u32,
45        supports_thinking: bool,
46        model_type: &str,
47        repo_root: Option<&Path>,
48        disable_template_overrides: bool,
49    ) -> Result<Self> {
50        let official_k3 = super::kimi_k3::uses_xtml(model_dir, model_type)?;
51        let tokenizer_path = model_dir.join("tokenizer.json");
52        let mut tokenizer = Tokenizer::from_file(&tokenizer_path)
53            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {e}"))?;
54        tokenizer
55            .with_truncation(None)
56            .map_err(|e| anyhow::anyhow!("Failed to disable tokenizer truncation: {e}"))?;
57
58        // 2026-09-26: Template source, first match wins: an official Kimi K3 checkpoint gets
59        // no chat template; then `jinja-templates/{model_type}.jinja` when the file exists and
60        // `--disable-template-overrides` is off; then the model's own template
61        // (`load_config_template`); then the ChatML default.
62        let override_tmpl = if disable_template_overrides {
63            None
64        } else {
65            super::jinja_helpers::load_override_template(model_type, repo_root)
66        };
67        let (chat_template, checkpoint_template) = if official_k3 {
68            tracing::warn!("Official Kimi K3: raw completions only; XTML chat is unavailable");
69            (String::new(), false)
70        } else if let Some(override_tmpl) = override_tmpl {
71            (override_tmpl, false)
72        } else if let Some(config_tmpl) = super::jinja_helpers::load_config_template(model_dir)? {
73            (config_tmpl, true)
74        } else {
75            tracing::warn!("No chat template found — using default ChatML");
76            (
77                super::jinja_helpers::default_chatml_template(supports_thinking),
78                false,
79            )
80        };
81
82        let jinja_env = super::jinja_helpers::build_jinja_env(&chat_template)?;
83
84        // 2026-09-26: A variant template that fails to compile is dropped without a log line.
85        let openai_jinja_env = super::jinja_helpers::load_openai_template(model_type, repo_root)
86            .and_then(|tmpl| {
87                tracing::info!("Loaded OpenAI-variant Jinja template for {model_type}");
88                super::jinja_helpers::build_jinja_env(&tmpl).ok()
89            });
90        let chat_encoding = if official_k3 {
91            ChatEncoding::KimiK3XtmlUnsupported
92        } else if model_type == "deepseek_v4" || model_type == "deepseek_v41" {
93            tracing::info!("Using checkpoint-native DeepSeek-V4 message encoding");
94            ChatEncoding::DeepseekV4
95        } else {
96            ChatEncoding::Jinja
97        };
98
99        let native_qwen_tool_template = checkpoint_owns_qwen_tool_prompt(
100            model_type,
101            checkpoint_template,
102            openai_jinja_env.is_some(),
103        ) && jinja_env
104            .get_template("chat")?
105            .undeclared_variables(false)
106            .contains("tools");
107        tracing::info!("Loaded tokenizer from {}", tokenizer_path.display());
108        Ok(Self {
109            tokenizer,
110            eos_token_id,
111            supports_thinking,
112            chat_encoding,
113            native_qwen_tool_template,
114            chat_template,
115            jinja_env,
116            openai_jinja_env,
117        })
118    }
119
120    pub(crate) fn uses_native_qwen_tool_template(&self) -> bool {
121        self.native_qwen_tool_template
122    }
123
124    /// 2026-09-26: The underlying Hugging Face tokenizer.
125    pub fn inner(&self) -> &tokenizers::Tokenizer {
126        &self.tokenizer
127    }
128
129    pub fn encode(&self, text: &str) -> Result<Vec<u32>> {
130        let encoding = self
131            .tokenizer
132            .encode(text, false)
133            .map_err(|e| anyhow::anyhow!("Tokenizer encode error: {e}"))?;
134        Ok(encoding.get_ids().to_vec())
135    }
136
137    pub fn decode(&self, ids: &[u32]) -> Result<String> {
138        self.tokenizer
139            .decode(ids, true)
140            .map_err(|e| anyhow::anyhow!("Tokenizer decode error: {e}"))
141    }
142
143    /// 2026-09-26: Decode without skipping special tokens; `decode` skips them.
144    pub fn decode_with_special(&self, ids: &[u32]) -> Result<String> {
145        self.tokenizer
146            .decode(ids, false)
147            .map_err(|e| anyhow::anyhow!("Tokenizer decode error: {e}"))
148    }
149
150    /// 2026-09-26: Incremental detokenizer. Returns the text `toks` added since the last call
151    /// and advances the offsets.
152    ///
153    /// Each call decodes only `toks[prefix_offset..read_offset]` and `toks[prefix_offset..]`
154    /// (with `decode`, so special tokens are skipped) and returns what the second adds to the
155    /// first. It returns an empty string and leaves the offsets alone when the second decode
156    /// is not longer, ends in U+FFFD, or does not split at a char boundary. Offsets that no
157    /// longer fit `toks` are reset to 0 first.
158    pub fn incremental_decode(
159        &self,
160        toks: &[u32],
161        prefix_offset: &mut usize,
162        read_offset: &mut usize,
163    ) -> String {
164        if *read_offset > toks.len() || *prefix_offset > *read_offset {
165            *prefix_offset = 0;
166            *read_offset = 0;
167        }
168        let prefix_text = self
169            .decode(&toks[*prefix_offset..*read_offset])
170            .unwrap_or_default();
171        let new_text = self.decode(&toks[*prefix_offset..]).unwrap_or_default();
172        if new_text.len() > prefix_text.len()
173            && !new_text.ends_with('\u{FFFD}')
174            && let Some(delta) = new_text.get(prefix_text.len()..)
175        {
176            let delta = delta.to_string();
177            *prefix_offset = *read_offset;
178            *read_offset = toks.len();
179            return delta;
180        }
181        String::new()
182    }
183
184    /// 2026-09-26: A new [`StreamingDecoder`] over this tokenizer.
185    pub fn streaming_decoder(&self, skip_special_tokens: bool) -> StreamingDecoder<'_> {
186        StreamingDecoder {
187            inner: self.tokenizer.decode_stream(skip_special_tokens),
188        }
189    }
190
191    /// 2026-09-26: `apply_chat_template_jinja_with_effort` with no reasoning effort and
192    /// `preserve_thinking` unset.
193    pub fn apply_chat_template_jinja(
194        &self,
195        messages: &[serde_json::Value],
196        tools: Option<&[serde_json::Value]>,
197        enable_thinking: bool,
198        disable_tool_steering: bool,
199    ) -> Result<Vec<u32>> {
200        self.apply_chat_template_jinja_with_effort(
201            messages,
202            tools,
203            enable_thinking,
204            disable_tool_steering,
205            None,
206            None,
207        )
208    }
209
210    pub fn apply_chat_template_jinja_with_effort(
211        &self,
212        messages: &[serde_json::Value],
213        tools: Option<&[serde_json::Value]>,
214        enable_thinking: bool,
215        disable_tool_steering: bool,
216        reasoning_effort: Option<&str>,
217        preserve_thinking: Option<bool>,
218    ) -> Result<Vec<u32>> {
219        super::kimi_k3::require_chat_support(self.chat_encoding)?;
220        if self.chat_encoding == ChatEncoding::DeepseekV4 {
221            let rendered = super::deepseek_v4::encode_messages(
222                messages,
223                tools,
224                enable_thinking,
225                reasoning_effort,
226            )?;
227            return self.encode(&rendered);
228        }
229
230        let rendered = super::chat_render::render_chat(
231            &self.jinja_env,
232            messages,
233            tools,
234            super::chat_render::RenderFlags {
235                enable_thinking,
236                disable_tool_steering,
237                reasoning_effort,
238                preserve_thinking,
239                allow_continue_final: true,
240            },
241        )?;
242
243        if rendered.len() < 2000 {
244            let tail_start = rendered.floor_char_boundary(rendered.len().saturating_sub(200));
245            tracing::info!(
246                "Jinja rendered ({} chars): {:?}",
247                rendered.len(),
248                &rendered[tail_start..]
249            );
250        }
251
252        self.encode(&rendered)
253    }
254
255    /// 2026-09-26: `apply_chat_template_openai_with_effort` with no reasoning effort and
256    /// `preserve_thinking` unset.
257    pub fn apply_chat_template_openai(
258        &self,
259        messages: &[serde_json::Value],
260        tools: Option<&[serde_json::Value]>,
261        enable_thinking: bool,
262        disable_tool_steering: bool,
263    ) -> Result<Vec<u32>> {
264        self.apply_chat_template_openai_with_effort(
265            messages,
266            tools,
267            enable_thinking,
268            disable_tool_steering,
269            None,
270            None,
271        )
272    }
273
274    pub fn apply_chat_template_openai_with_effort(
275        &self,
276        messages: &[serde_json::Value],
277        tools: Option<&[serde_json::Value]>,
278        enable_thinking: bool,
279        disable_tool_steering: bool,
280        reasoning_effort: Option<&str>,
281        preserve_thinking: Option<bool>,
282    ) -> Result<Vec<u32>> {
283        super::kimi_k3::require_chat_support(self.chat_encoding)?;
284        if self.chat_encoding == ChatEncoding::DeepseekV4 {
285            return self.apply_chat_template_jinja_with_effort(
286                messages,
287                tools,
288                enable_thinking,
289                disable_tool_steering,
290                reasoning_effort,
291                preserve_thinking,
292            );
293        }
294        if let Some(ref env) = self.openai_jinja_env {
295            let rendered = super::chat_render::render_chat(
296                env,
297                messages,
298                tools,
299                super::chat_render::RenderFlags {
300                    enable_thinking,
301                    disable_tool_steering,
302                    reasoning_effort,
303                    preserve_thinking,
304                    allow_continue_final: false,
305                },
306            )
307            .map_err(|e| anyhow::anyhow!("Failed to render OpenAI Jinja template: {e}"))?;
308            self.encode(&rendered)
309        } else {
310            self.apply_chat_template_jinja_with_effort(
311                messages,
312                tools,
313                enable_thinking,
314                disable_tool_steering,
315                reasoning_effort,
316                preserve_thinking,
317            )
318        }
319    }
320
321    /// 2026-09-26: Render `(role, content)` pairs through `apply_chat_template_jinja`, with no
322    /// tools; `_image_pad_counts` is unused.
323    pub fn apply_chat_template(
324        &self,
325        messages: &[(String, String)],
326        enable_thinking: bool,
327        _image_pad_counts: &[usize],
328    ) -> Result<Vec<u32>> {
329        let json_messages: Vec<serde_json::Value> = messages
330            .iter()
331            .map(|(role, content)| {
332                serde_json::json!({
333                    "role": role,
334                    "content": content,
335                })
336            })
337            .collect();
338
339        self.apply_chat_template_jinja(&json_messages, None, enable_thinking, false)
340    }
341
342    pub fn eos_token_id(&self) -> u32 {
343        self.eos_token_id
344    }
345
346    pub fn think_end_token_id(&self) -> Option<u32> {
347        if !self.supports_thinking {
348            return None;
349        }
350        match self.encode("</think>") {
351            Ok(ids) if ids.len() == 1 => Some(ids[0]),
352            _ => None,
353        }
354    }
355
356    pub fn supports_thinking(&self) -> bool {
357        self.supports_thinking
358    }
359
360    pub fn uses_deepseek_v4_encoding(&self) -> bool {
361        self.chat_encoding == ChatEncoding::DeepseekV4
362    }
363}
364
365/// 2026-09-26: True only for the checkpoint's own template (not an override or the ChatML
366/// default), with no OpenAI variant, on model types `qwen3_5`, `qwen3_5_moe`, `qwen3_6` and
367/// `qwen3_6_moe`. `from_model_dir` also requires the template to use `tools`.
368fn checkpoint_owns_qwen_tool_prompt(
369    model_type: &str,
370    checkpoint: bool,
371    openai_override: bool,
372) -> bool {
373    checkpoint
374        && !openai_override
375        && matches!(
376            model_type,
377            "qwen3_5" | "qwen3_5_moe" | "qwen3_6" | "qwen3_6_moe"
378        )
379}
metrale_server/tokenizer/chat_impl.rs

metrale_server/tokenizer/
chat_impl.rs
