xgrammar 0.5.0

Rust bindings for XGrammar
Documentation 
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

mod common;

use xgrammar::{TokenizerInfo, VocabType};

// Shared test cases (tokenizer_id, expected_vocab_type, expected_add_prefix_space)
const TEST_TOKENIZER_CASES: &[(&str, VocabType, bool)] = &[
    ("luodian/llama-7b-hf", VocabType::ByteFallback, true),
    ("meta-llama/Llama-2-7b-chat-hf", VocabType::ByteFallback, true),
    ("meta-llama/Meta-Llama-3-8B-Instruct", VocabType::ByteLevel, false),
    ("meta-llama/Meta-Llama-3.1-8B-Instruct", VocabType::ByteLevel, false),
    // ("lmsys/vicuna-7b-v1.5", VocabType::ByteFallback, true), // no tokenizer.json
    ("NousResearch/Hermes-2-Theta-Llama-3-70B", VocabType::ByteLevel, false),
    ("NousResearch/Hermes-3-Llama-3.1-8B", VocabType::ByteLevel, false),
    ("google/gemma-2b-it", VocabType::ByteFallback, false),
    ("CohereForAI/aya-23-8B", VocabType::ByteLevel, false),
    ("deepseek-ai/DeepSeek-Coder-V2-Instruct", VocabType::ByteLevel, false),
    ("deepseek-ai/DeepSeek-V2-Chat-0628", VocabType::ByteLevel, false),
    ("deepseek-ai/deepseek-coder-7b-instruct-v1.5", VocabType::ByteLevel, false),
    ("microsoft/phi-2", VocabType::ByteLevel, false),
    ("microsoft/Phi-3-mini-4k-instruct", VocabType::ByteFallback, true),
    ("microsoft/Phi-3.5-mini-instruct", VocabType::ByteFallback, true),
    ("Qwen/Qwen1.5-4B-Chat", VocabType::ByteLevel, false),
    // ("Qwen/Qwen2-7B-Instruct", VocabType::ByteLevel, false), // no tokenizer.json
    // ("microsoft/Phi-3-small-8k-instruct", VocabType::Raw, false), // no tokenizer.json
    // ("Qwen/Qwen-7B-Chat", VocabType::Raw, false), // no tokenizer.json
    ("meta-llama/Llama-3.2-1B", VocabType::ByteLevel, false),
    ("google/gemma-2-2b-it", VocabType::ByteFallback, false),
    ("deepseek-ai/DeepSeek-V2.5", VocabType::ByteLevel, false),
    ("Qwen/Qwen2.5-1.5B", VocabType::ByteLevel, false),
    // ("internlm/internlm2_5-7b-chat", VocabType::ByteFallback, false), // no tokenizer.json
    ("mistralai/Mixtral-8x22B-Instruct-v0.1", VocabType::ByteFallback, true),
    // ("THUDM/glm-4-9b-chat", VocabType::Raw, false), // no tokenizer.json
    // ("THUDM/chatglm3-6b", VocabType::ByteFallback, true), // no tokenizer.json
    ("deepseek-ai/DeepSeek-R1", VocabType::ByteLevel, false),
    ("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", VocabType::ByteLevel, false),
    ("deepseek-ai/DeepSeek-R1-Distill-Llama-8B", VocabType::ByteLevel, false),
    ("LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct", VocabType::ByteLevel, false),
    ("LGAI-EXAONE/EXAONE-4.0-32B-FP8", VocabType::ByteLevel, false),
];

fn assert_metadata(
    tokenizer_info: &TokenizerInfo,
    expected_vocab_type: VocabType,
    expected_add_prefix_space: bool,
) {
    assert_eq!(tokenizer_info.get_vocab_type(), expected_vocab_type);
    assert_eq!(tokenizer_info.get_add_prefix_space(), expected_add_prefix_space);
}

/// Test to verify vocab type and add_prefix_space from tokenizer metadata
#[test]
#[ignore = "Ignored by default to avoid frequent HF hub downloads"]
fn test_tokenizer_info() {
    for &(tokenizer_id, vocab_type, add_prefix_space) in TEST_TOKENIZER_CASES {
        tracing::info!("Testing tokenizer: {}", tokenizer_id);

        let tokenizer = common::load_tokenizer(tokenizer_id).expect("Failed to load tokenizer");

        let tokenizer_info = TokenizerInfo::from_pretrained(tokenizer_id, None, None, None)
            .expect("Failed to get tokenizer info");
        assert_metadata(&tokenizer_info, vocab_type, add_prefix_space);

        assert_eq!(tokenizer.get_vocab_size(true), tokenizer_info.get_vocab_size() as usize);
    }
}

/// Round-trip TokenizerInfo through serialize_json / deserialize_json.
#[test]
fn test_tokenizer_info_serialize_roundtrip() {
    let tokenizer_info = TokenizerInfo::from_pretrained("openai/gpt-oss-20b", None, None, None)
        .expect("Failed to load tokenizer info");

    let json = tokenizer_info.serialize_json();
    assert!(!json.is_empty());

    let restored = TokenizerInfo::deserialize_json(&json).expect("deserialization should succeed");
    assert_eq!(restored.get_vocab_size(), tokenizer_info.get_vocab_size());
    assert_eq!(restored.get_vocab_type(), tokenizer_info.get_vocab_type());
    assert_eq!(restored.get_add_prefix_space(), tokenizer_info.get_add_prefix_space());

    // Serializing the restored info must reproduce the same JSON.
    assert_eq!(json, restored.serialize_json());

    // The vocab survives byte-for-byte except for tokens containing NUL:
    // xgrammar v0.2.3 truncates strings at the first NUL byte during
    // serialization (ByteToLatin1 in cpp/support/encoding.h iterates with
    // C-string semantics), so such tokens come back truncated.
    let original_vocab = tokenizer_info.get_decoded_vocab();
    let restored_vocab = restored.get_decoded_vocab();
    assert_eq!(original_vocab.len(), restored_vocab.len());
    for (i, (original, restored)) in original_vocab.iter().zip(restored_vocab.iter()).enumerate() {
        let expected = match original.find('\0') {
            Some(nul_idx) => &original[..nul_idx],
            None => original.as_str(),
        };
        assert_eq!(restored.as_str(), expected, "vocab mismatch at token id {i}");
    }
}

#[test]
fn test_tokenizer_info_deserialize_error() {
    use xgrammar::XGrammarErr;

    let Err(XGrammarErr::InvalidJson(err_msg)) = TokenizerInfo::deserialize_json("not json") else {
        panic!("Expected InvalidJson");
    };
    assert!(err_msg.contains("Invalid JSON error"), "unexpected message: {err_msg}");
}

#[test]
fn test_tokenizer_info_deserialize_untyped_error() {
    use xgrammar::XGrammarErr;

    // Valid JSON that is not an object bypasses the typed error construction
    // upstream (a plain runtime_error from picojson) and surfaces as the
    // untyped fallback for this type.
    let Err(XGrammarErr::TokenizerInfoError(_)) = TokenizerInfo::deserialize_json("[1, 2, 3]")
    else {
        panic!("Expected TokenizerInfoError");
    };
}