laya_core/
tokenizer.rs

1//! Checkpoint tokenizer.
2//!
3//! The Python package wraps the very same Rust `tokenizers` crate through
4//! PyO3. Here it is used directly, so encoding no longer crosses an FFI
5//! boundary or allocates a Python dictionary per call.
6
7use crate::error::{Error, Result};
8use serde_json::Value;
9use std::path::Path;
10use tokenizers::Tokenizer as Backend;
11
12/// A loaded tokenizer plus the four special tokens upstream requires.
13pub struct Tokenizer {
14    backend: Backend,
15    pub cls_token: String,
16    pub sep_token: String,
17    pub pad_token: String,
18    pub mask_token: String,
19    pub cls_token_id: u32,
20    pub sep_token_id: u32,
21    pub pad_token_id: u32,
22    pub mask_token_id: u32,
23}
24
25impl std::fmt::Debug for Tokenizer {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        f.debug_struct("Tokenizer")
28            .field("mask_token", &self.mask_token)
29            .field("vocab_size", &self.backend.get_vocab_size(true))
30            .finish_non_exhaustive()
31    }
32}
33
34fn special_token(config: &Value, name: &'static str) -> Option<String> {
35    match config.get(name)? {
36        Value::String(text) => Some(text.clone()),
37        // `tokenizer_config.json` sometimes stores an AddedToken object.
38        Value::Object(map) => map.get("content")?.as_str().map(str::to_string),
39        _ => None,
40    }
41}
42
43impl Tokenizer {
44    /// Load `tokenizer.json` and `tokenizer_config.json` from a directory.
45    pub fn from_dir(dir: impl AsRef<Path>) -> Result<Self> {
46        let dir = dir.as_ref();
47        let tokenizer_path = dir.join("tokenizer.json");
48        let mut backend =
49            Backend::from_file(&tokenizer_path).map_err(|error| Error::Tokenizer {
50                path: tokenizer_path,
51                message: error.to_string(),
52            })?;
53        // Upstream disables both explicitly; padding and truncation are handled
54        // by the collator so that marker positions stay valid.
55        backend.with_padding(None);
56        backend
57            .with_truncation(None)
58            .map_err(|error| Error::Tokenizer {
59                path: dir.join("tokenizer.json"),
60                message: error.to_string(),
61            })?;
62
63        let config_path = dir.join("tokenizer_config.json");
64        let config: Value = serde_json::from_str(&std::fs::read_to_string(&config_path)?)?;
65
66        let mut resolved = Vec::with_capacity(4);
67        for name in ["cls_token", "sep_token", "pad_token", "mask_token"] {
68            let token = special_token(&config, name).ok_or(Error::MissingSpecialToken(name))?;
69            let id = backend
70                .token_to_id(&token)
71                .ok_or(Error::MissingSpecialToken(name))?;
72            resolved.push((token, id));
73        }
74        let [cls, sep, pad, mask]: [(String, u32); 4] =
75            resolved.try_into().expect("exactly four special tokens");
76
77        Ok(Self {
78            backend,
79            cls_token: cls.0,
80            sep_token: sep.0,
81            pad_token: pad.0,
82            mask_token: mask.0,
83            cls_token_id: cls.1,
84            sep_token_id: sep.1,
85            pad_token_id: pad.1,
86            mask_token_id: mask.1,
87        })
88    }
89
90    /// Encode without special tokens, as every upstream call site does.
91    pub fn encode(&self, text: &str) -> Result<Vec<u32>> {
92        self.backend
93            .encode_fast(text, false)
94            .map(|encoding| encoding.get_ids().to_vec())
95            .map_err(|error| Error::Tokenizer {
96                path: Default::default(),
97                message: error.to_string(),
98            })
99    }
100
101    pub fn vocab_size(&self) -> usize {
102        self.backend.get_vocab_size(true)
103    }
104}
laya_core/tokenizer.rs

laya_core/
tokenizer.rs
