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
//! Donor-side token→text display support. In pipeline sharding every donor already RECEIVES the
//! token ids it processes (`BatchCol.token`) and, on the last stage, the ids it GENERATES — but a
//! donor has no tokenizer (the browser build drops it; `onig_sys` can't target wasm). Decoding
//! ids→text, however, needs neither the tokenizer nor `onig`: it only needs a compact `id→bytes`
//! table plus concatenation. The coordinator (which HAS the tokenizer) builds that table once from
//! `tokenizer.json` ([`build_vocab_blob`]) and ships it; each donor decodes with [`ByteVocab`].
//!
//! This is DISPLAY-ONLY and off the hot path: donors buffer ids per step (O(1)) and decode a short
//! rolling window only on their throttled stat tick, so it never slows inference.
use std::collections::HashMap;
/// GPT-2 / Qwen byte-level alphabet — the reversible byte↔unicode map HF's `ByteLevel` uses.
/// Returns the DECODE direction (`char → byte`): each byte-level vocab piece is a string over these
/// chars, so mapping every char back to its byte and UTF-8-decoding the concatenation reconstructs
/// the original text (including multi-byte code points that straddle token boundaries).
fn byte_decoder() -> HashMap<char, u8> {
// Printable ranges kept as their own code point; everything else is offset by 256.
let mut bs: Vec<u32> = Vec::new();
bs.extend((b'!' as u32)..=(b'~' as u32));
bs.extend(0xA1u32..=0xAC);
bs.extend(0xAEu32..=0xFF);
let mut cs = bs.clone();
let mut n = 0u32;
for b in 0u32..256 {
if !bs.contains(&b) {
bs.push(b);
cs.push(256 + n);
n += 1;
}
}
bs.iter()
.zip(cs.iter())
.filter_map(|(&b, &c)| char::from_u32(c).map(|ch| (ch, b as u8)))
.collect()
}
/// Build the compact `id→bytes` vocab blob a coordinator ships to donors, from a `tokenizer.json`'s
/// `model.vocab` (`{piece: id}`). Byte-level pieces are remapped to their raw bytes so a donor
/// decodes purely by concatenation — no tokenizer, no `onig`, wasm-friendly. Special/control tokens
/// (from `added_tokens`) render as empty so generated text reads cleanly. Returns `None` if the
/// JSON carries no recognizable byte-level vocab.
///
/// Blob format: `u32 count` then, for each id `0..count`, `u32 len` + `len` bytes.
pub fn build_vocab_blob(tokenizer_json: &[u8]) -> Option<Vec<u8>> {
let v: serde_json::Value = serde_json::from_slice(tokenizer_json).ok()?;
let vocab = v.get("model")?.get("vocab")?.as_object()?;
let dec = byte_decoder();
let max_id = vocab.values().filter_map(serde_json::Value::as_u64).max()? as usize;
let mut id_to_bytes: Vec<Vec<u8>> = vec![Vec::new(); max_id + 1];
for (piece, id) in vocab {
if let Some(id) = id.as_u64() {
id_to_bytes[id as usize] = piece.chars().filter_map(|c| dec.get(&c).copied()).collect();
}
}
// Control tokens (e.g. `<|im_start|>`, `<|im_end|>`) are stored verbatim, not byte-level.
// Render them as empty so a donor's live text isn't cluttered by scaffolding.
if let Some(added) = v.get("added_tokens").and_then(|a| a.as_array()) {
for t in added {
let Some(id) = t.get("id").and_then(serde_json::Value::as_u64) else {
continue;
};
let Some(slot) = id_to_bytes.get_mut(id as usize) else {
continue;
};
let special = t
.get("special")
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
*slot = if special {
Vec::new()
} else {
t.get("content")
.and_then(serde_json::Value::as_str)
.map(|s| s.as_bytes().to_vec())
.unwrap_or_default()
};
}
}
let mut blob = Vec::with_capacity(4 + id_to_bytes.iter().map(|b| 4 + b.len()).sum::<usize>());
blob.extend_from_slice(&(id_to_bytes.len() as u32).to_le_bytes());
for b in &id_to_bytes {
blob.extend_from_slice(&(b.len() as u32).to_le_bytes());
blob.extend_from_slice(b);
}
Some(blob)
}
/// Donor-side token→text decoder: `id → raw bytes` (from a shipped [`build_vocab_blob`]), decoding
/// by pure concatenation + UTF-8. No tokenizer, no `onig` — compiles to wasm. Display-only.
pub struct ByteVocab {
pieces: Vec<Vec<u8>>,
}
impl ByteVocab {
/// Parse a shipped vocab blob. Returns `None` on a malformed/truncated blob.
pub fn from_blob(blob: &[u8]) -> Option<Self> {
let count = u32::from_le_bytes(blob.get(0..4)?.try_into().ok()?) as usize;
let mut pieces = Vec::with_capacity(count);
let mut o = 4;
for _ in 0..count {
let len = u32::from_le_bytes(blob.get(o..o + 4)?.try_into().ok()?) as usize;
o += 4;
pieces.push(blob.get(o..o + len)?.to_vec());
o += len;
}
Some(Self { pieces })
}
/// Decode a token-id sequence to text (unknown ids contribute nothing).
pub fn decode(&self, ids: &[u32]) -> String {
let mut bytes = Vec::new();
for &id in ids {
if let Some(p) = self.pieces.get(id as usize) {
bytes.extend_from_slice(p);
}
}
String::from_utf8_lossy(&bytes).into_owned()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn should_round_trip_byte_level_pieces_through_ship_and_decode() {
// A tiny byte-level vocab: 'Ġ' (U+0120) is the byte-level char for a space (0x20), so
// "ĠParis" must decode to " Paris"; a control token renders as empty.
let tok = r#"{
"model": { "vocab": { "ĠParis": 0, "Hello": 1, "Ġworld": 2 } },
"added_tokens": [ { "id": 3, "content": "<|end|>", "special": true } ]
}"#;
let blob = build_vocab_blob(tok.as_bytes()).expect("blob");
let vocab = ByteVocab::from_blob(&blob).expect("parse");
assert_eq!(vocab.decode(&[1, 0]), "Hello Paris");
assert_eq!(vocab.decode(&[1, 2]), "Hello world");
assert_eq!(vocab.decode(&[1, 3, 0]), "Hello Paris"); // control token drops out
}
#[test]
fn should_reconstruct_multibyte_utf8_split_across_tokens() {
// 'é' is UTF-8 0xC3 0xA9 → byte-level chars 'Ã'(0xC3) '©'(0xA9). Split across two tokens,
// concatenation-then-UTF-8 must still yield "é".
let tok = r#"{ "model": { "vocab": { "Ã": 0, "©": 1 } } }"#;
let blob = build_vocab_blob(tok.as_bytes()).expect("blob");
let vocab = ByteVocab::from_blob(&blob).expect("parse");
assert_eq!(vocab.decode(&[0, 1]), "é");
}
}