Skip to main content

jev_bridge/
tokenizer.rs

1//! Local tokenization through the Hugging Face `tokenizers` crate.
2//!
3//! The serving runtime's `/tokenize` endpoint is enough when it exists, and the
4//! bridge prefers it because it cannot drift from the runtime's own vocabulary.
5//! A `tokenizer.json` file makes the whole contract independent of the runtime:
6//! the answer-slot check, the boundary check and the prompt length all run
7//! in-process, which matters for runtimes without a tokenize endpoint and for
8//! deployments that must not depend on extra HTTP round trips.
9//!
10//! The file must be the model's own `tokenizer.json`. A different vocabulary
11//! would silently score the wrong token ids.
12
13use std::path::Path;
14
15use anyhow::{bail, Result};
16use tokenizers::Tokenizer;
17
18/// A tokenizer loaded from a local `tokenizer.json`.
19pub struct LocalTokenizer {
20    tokenizer: Tokenizer,
21}
22
23impl LocalTokenizer {
24    /// Load a `tokenizer.json` from disk.
25    pub fn from_file(path: &Path) -> Result<Self> {
26        let tokenizer = Tokenizer::from_file(path).map_err(|error| {
27            anyhow::anyhow!("loading {} failed: {error}", path.display())
28        })?;
29        Ok(Self { tokenizer })
30    }
31
32    /// Load a `tokenizer.json` already in memory.
33    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
34        let tokenizer = Tokenizer::from_bytes(bytes)
35            .map_err(|error| anyhow::anyhow!("parsing the tokenizer failed: {error}"))?;
36        Ok(Self { tokenizer })
37    }
38
39    /// Encode without special tokens, matching fastjev's `encode_prompt`.
40    pub fn tokenize(&self, content: &str) -> Result<Vec<u32>> {
41        let encoding = self
42            .tokenizer
43            .encode(content, false)
44            .map_err(|error| anyhow::anyhow!("tokenizing failed: {error}"))?;
45        Ok(encoding.get_ids().to_vec())
46    }
47
48    /// Decode skipping special tokens, matching fastjev's round-trip check.
49    pub fn detokenize(&self, tokens: &[u32]) -> Result<String> {
50        self.tokenizer
51            .decode(tokens, true)
52            .map_err(|error| anyhow::anyhow!("detokenizing failed: {error}"))
53    }
54
55    /// The vocabulary size, for startup diagnostics.
56    pub fn vocab_size(&self) -> usize {
57        self.tokenizer.get_vocab_size(true)
58    }
59
60    /// Refuse a tokenizer whose special tokens disagree with the template.
61    ///
62    /// A template that emits `{{ bos_token }}` and a tokenizer whose BOS token
63    /// is something else would produce a prompt the model never saw in
64    /// training, so this is checked rather than assumed.
65    pub fn verify_special_tokens(&self, expected: &serde_json::Map<String, serde_json::Value>) -> Result<()> {
66        for (name, value) in expected {
67            let Some(text) = value.as_str() else {
68                continue;
69            };
70            if !name.ends_with("_token") {
71                continue;
72            }
73            let known = self.tokenizer.token_to_id(text).is_some();
74            if !known {
75                bail!(
76                    "the chat template context sets {name} = {text:?}, but this tokenizer has no \
77                     such token; the rendered prompt would not match the model"
78                );
79            }
80        }
81        Ok(())
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn a_missing_file_fails_with_a_clear_message() {
91        let error = LocalTokenizer::from_file(Path::new("/nonexistent/tokenizer.json"))
92            .err()
93            .expect("a missing file must not load");
94        assert!(
95            format!("{error:#}").contains("/nonexistent/tokenizer.json"),
96            "{error:#}"
97        );
98    }
99
100    #[test]
101    fn garbage_bytes_are_rejected() {
102        let error = LocalTokenizer::from_bytes(b"not a tokenizer")
103            .err()
104            .expect("garbage must not parse");
105        assert!(
106            format!("{error:#}").contains("parsing the tokenizer failed"),
107            "{error:#}"
108        );
109    }
110}