metrale_grammar/tokenizer/
vocab_type.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// VocabType — port of the `enum class VocabType` in
4// `include/xgrammar/tokenizer_info.h`.
5//
6// The vocabulary type controls how raw (encoded) vocabulary strings are
7// turned into byte sequences. `ByteLevel` is the F68-critical type:
8// ByteLevel-BPE tokenizers (Qwen, MiniMax, Mistral 4) MUST be detected
9// as `ByteLevel` or constrained decoding silently corrupts.
10
11use std::fmt;
12
13/// How a tokenizer encodes its vocabulary strings.
14///
15/// Discriminants match the C++ `enum class VocabType : int` exactly so
16/// that serialized metadata (`vocab_type` integer) round-trips.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum VocabType {
19    /// The vocabulary string is the literal token bytes (no decoding).
20    Raw = 0,
21    /// SentencePiece-style: `<0xHH>` byte tokens and `▁` space markers.
22    ByteFallback = 1,
23    /// GPT-2 byte-level BPE: each byte is mapped to a printable unicode
24    /// codepoint via the bytes-to-unicode table.
25    ByteLevel = 2,
26}
27
28impl VocabType {
29    /// Convert from the integer discriminant used in serialized metadata.
30    ///
31    /// Returns `None` for any value other than `0`, `1`, `2` — faithful
32    /// to the C++ `XGRAMMAR_CHECK(vocab_type_int == 0 || 1 || 2)`.
33    pub fn from_int(value: i64) -> Option<Self> {
34        match value {
35            0 => Some(VocabType::Raw),
36            1 => Some(VocabType::ByteFallback),
37            2 => Some(VocabType::ByteLevel),
38            _ => None,
39        }
40    }
41
42    /// The integer discriminant used in serialized metadata.
43    pub fn as_int(self) -> i64 {
44        self as i64
45    }
46}
47
48impl Default for VocabType {
49    /// The C++ constructor defaults `vocab_type` to `VocabType::RAW`.
50    fn default() -> Self {
51        VocabType::Raw
52    }
53}
54
55impl fmt::Display for VocabType {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        let s = match self {
58            VocabType::Raw => "RAW",
59            VocabType::ByteFallback => "BYTE_FALLBACK",
60            VocabType::ByteLevel => "BYTE_LEVEL",
61        };
62        f.write_str(s)
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn discriminants_match_cpp() {
72        assert_eq!(VocabType::Raw as i64, 0);
73        assert_eq!(VocabType::ByteFallback as i64, 1);
74        assert_eq!(VocabType::ByteLevel as i64, 2);
75    }
76
77    #[test]
78    fn from_int_round_trip() {
79        for v in [
80            VocabType::Raw,
81            VocabType::ByteFallback,
82            VocabType::ByteLevel,
83        ] {
84            assert_eq!(VocabType::from_int(v.as_int()), Some(v));
85        }
86    }
87
88    #[test]
89    fn from_int_rejects_invalid() {
90        assert_eq!(VocabType::from_int(-1), None);
91        assert_eq!(VocabType::from_int(3), None);
92        assert_eq!(VocabType::from_int(999), None);
93    }
94
95    #[test]
96    fn default_is_raw() {
97        assert_eq!(VocabType::default(), VocabType::Raw);
98    }
99
100    #[test]
101    fn display_strings() {
102        assert_eq!(VocabType::Raw.to_string(), "RAW");
103        assert_eq!(VocabType::ByteFallback.to_string(), "BYTE_FALLBACK");
104        assert_eq!(VocabType::ByteLevel.to_string(), "BYTE_LEVEL");
105    }
106}
metrale_grammar/tokenizer/vocab_type.rs

metrale_grammar/tokenizer/
vocab_type.rs
