metrale_grammar/tokenizer/
decoder.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// Token decoders — port of `class TokenDecoder` in
4// `cpp/tokenizer_info.cc`.
5//
6// Turns a raw (encoded) vocabulary string into the actual token byte
7// sequence, according to the tokenizer's `VocabType`:
8//
9//   * Raw          — the string is the literal token bytes.
10//   * ByteFallback — `<0xHH>` tokens decode to a single byte; the
11//                    `▁` (U+2581) marker decodes back to a space.
12//   * ByteLevel    — the GPT-2 bytes-to-unicode transform is inverted.
13//
14// The C++ returns `std::string` (raw bytes). We return `Vec<u8>` since
15// a decoded token need not be valid UTF-8.
16
17use crate::support::encoding::{char_handling_error, parse_utf8};
18use crate::tokenizer::vocab_type::VocabType;
19
20/// Decode one raw vocabulary token to its byte sequence for `vocab_type`.
21pub fn decode_token(token: &str, vocab_type: VocabType) -> Vec<u8> {
22    match vocab_type {
23        VocabType::Raw => token.as_bytes().to_vec(),
24        VocabType::ByteFallback => space_replacer_decoder(&byte_fallback_decoder(token)),
25        VocabType::ByteLevel => byte_level_decoder(token),
26    }
27}
28
29/// ByteFallback decoder: transform a token like `<0x1B>` into the single
30/// raw byte `0x1B`. Any other token is returned unchanged.
31///
32/// Faithful to C++ `ByteFallbackDecoder`: the match requires length 6,
33/// a `<0x` prefix, a `>` suffix, and two uppercase-or-digit hex chars.
34fn byte_fallback_decoder(token: &str) -> Vec<u8> {
35    let bytes = token.as_bytes();
36    if bytes.len() == 6 && &bytes[0..3] == b"<0x" && bytes[5] == b'>' {
37        let mut byte: i32 = 0;
38        for &c in &bytes[3..5] {
39            // C++ only handles `0-9` and `A-F` here (uppercase hex).
40            let digit = if c.is_ascii_digit() {
41                (c - b'0') as i32
42            } else {
43                (c - b'A') as i32 + 10
44            };
45            byte = byte * 16 + digit;
46        }
47        // C++ `XGRAMMAR_CHECK(byte >= 0 && byte < 256)`. With two hex
48        // chars this always holds, so we keep the value directly.
49        return vec![byte as u8];
50    }
51    token.as_bytes().to_vec()
52}
53
54/// SpaceReplacer decoder: transform the UTF-8 of `▁` (U+2581, "lower one
55/// eighth block", bytes `E2 96 81`) back into an ASCII space.
56///
57/// Faithful to C++ `SpaceReplacerDecoder`, which scans bytes directly.
58fn space_replacer_decoder(token: &[u8]) -> Vec<u8> {
59    let mut result = Vec::with_capacity(token.len());
60    let mut i = 0;
61    while i < token.len() {
62        // C++ guard is `i + 2 < size`, i.e. all three bytes addressable.
63        if i + 2 < token.len() && token[i] == 0xE2 && token[i + 1] == 0x96 && token[i + 2] == 0x81 {
64            result.push(b' ');
65            i += 3;
66        } else {
67            result.push(token[i]);
68            i += 1;
69        }
70    }
71    result
72}
73
74/// The inverse of the GPT-2 bytes-to-unicode map.
75///
76/// `char_to_byte_map()[codepoint]` is the original byte, or `-1` when
77/// no byte maps to that codepoint. The table has 324 entries, matching
78/// `std::array<int, 324>` in the C++ source verbatim.
79pub fn char_to_byte_map() -> &'static [i16; 324] {
80    &CHAR_TO_BYTE_MAP
81}
82
83/// The GPT-2 bytes-to-unicode forward map: byte -> unicode codepoint.
84///
85/// This is the table HuggingFace uses to *encode* the vocabulary; it is
86/// the exact inverse of `CHAR_TO_BYTE_MAP`. Derived (not duplicated)
87/// from the inverse table to keep a single source of truth.
88pub fn byte_to_char_map() -> &'static [u32; 256] {
89    &BYTE_TO_CHAR_MAP
90}
91
92/// GPT-2 byte-decoder inverse table.
93///
94/// STATIC, DELIBERATELY — compile-time data. A fixed table defined by the
95/// GPT-2 byte-level BPE specification, identical for every tokenizer that
96/// uses that scheme; it is not derived from any model's vocabulary and has
97/// no interior mutability. `static` rather than `const` so the lookups
98/// borrow one copy instead of inlining 324 halfwords at each call site.
99#[rustfmt::skip]
100const CHAR_TO_BYTE_MAP: [i16; 324] = [
101    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
102    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
103    46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68,
104    69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91,
105    92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
106    112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, -1, -1, -1, -1,
107    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
108    -1, -1, -1, -1, -1, -1, -1, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, -1,
109    174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191,
110    192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209,
111    210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227,
112    228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245,
113    246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
114    13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128,
115    129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146,
116    147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 173,
117];
118
119/// Forward map derived once from `CHAR_TO_BYTE_MAP` (SSOT: the inverse
120/// table is authoritative; this is `byte -> codepoint` such that
121/// `CHAR_TO_BYTE_MAP[byte_to_char(byte)] == byte`).
122const BYTE_TO_CHAR_MAP: [u32; 256] = build_byte_to_char_map();
123
124const fn build_byte_to_char_map() -> [u32; 256] {
125    let mut out = [0u32; 256];
126    let mut cp = 0usize;
127    while cp < CHAR_TO_BYTE_MAP.len() {
128        let b = CHAR_TO_BYTE_MAP[cp];
129        if b >= 0 {
130            out[b as usize] = cp as u32;
131        }
132        cp += 1;
133    }
134    out
135}
136
137/// ByteLevel decoder: invert the GPT-2 bytes-to-unicode transform.
138///
139/// Faithful to C++ `ByteLevelDecoder`. If the token is not valid UTF-8,
140/// or contains a codepoint with no inverse mapping, the original token
141/// bytes are returned unchanged.
142fn byte_level_decoder(token: &str) -> Vec<u8> {
143    let codepoints = parse_utf8(token.as_bytes(), false);
144    if codepoints.len() == 1 && codepoints[0] == char_handling_error::INVALID_UTF8 {
145        return token.as_bytes().to_vec();
146    }
147
148    let map = char_to_byte_map();
149    let mut decoded = Vec::with_capacity(codepoints.len());
150    for cp in codepoints {
151        // C++ `XGRAMMAR_CHECK(unicode_codepoint >= 0)` — parse_utf8 with
152        // preserve=false never yields negatives other than the invalid
153        // sentinel handled above, so any negative here means no mapping.
154        if cp < 0 || cp as usize >= map.len() || map[cp as usize] == -1 {
155            return token.as_bytes().to_vec();
156        }
157        decoded.push(map[cp as usize] as u8);
158    }
159    decoded
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn raw_is_literal() {
168        assert_eq!(decode_token("hello", VocabType::Raw), b"hello");
169        assert_eq!(decode_token("", VocabType::Raw), b"");
170    }
171
172    #[test]
173    fn byte_fallback_hex_token() {
174        // <0x1B> -> ESC byte 0x1B
175        assert_eq!(decode_token("<0x1B>", VocabType::ByteFallback), vec![0x1B]);
176        // <0x00> -> NUL
177        assert_eq!(decode_token("<0x00>", VocabType::ByteFallback), vec![0x00]);
178        // <0xFF> -> 0xFF
179        assert_eq!(decode_token("<0xFF>", VocabType::ByteFallback), vec![0xFF]);
180    }
181
182    #[test]
183    fn byte_fallback_non_hex_passthrough() {
184        // Not the <0xHH> shape -> unchanged.
185        assert_eq!(decode_token("er", VocabType::ByteFallback), b"er");
186        // Lowercase hex is NOT matched by the C++ uppercase-only path,
187        // but it still has the <0x..> shape; C++ computes c-'A'+10 on
188        // 'b' which is wrong — we faithfully reproduce: keep length 6
189        // shape match. Use a clearly-non-matching token instead.
190        assert_eq!(decode_token("<0x1>", VocabType::ByteFallback), b"<0x1>");
191    }
192
193    #[test]
194    fn byte_fallback_space_marker() {
195        // U+2581 "▁" -> space
196        let token = "\u{2581}hello";
197        assert_eq!(decode_token(token, VocabType::ByteFallback), b" hello");
198    }
199
200    #[test]
201    fn byte_fallback_double_space_marker() {
202        let token = "\u{2581}\u{2581}";
203        assert_eq!(decode_token(token, VocabType::ByteFallback), b"  ");
204    }
205
206    #[test]
207    fn byte_level_round_trip_ascii() {
208        // For printable ASCII, byte_to_char is the identity, so the
209        // encoded string equals the literal — decode returns the bytes.
210        let map = byte_to_char_map();
211        // 'A' (0x41) maps to itself.
212        assert_eq!(map[0x41], 0x41);
213        let mut enc = String::new();
214        for &b in b"automotive" {
215            enc.push(char::from_u32(map[b as usize]).unwrap());
216        }
217        assert_eq!(decode_token(&enc, VocabType::ByteLevel), b"automotive");
218    }
219
220    #[test]
221    fn byte_level_space_is_g_with_dot() {
222        // GPT-2 maps byte 0x20 (space) to U+0120 'Ġ'.
223        let map = byte_to_char_map();
224        assert_eq!(map[0x20], 0x0120);
225        // " automotive" -> "Ġautomotive"
226        let mut enc = String::from("\u{0120}");
227        for &b in b"automotive" {
228            enc.push(char::from_u32(map[b as usize]).unwrap());
229        }
230        assert_eq!(decode_token(&enc, VocabType::ByteLevel), b" automotive");
231    }
232
233    #[test]
234    fn byte_level_high_bytes() {
235        // The UTF-8 of '我' is E6 88 91 — encode each byte via the map.
236        let map = byte_to_char_map();
237        let mut enc = String::new();
238        for &b in "我".as_bytes() {
239            enc.push(char::from_u32(map[b as usize]).unwrap());
240        }
241        assert_eq!(decode_token(&enc, VocabType::ByteLevel), "我".as_bytes());
242    }
243
244    #[test]
245    fn byte_level_unmapped_codepoint_passthrough() {
246        // U+4E2D '中' has no inverse mapping -> token returned as-is.
247        assert_eq!(decode_token("中", VocabType::ByteLevel), "中".as_bytes());
248    }
249
250    #[test]
251    fn inverse_table_is_consistent() {
252        // Every mapped entry of the inverse table must round-trip.
253        let inv = char_to_byte_map();
254        let fwd = byte_to_char_map();
255        for byte in 0u32..256 {
256            let cp = fwd[byte as usize];
257            assert_eq!(inv[cp as usize], byte as i16);
258        }
259        // All 256 bytes have a forward mapping (the map is a bijection
260        // onto 256 distinct codepoints).
261        let mut seen = std::collections::HashSet::new();
262        for byte in 0u32..256 {
263            assert!(seen.insert(fwd[byte as usize]));
264        }
265    }
266
267    #[test]
268    fn byte_level_invalid_utf8_passthrough() {
269        // A literal char that decodes to a non-mapped point falls back.
270        // Codepoint 0 maps to -1 in the inverse table.
271        assert_eq!(char_to_byte_map()[0], -1);
272    }
273}
metrale_grammar/tokenizer/decoder.rs

metrale_grammar/tokenizer/
decoder.rs
