1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
//! Byte-level BPE tokenizer for GPT-2 style models.
//!
//! Supports encoding text into token IDs and decoding back, using a vocabulary
//! of byte tokens (0–255) plus merge-derived tokens. This is the foundation for
//! real text-in/text-out transformer inference.
use std::collections::HashMap;
/// A byte-level BPE tokenizer.
pub struct BpeTokenizer {
/// token_id → bytes
vocab: Vec<Vec<u8>>,
/// (first_token_id, second_token_id) → merged_token_id, in priority order
merges: Vec<(u32, u32, u32)>,
/// Fast lookup for merge rules: (a, b) → merged_id
merge_map: HashMap<(u32, u32), u32>,
/// single-byte → token_id (for initial encoding)
byte_to_id: HashMap<u8, u32>,
/// Pre-computed token bytes for every vocab entry. Avoids rebuilding this
/// Vec (50K small allocations) on every constrained sampling step.
token_bytes_cache: Vec<Option<Vec<u8>>>,
}
impl BpeTokenizer {
/// Create a tokenizer from an explicit vocab and ordered merge list.
///
/// `vocab[i]` is the byte sequence for token ID `i`.
/// `merges` is a list of `(first_id, second_id, merged_id)` in priority order.
pub fn new(vocab: Vec<Vec<u8>>, merges: Vec<(u32, u32, u32)>) -> Self {
let merge_map = merges.iter().map(|(a, b, m)| ((*a, *b), *m)).collect();
let mut byte_to_id = HashMap::new();
let mut token_bytes_cache = Vec::with_capacity(vocab.len());
for (id, seq) in vocab.iter().enumerate() {
if seq.len() == 1 {
byte_to_id.insert(seq[0], id as u32);
}
token_bytes_cache.push(Some(seq.clone()));
}
Self {
vocab,
merges,
merge_map,
byte_to_id,
token_bytes_cache,
}
}
/// A minimal tokenizer with only the 256 byte tokens (no merges).
/// Useful as a fallback or for testing.
pub fn byte_fallback() -> Self {
let vocab: Vec<Vec<u8>> = (0..=255u8).map(|b| vec![b]).collect();
Self::new(vocab, Vec::new())
}
/// Encode `text` into token IDs using BPE.
pub fn encode(&self, text: &str) -> Vec<u32> {
let bytes = text.as_bytes();
if bytes.is_empty() {
return Vec::new();
}
// Start with each byte mapped to its vocab token ID.
let mut tokens: Vec<u32> = bytes
.iter()
.map(|b| self.byte_to_id.get(b).copied().unwrap_or(*b as u32))
.collect();
// Greedily apply the highest-priority merge that exists in the current sequence.
// Repeat until no more merges are possible.
loop {
let mut best_idx = None;
let mut best_rank: Option<usize> = None;
for (i, window) in tokens.windows(2).enumerate() {
let key = (window[0], window[1]);
if let Some(&merged) = self.merge_map.get(&key) {
// Find the rank (priority) of this merge.
let rank = self
.merges
.iter()
.position(|(a, b, _)| *a == key.0 && *b == key.1);
if let Some(r) = rank
&& best_rank.is_none_or(|br| r < br)
{
best_rank = Some(r);
best_idx = Some((i, merged));
}
}
}
if let Some((idx, merged)) = best_idx {
tokens.splice(idx..=idx + 1, [merged]);
} else {
break;
}
}
tokens
}
/// Decode token IDs back into a string.
///
/// Invalid token IDs are skipped (they have no entry in the vocab).
pub fn decode(&self, tokens: &[u32]) -> String {
let mut bytes = Vec::new();
for &id in tokens {
if let Some(seq) = self.vocab.get(id as usize) {
bytes.extend_from_slice(seq);
}
}
String::from_utf8_lossy(&bytes).into_owned()
}
/// Raw byte sequence for a token ID, or `None` if the ID is out of range.
/// Useful for emitting per-token `bytes` fields (e.g. OpenAI logprobs payloads).
pub fn token_bytes(&self, id: u32) -> Option<&[u8]> {
self.vocab.get(id as usize).map(Vec::as_slice)
}
/// Pre-computed `Option<Vec<u8>>` for every vocab entry.
/// Eliminates rebuilding this on every constrained sampling step.
pub fn cached_token_bytes(&self) -> &[Option<Vec<u8>>] {
&self.token_bytes_cache
}
/// Vocabulary size (number of distinct tokens).
pub fn vocab_size(&self) -> usize {
self.vocab.len()
}
/// Build a `BpeTokenizer` from GGUF tokenizer metadata.
///
/// Each vocab entry is converted to bytes via UTF-8. Merge rules are parsed
/// as "first second" strings; their IDs are looked up in the vocab. The
/// merged token is the concatenation of the two byte sequences, and its ID
/// is found in the vocab (standard BPE vocab contains all merged tokens).
pub fn from_gguf(meta: &crate::parsers::gguf::GgufTokenizerMetadata) -> anyhow::Result<Self> {
use anyhow::Context;
if meta.vocab.is_empty() {
anyhow::bail!("GGUF tokenizer metadata has no vocab");
}
let vocab: Vec<Vec<u8>> = meta.vocab.iter().map(|s| s.as_bytes().to_vec()).collect();
// Build token string → ID map for merge lookups.
let mut token_to_id: std::collections::HashMap<Vec<u8>, u32> =
std::collections::HashMap::new();
for (id, seq) in vocab.iter().enumerate() {
token_to_id.insert(seq.clone(), id as u32);
}
let mut merges = Vec::new();
for rule in &meta.merges {
let mut parts = rule.splitn(2, ' ');
let first_str = parts.next().context("empty merge rule")?;
let second_str = parts.next().context("merge rule missing second token")?;
let first_bytes = first_str.as_bytes().to_vec();
let second_bytes = second_str.as_bytes().to_vec();
let first_id = *token_to_id
.get(&first_bytes)
.with_context(|| format!("merge rule first token not in vocab: {first_str:?}"))?;
let second_id = *token_to_id
.get(&second_bytes)
.with_context(|| format!("merge rule second token not in vocab: {second_str:?}"))?;
let merged_bytes: Vec<u8> = first_bytes.iter().chain(&second_bytes).copied().collect();
let merged_id = *token_to_id.get(&merged_bytes).with_context(|| {
format!("merged token not in vocab: {first_str:?}+{second_str:?}")
})?;
merges.push((first_id, second_id, merged_id));
}
Ok(Self::new(vocab, merges))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn byte_fallback_roundtrip() {
let tok = BpeTokenizer::byte_fallback();
let text = "hello world 123 !@#";
let ids = tok.encode(text);
assert!(!ids.is_empty());
let decoded = tok.decode(&ids);
assert_eq!(decoded, text);
}
#[test]
fn empty_string_encodes_to_empty() {
let tok = BpeTokenizer::byte_fallback();
assert!(tok.encode("").is_empty());
assert_eq!(tok.decode(&[]), "");
}
#[test]
fn bpe_merge_combines_tokens() {
// Vocab: 0='a', 1='b', 2='c', 3="ab" (merged)
let vocab = vec![vec![b'a'], vec![b'b'], vec![b'c'], vec![b'a', b'b']];
// Merge 'a' (0) + 'b' (1) → "ab" (3)
let merges = vec![(0, 1, 3)];
let tok = BpeTokenizer::new(vocab, merges);
let ids = tok.encode("ab");
assert_eq!(ids, vec![3]);
let ids2 = tok.encode("abc");
assert_eq!(ids2, vec![3, 2]);
}
#[test]
fn bpe_multiple_merges_priority() {
// Vocab: 0='a', 1='b', 2='c', 3="ab", 4="bc", 5="abc"
let vocab = vec![
vec![b'a'],
vec![b'b'],
vec![b'c'],
vec![b'a', b'b'],
vec![b'b', b'c'],
vec![b'a', b'b', b'c'],
];
// Priority: "ab" first, then "bc", then "abc"
let merges = vec![(0, 1, 3), (1, 2, 4), (3, 2, 5)];
let tok = BpeTokenizer::new(vocab, merges);
let ids = tok.encode("abc");
// "ab" merges first (priority 0) → [3, 2], then "abc" → 5 merges again.
assert_eq!(ids, vec![5]);
}
#[test]
fn invalid_token_id_skipped() {
let tok = BpeTokenizer::byte_fallback();
let decoded = tok.decode(&[0, 9999, 1]); // 9999 is out of range
assert_eq!(decoded, "\x00\x01");
}
#[test]
fn from_gguf_builds_tokenizer() {
let meta = crate::parsers::gguf::GgufTokenizerMetadata {
model: "gpt2".to_string(),
vocab: vec!["a".to_string(), "b".to_string(), "ab".to_string()],
merges: vec!["a b".to_string()],
bos_token_id: None,
eos_token_id: None,
chat_template: None,
};
let tok = BpeTokenizer::from_gguf(&meta).expect("from_gguf should succeed");
assert_eq!(tok.vocab_size(), 3);
let ids = tok.encode("ab");
assert_eq!(ids, vec![2]); // "ab" merged to token 2
}
#[test]
fn from_gguf_fails_on_missing_vocab() {
let meta = crate::parsers::gguf::GgufTokenizerMetadata::default();
assert!(BpeTokenizer::from_gguf(&meta).is_err());
}
}