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}