cherenkov/
tok.rs

1use anyhow::{Context, Result};
2use std::path::Path;
3use tokenizers::Tokenizer;
4
5pub struct ChatTokenizer {
6    pub inner: Tokenizer,
7    pub template: Option<crate::prompt::ChatTemplate>,
8    pub im_end: u32,
9    pub endoftext: u32,
10}
11
12impl ChatTokenizer {
13    pub fn load(model_dir: &Path) -> Result<Self> {
14        let path = model_dir.join("tokenizer.json");
15        let mut inner = Tokenizer::from_file(&path)
16            .map_err(|e| anyhow::anyhow!("loading {}: {e}", path.display()))?;
17
18        // Checkpoints may save training-time padding and truncation. Inference
19        // uses actual prompt lengths and checks its own context budget.
20        inner.with_padding(None);
21        inner
22            .with_truncation(None)
23            .map_err(|e| anyhow::anyhow!("tokenizer truncation: {e}"))?;
24
25        let tok = |s: &str| -> Result<u32> {
26            inner
27                .token_to_id(s)
28                .with_context(|| format!("special token {s:?} missing"))
29        };
30        let im_end = tok("<|im_end|>")?;
31        let endoftext = tok("<|endoftext|>")?;
32
33        Ok(ChatTokenizer {
34            inner,
35            template: crate::prompt::ChatTemplate::load(model_dir)?,
36            im_end,
37            endoftext,
38        })
39    }
40
41    pub fn encode(&self, text: &str) -> Result<Vec<u32>> {
42        // ChatML scaffolding is already present; do not add template tokens.
43        let enc = self
44            .inner
45            .encode(text, false)
46            .map_err(|e| anyhow::anyhow!("encode: {e}"))?;
47
48        Ok(enc.get_ids().to_vec())
49    }
50
51    pub fn decode(&self, ids: &[u32]) -> Result<String> {
52        self.inner
53            .decode(ids, false)
54            .map_err(|e| anyhow::anyhow!("decode: {e}"))
55    }
56}
cherenkov/tok.rs

cherenkov/
tok.rs
