Skip to main content

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}
18
19impl TokenizerSpec {
20    /// Looks up the id of an added special token by its string, e.g.
21    /// `spec.special_token_id("<|im_end|>")`.
22    pub fn special_token_id(&self, token: &str) -> Option<u32> {
23        self.added_tokens
24            .iter()
25            .find(|(_, s)| s.as_str() == token)
26            .map(|(id, _)| *id)
27    }
28
29    /// Wraps a user prompt in the ChatML template used by SmolLM2-style
30    /// instruction models. Returns `None` if the tokenizer does not define
31    /// `<|im_start|>` / `<|im_end|>` tokens.
32    pub fn chatml_wrap(&self, user_prompt: &str) -> Option<String> {
33        let start = self.special_token_id("<|im_start|>")?;
34        self.special_token_id("<|im_end|>")?;
35        let _ = start;
36        Some(format!(
37            "<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
38        ))
39    }
40}