hermes_tokenizer/
lib.rs

1//! Stable-Rust, byte-level BPE tokenization for Hermes.
2//!
3//! The optimized merge engine and pretokenizers are derived from GigaToken
4//! 0.9.0. The public wrapper is intentionally narrower: local or in-memory
5//! Hugging Face `tokenizer.json` artifacts, persistent per-thread caches,
6//! deterministic batch order, and the vocabulary operations required by
7//! Hermes training, inference, and visualization.
8
9mod bpe;
10mod hf;
11mod pretokenize;
12mod token;
13
14use std::collections::{HashMap, HashSet};
15use std::path::Path;
16use std::sync::{Arc, OnceLock};
17
18use anyhow::{Result, ensure};
19use parking_lot::Mutex;
20use rayon::prelude::*;
21
22use crate::bpe::Tokenizer as BpeEngine;
23use crate::token::TokenId;
24
25/// GigaToken revision from which the optimized byte-level core was extracted.
26pub const UPSTREAM_GIGATOKEN_REVISION: &str = "542367a3efed134883fb4f1140b49c04e6fad3a3";
27
28/// Thread-safe tokenizer with persistent single-document and Rayon-worker
29/// caches. Token IDs remain those of the source `tokenizer.json`.
30pub struct Tokenizer {
31    primary: Mutex<BpeEngine>,
32    workers: OnceLock<Vec<Mutex<BpeEngine>>>,
33    pieces: Arc<[Option<String>]>,
34    token_to_id: Arc<HashMap<String, u32>>,
35    special_ids: Arc<HashSet<u32>>,
36    prefix_ids: Arc<[u32]>,
37    suffix_ids: Arc<[u32]>,
38}
39
40impl Tokenizer {
41    /// Load a byte-level BPE tokenizer from a Hugging Face `tokenizer.json`.
42    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
43        Self::from_loaded(hf::load_file(path)?)
44    }
45
46    /// Load a byte-level BPE tokenizer from in-memory `tokenizer.json` bytes.
47    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
48        Self::from_loaded(hf::load_slice(bytes)?)
49    }
50
51    fn from_loaded(loaded: hf::LoadedTokenizer) -> Result<Self> {
52        ensure!(!loaded.pieces.is_empty(), "tokenizer vocabulary is empty");
53        Ok(Self {
54            primary: Mutex::new(loaded.engine),
55            workers: OnceLock::new(),
56            pieces: loaded.pieces.into(),
57            token_to_id: Arc::new(loaded.token_to_id),
58            special_ids: Arc::new(loaded.special_ids),
59            prefix_ids: loaded.prefix_ids.into(),
60            suffix_ids: loaded.suffix_ids.into(),
61        })
62    }
63
64    /// Encode one UTF-8 string, optionally applying the tokenizer's supported
65    /// single-sequence post-processor.
66    pub fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>> {
67        let mut engine = self.primary.lock();
68        Ok(encode_with(
69            &mut engine,
70            text,
71            add_special_tokens,
72            &self.prefix_ids,
73            &self.suffix_ids,
74        ))
75    }
76
77    /// Encode documents in input order. Small batches stay on the caller;
78    /// larger batches use persistent per-Rayon-thread tokenizer forks.
79    pub fn encode_batch(
80        &self,
81        texts: Vec<String>,
82        add_special_tokens: bool,
83    ) -> Result<Vec<Vec<u32>>> {
84        if texts.is_empty() {
85            return Ok(Vec::new());
86        }
87        let total_bytes = texts.iter().map(String::len).sum::<usize>();
88        if rayon::current_num_threads() == 1 || total_bytes < 1 << 20 {
89            let mut engine = self.primary.lock();
90            return Ok(texts
91                .iter()
92                .map(|text| {
93                    encode_with(
94                        &mut engine,
95                        text,
96                        add_special_tokens,
97                        &self.prefix_ids,
98                        &self.suffix_ids,
99                    )
100                })
101                .collect());
102        }
103
104        let workers = self.worker_engines();
105        let task_count = workers.len() * 4;
106        let chunk_size = texts.len().div_ceil(task_count).max(1);
107        let chunks: Vec<Vec<Vec<u32>>> = texts
108            .par_chunks(chunk_size)
109            .map(|chunk| {
110                let worker = rayon::current_thread_index().unwrap_or(0) % workers.len();
111                let mut engine = workers[worker].lock();
112                chunk
113                    .iter()
114                    .map(|text| {
115                        encode_with(
116                            &mut engine,
117                            text,
118                            add_special_tokens,
119                            &self.prefix_ids,
120                            &self.suffix_ids,
121                        )
122                    })
123                    .collect()
124            })
125            .collect();
126        Ok(chunks.into_iter().flatten().collect())
127    }
128
129    fn worker_engines(&self) -> &[Mutex<BpeEngine>] {
130        self.workers.get_or_init(|| {
131            let engine = self.primary.lock();
132            (0..rayon::current_num_threads().max(1))
133                .map(|_| Mutex::new(engine.fork()))
134                .collect()
135        })
136    }
137
138    /// Decode IDs through the byte-level vocabulary. Invalid UTF-8 fragments
139    /// use the replacement character, matching the display behavior of the
140    /// previous tokenizer backend.
141    pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result<String> {
142        let mut tokens = Vec::with_capacity(ids.len());
143        for &id in ids {
144            if skip_special_tokens && self.special_ids.contains(&id) {
145                continue;
146            }
147            ensure!(
148                self.pieces.get(id as usize).is_some_and(Option::is_some),
149                "tokenizer has no vocabulary entry for token ID {id}"
150            );
151            tokens.push(TokenId::from(id));
152        }
153        let engine = self.primary.lock();
154        let bytes: Vec<u8> = engine.decode(&tokens).collect();
155        Ok(String::from_utf8_lossy(&bytes).into_owned())
156    }
157
158    /// Return the exact vocabulary spelling stored in `tokenizer.json`.
159    pub fn id_to_token(&self, id: u32) -> Option<String> {
160        self.pieces.get(id as usize)?.clone()
161    }
162
163    /// Resolve an exact model or added-token spelling to its ID.
164    pub fn token_to_id(&self, token: &str) -> Option<u32> {
165        self.token_to_id.get(token).copied()
166    }
167
168    /// Vocabulary extent including added tokens.
169    pub fn vocab_size(&self) -> usize {
170        self.pieces.len()
171    }
172}
173
174fn encode_with(
175    engine: &mut BpeEngine,
176    text: &str,
177    add_special_tokens: bool,
178    prefix_ids: &[u32],
179    suffix_ids: &[u32],
180) -> Vec<u32> {
181    let mut ids = Vec::new();
182    engine.encode_with_added_tokens_flat(text.as_bytes(), &mut ids);
183    if !add_special_tokens || (prefix_ids.is_empty() && suffix_ids.is_empty()) {
184        return ids;
185    }
186    let mut wrapped = Vec::with_capacity(prefix_ids.len() + ids.len() + suffix_ids.len());
187    wrapped.extend_from_slice(prefix_ids);
188    wrapped.extend(ids);
189    wrapped.extend_from_slice(suffix_ids);
190    wrapped
191}
hermes_tokenizer/lib.rs

hermes_tokenizer/
lib.rs
