combs_formats/
tokenizer.rs

1//! Tokenizer specification returned by [`crate::ModelSource::tokenizer`].
2
3use std::collections::HashMap;
4use std::path::PathBuf;
5
6/// Where to find the tokenizer and which special tokens it defines.
7#[derive(Debug, Clone)]
8pub struct TokenizerSpec {
9    /// Path to the HuggingFace `tokenizer.json`.
10    pub tokenizer_json: PathBuf,
11    /// Added special tokens parsed from `tokenizer_config.json`
12    /// (`added_tokens_decoder`): token id → token string (e.g. `<|im_end|>`).
13    pub added_tokens: HashMap<u32, String>,
14    /// Raw chat template string (Jinja) if present in `tokenizer_config.json`.
15    /// Phase 1 uses a built-in ChatML wrap instead of evaluating Jinja.
16    pub chat_template: Option<String>,
17    /// Whether prompts should be prefixed with BOS (HF `add_bos_token` /
18    /// GGUF `tokenizer.ggml.add_bos_token`). `None` = unspecified, which the
19    /// engine treats as "prepend when the model declares a BOS id". Qwen2
20    /// declares a BOS id but sets this to `false`; ignoring it prepends
21    /// `<|endoftext|>` to every prompt.
22    pub add_bos: Option<bool>,
23}
24
25impl TokenizerSpec {
26    /// Dummy spec used by weight-only safetensors sources that have no
27    /// tokenizer. `tokenizer()` on those sources errors before this is used.
28    pub(crate) fn placeholder() -> Self {
29        Self {
30            tokenizer_json: PathBuf::new(),
31            added_tokens: HashMap::new(),
32            chat_template: None,
33            add_bos: None,
34        }
35    }
36
37    /// Looks up the id of an added special token by its string, e.g.
38    /// `spec.special_token_id("<|im_end|>")`.
39    pub fn special_token_id(&self, token: &str) -> Option<u32> {
40        self.added_tokens
41            .iter()
42            .find(|(_, s)| s.as_str() == token)
43            .map(|(id, _)| *id)
44    }
45
46    /// Wraps a user prompt in the ChatML template used by SmolLM2-style
47    /// instruction models. Returns `None` if the tokenizer does not define
48    /// `<|im_start|>` / `<|im_end|>` tokens.
49    pub fn chatml_wrap(&self, user_prompt: &str) -> Option<String> {
50        let start = self.special_token_id("<|im_start|>")?;
51        self.special_token_id("<|im_end|>")?;
52        let _ = start;
53        Some(format!(
54            "<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
55        ))
56    }
57
58    /// Which chat template this tokenizer's special tokens imply.
59    pub fn chat_template_kind(&self) -> ChatTemplate {
60        if self.special_token_id("<start_of_turn>").is_some() {
61            ChatTemplate::Gemma
62        } else {
63            ChatTemplate::Chatml
64        }
65    }
66
67    /// Wraps (role, content) message pairs into the model's chat format,
68    /// ending with the assistant turn left open for generation. Unknown
69    /// roles are coerced to `user` (same convention as serve).
70    pub fn wrap_messages(&self, messages: &[(String, String)]) -> String {
71        match self.chat_template_kind() {
72            ChatTemplate::Chatml => {
73                let mut out = String::new();
74                for (role, content) in messages {
75                    let role = match role.as_str() {
76                        "system" | "user" | "assistant" => role.as_str(),
77                        _ => "user",
78                    };
79                    out.push_str(&format!("<|im_start|>{role}\n{content}<|im_end|>\n"));
80                }
81                out.push_str("<|im_start|>assistant\n");
82                out
83            }
84            ChatTemplate::Gemma => {
85                // Gemma-3 template: <bos> prefix, turns are
86                // `<start_of_turn>{user|model}\n{content}<end_of_turn>\n`,
87                // and system content is folded into the first user turn
88                // (the HF template prepends it with a blank line).
89                let mut out = String::from("<bos>");
90                let mut system_prefix = String::new();
91                let mut first_user_seen = false;
92                for (role, content) in messages {
93                    match role.as_str() {
94                        "system" => {
95                            system_prefix.push_str(content);
96                            system_prefix.push_str("\n\n");
97                        }
98                        _ => {
99                            let turn_role = match role.as_str() {
100                                "assistant" | "model" => "model",
101                                _ => "user",
102                            };
103                            let mut body = String::new();
104                            if turn_role == "user" && !first_user_seen {
105                                body.push_str(&system_prefix);
106                                first_user_seen = true;
107                            }
108                            body.push_str(content);
109                            out.push_str(&format!(
110                                "<start_of_turn>{turn_role}\n{body}<end_of_turn>\n"
111                            ));
112                        }
113                    }
114                }
115                out.push_str("<start_of_turn>model\n");
116                out
117            }
118        }
119    }
120}
121
122/// Chat template flavor detected from the tokenizer's special tokens.
123#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub enum ChatTemplate {
125    /// `<|im_start|>role\n…<|im_end|>` (SmolLM2 / Qwen style).
126    Chatml,
127    /// `<start_of_turn>user\n…<end_of_turn>` (Gemma style).
128    Gemma,
129}
combs_formats/tokenizer.rs

combs_formats/
tokenizer.rs
