metrale_grammar/tokenizer/
hf_metadata.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// HuggingFace metadata detection — port of `class HFTokenizerAnalyzer`
4// and `TokenizerInfo::Impl::DetectMetadataFromHF` in
5// `cpp/tokenizer_info.cc`.
6//
7// Given the JSON text of a HuggingFace `tokenizer.json`, infer:
8//   * the `VocabType` (from the `decoder` field), and
9//   * whether the tokenizer adds a prefix space (from the `normalizer`
10//     or `pre_tokenizer` field).
11//
12// The C++ uses picojson; we use `serde_json`. On any structural failure
13// the C++ logs a warning and falls back to `VocabType::RAW` — we mirror
14// that fall-back, returning the metadata rather than an error.
15
16use serde_json::Value;
17
18use crate::tokenizer::vocab_type::VocabType;
19
20/// Metadata inferred from a HuggingFace `tokenizer.json`.
21///
22/// Mirrors the JSON object produced by C++ `DetectMetadataFromHF`,
23/// which serializes `{"vocab_type": <int>, "add_prefix_space": <bool>}`.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct HfMetadata {
26    /// The detected vocabulary type (`Raw` on detection failure).
27    pub vocab_type: VocabType,
28    /// Whether the tokenizer prepends a space to the first token.
29    pub add_prefix_space: bool,
30}
31
32impl HfMetadata {
33    /// Serialize to the compact JSON form C++ `DetectMetadataFromHF`
34    /// produces: `{"vocab_type":N,"add_prefix_space":bool}`.
35    pub fn to_json(self) -> String {
36        format!(
37            "{{\"vocab_type\":{},\"add_prefix_space\":{}}}",
38            self.vocab_type.as_int(),
39            self.add_prefix_space
40        )
41    }
42}
43
44/// Detect [`HfMetadata`] from the raw text of a `tokenizer.json`.
45///
46/// Returns `Err` only when `backend_str` is not parseable JSON or is not
47/// a JSON object — faithful to the C++ `XGRAMMAR_CHECK(err.empty() &&
48/// v.is<object>())`. A well-formed object that simply lacks the expected
49/// fields yields `HfMetadata { vocab_type: Raw, add_prefix_space: false }`.
50pub fn detect_metadata_from_hf(backend_str: &str) -> Result<HfMetadata, String> {
51    let value: Value = serde_json::from_str(backend_str)
52        .map_err(|e| format!("Failed to parse JSON object: {e}"))?;
53    let obj = value
54        .as_object()
55        .ok_or_else(|| "Failed to parse JSON object: not an object".to_string())?;
56
57    Ok(HfMetadata {
58        vocab_type: detect_vocab_type(obj),
59        add_prefix_space: detect_add_prefix_space(obj),
60    })
61}
62
63type Obj = serde_json::Map<String, Value>;
64
65/// Detect the vocabulary type from the `decoder` field.
66///
67/// Finds `{"type": "ByteFallback"}` or `{"type": "ByteLevel"}` in the
68/// `decoder` field (descending into a `Sequence` decoder's `decoders`
69/// array). Any structural mismatch falls back to `VocabType::Raw`.
70fn detect_vocab_type(obj: &Obj) -> VocabType {
71    let decoder = match obj.get("decoder").and_then(Value::as_object) {
72        Some(d) => d,
73        None => return VocabType::Raw,
74    };
75    let top_type = match decoder.get("type").and_then(Value::as_str) {
76        Some(t) => t,
77        None => return VocabType::Raw,
78    };
79
80    // A `Sequence` decoder nests a `decoders` array; otherwise the
81    // single decoder object is examined directly.
82    let decoders: Vec<&Value> = if top_type == "Sequence" {
83        match decoder.get("decoders").and_then(Value::as_array) {
84            Some(arr) => arr.iter().collect(),
85            None => return VocabType::Raw,
86        }
87    } else {
88        vec![obj.get("decoder").expect("decoder present above")]
89    };
90
91    for d in decoders {
92        let d_obj = match d.as_object() {
93            Some(o) => o,
94            None => return VocabType::Raw,
95        };
96        let ty = match d_obj.get("type").and_then(Value::as_str) {
97            Some(t) => t,
98            None => return VocabType::Raw,
99        };
100        match ty {
101            "ByteLevel" => return VocabType::ByteLevel,
102            "ByteFallback" => return VocabType::ByteFallback,
103            _ => {}
104        }
105    }
106    VocabType::Raw
107}
108
109/// Detect whether the tokenizer adds a prefix space: true if either a
110/// `Prepend` normalizer with `▁`, or a `Metaspace` pre-tokenizer with
111/// `prepend_scheme` of `always`/`first` is present.
112fn detect_add_prefix_space(obj: &Obj) -> bool {
113    detect_prepend_normalizer(obj) || detect_metaspace_pretokenizer(obj)
114}
115
116/// Find `{"type": "Prepend", "prepend": "▁"}` in the `normalizer` field
117/// (descending into a `Sequence` normalizer's `normalizers` array).
118fn detect_prepend_normalizer(obj: &Obj) -> bool {
119    let normalizer_value = match obj.get("normalizer") {
120        Some(v) if v.is_object() => v,
121        _ => return false,
122    };
123    let normalizer = normalizer_value.as_object().expect("checked is_object");
124    let top_type = match normalizer.get("type").and_then(Value::as_str) {
125        Some(t) => t,
126        None => return false,
127    };
128
129    let normalizers: Vec<&Value> = if top_type == "Sequence" {
130        match normalizer.get("normalizers").and_then(Value::as_array) {
131            Some(arr) => arr.iter().collect(),
132            None => return false,
133        }
134    } else {
135        vec![normalizer_value]
136    };
137
138    for n in normalizers {
139        let n_obj = match n.as_object() {
140            Some(o) => o,
141            None => continue,
142        };
143        if n_obj.get("type").and_then(Value::as_str) != Some("Prepend") {
144            continue;
145        }
146        if n_obj.get("prepend").and_then(Value::as_str) == Some("\u{2581}") {
147            return true;
148        }
149    }
150    false
151}
152
153/// Find `pre_tokenizer: {"type": "Metaspace", "prepend_scheme":
154/// "always" | "first"}`.
155fn detect_metaspace_pretokenizer(obj: &Obj) -> bool {
156    let pre = match obj.get("pre_tokenizer").and_then(Value::as_object) {
157        Some(p) => p,
158        None => return false,
159    };
160    let ty = pre.get("type").and_then(Value::as_str);
161    let scheme = pre.get("prepend_scheme").and_then(Value::as_str);
162    match (ty, scheme) {
163        (Some("Metaspace"), Some(s)) => s == "always" || s == "first",
164        _ => false,
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[test]
173    fn invalid_json_is_err() {
174        assert!(detect_metadata_from_hf("{not json").is_err());
175    }
176
177    #[test]
178    fn non_object_json_is_err() {
179        assert!(detect_metadata_from_hf("[1,2,3]").is_err());
180        assert!(detect_metadata_from_hf("42").is_err());
181    }
182
183    #[test]
184    fn empty_object_falls_back_to_raw() {
185        let m = detect_metadata_from_hf("{}").unwrap();
186        assert_eq!(m.vocab_type, VocabType::Raw);
187        assert!(!m.add_prefix_space);
188    }
189
190    #[test]
191    fn byte_level_decoder_detected() {
192        let json = r#"{"decoder":{"type":"ByteLevel"}}"#;
193        let m = detect_metadata_from_hf(json).unwrap();
194        assert_eq!(m.vocab_type, VocabType::ByteLevel);
195    }
196
197    #[test]
198    fn byte_fallback_decoder_detected() {
199        let json = r#"{"decoder":{"type":"ByteFallback"}}"#;
200        let m = detect_metadata_from_hf(json).unwrap();
201        assert_eq!(m.vocab_type, VocabType::ByteFallback);
202    }
203
204    #[test]
205    fn sequence_decoder_descends() {
206        let json = r#"{"decoder":{"type":"Sequence","decoders":[
207            {"type":"Replace"},{"type":"ByteFallback"},{"type":"Fuse"}]}}"#;
208        let m = detect_metadata_from_hf(json).unwrap();
209        assert_eq!(m.vocab_type, VocabType::ByteFallback);
210    }
211
212    #[test]
213    fn sequence_decoder_byte_level_inside() {
214        let json = r#"{"decoder":{"type":"Sequence","decoders":[
215            {"type":"ByteLevel","add_prefix_space":true}]}}"#;
216        let m = detect_metadata_from_hf(json).unwrap();
217        assert_eq!(m.vocab_type, VocabType::ByteLevel);
218    }
219
220    #[test]
221    fn unknown_decoder_is_raw() {
222        let json = r#"{"decoder":{"type":"Replace"}}"#;
223        let m = detect_metadata_from_hf(json).unwrap();
224        assert_eq!(m.vocab_type, VocabType::Raw);
225    }
226
227    #[test]
228    fn sequence_missing_decoders_is_raw() {
229        let json = r#"{"decoder":{"type":"Sequence"}}"#;
230        let m = detect_metadata_from_hf(json).unwrap();
231        assert_eq!(m.vocab_type, VocabType::Raw);
232    }
233
234    #[test]
235    fn prepend_normalizer_sets_prefix_space() {
236        let json = r#"{"normalizer":{"type":"Prepend","prepend":"▁"}}"#;
237        let m = detect_metadata_from_hf(json).unwrap();
238        assert!(m.add_prefix_space);
239    }
240
241    #[test]
242    fn prepend_normalizer_wrong_marker_no_prefix() {
243        let json = r#"{"normalizer":{"type":"Prepend","prepend":" "}}"#;
244        let m = detect_metadata_from_hf(json).unwrap();
245        assert!(!m.add_prefix_space);
246    }
247
248    #[test]
249    fn sequence_normalizer_descends() {
250        let json = r#"{"normalizer":{"type":"Sequence","normalizers":[
251            {"type":"Strip"},{"type":"Prepend","prepend":"▁"}]}}"#;
252        let m = detect_metadata_from_hf(json).unwrap();
253        assert!(m.add_prefix_space);
254    }
255
256    #[test]
257    fn metaspace_pretokenizer_always() {
258        let json = r#"{"pre_tokenizer":{"type":"Metaspace","prepend_scheme":"always"}}"#;
259        let m = detect_metadata_from_hf(json).unwrap();
260        assert!(m.add_prefix_space);
261    }
262
263    #[test]
264    fn metaspace_pretokenizer_first() {
265        let json = r#"{"pre_tokenizer":{"type":"Metaspace","prepend_scheme":"first"}}"#;
266        let m = detect_metadata_from_hf(json).unwrap();
267        assert!(m.add_prefix_space);
268    }
269
270    #[test]
271    fn metaspace_pretokenizer_never_no_prefix() {
272        let json = r#"{"pre_tokenizer":{"type":"Metaspace","prepend_scheme":"never"}}"#;
273        let m = detect_metadata_from_hf(json).unwrap();
274        assert!(!m.add_prefix_space);
275    }
276
277    #[test]
278    fn to_json_format() {
279        let m = HfMetadata {
280            vocab_type: VocabType::ByteLevel,
281            add_prefix_space: false,
282        };
283        assert_eq!(m.to_json(), r#"{"vocab_type":2,"add_prefix_space":false}"#);
284        let m2 = HfMetadata {
285            vocab_type: VocabType::ByteFallback,
286            add_prefix_space: true,
287        };
288        assert_eq!(m2.to_json(), r#"{"vocab_type":1,"add_prefix_space":true}"#);
289    }
290
291    #[test]
292    fn qwen_style_byte_level_with_metaspace_absent() {
293        // Qwen/MiniMax: ByteLevel decoder, no prefix space (F68-critical).
294        let json = r#"{"decoder":{"type":"ByteLevel"},
295            "pre_tokenizer":{"type":"ByteLevel","add_prefix_space":false}}"#;
296        let m = detect_metadata_from_hf(json).unwrap();
297        assert_eq!(m.vocab_type, VocabType::ByteLevel);
298        assert!(!m.add_prefix_space);
299    }
300}
metrale_grammar/tokenizer/hf_metadata.rs

metrale_grammar/tokenizer/
hf_metadata.rs
