use super::*;
use crate::core::policy::PolicyError;
use crate::core::tokenize::Tokenize;
use crate::core::Backend;
#[test]
fn dispatches_bpe_byte_level() {
let json = r#"{
"added_tokens": [{"id": 2, "content": "<|endoftext|>", "special": true}],
"normalizer": null,
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "Ġa": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("bpe ok");
assert_eq!(tok.family(), "BPE");
assert!(matches!(tok.backend(), Backend::Bpe(_)));
if let Backend::Bpe(t) = tok.backend() {
assert_eq!(t.encode_with_special("<|endoftext|>"), vec![2]);
}
}
#[test]
fn bpe_split_regex_is_read_from_json() {
let json = r#"{
"added_tokens": [],
"pre_tokenizer": {"type": "Sequence", "pretokenizers": [
{"type": "Split", "pattern": {"Regex": "\\w+|[^\\w\\s]+"}, "behavior": "Isolated"},
{"type": "ByteLevel", "add_prefix_space": false}
]},
"model": {"type": "BPE", "vocab": {"a": 0, "Ġa": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("bpe ok");
assert!(matches!(tok.backend(), Backend::Bpe(_)));
}
#[test]
fn bpe_applies_nfc_normalizer_to_content() {
let json = r#"{
"added_tokens": [],
"normalizer": {"type": "NFC"},
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE",
"vocab": {"a": 0, "Ã": 1, "©": 2, "é": 3},
"merges": ["Ã ©"]}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("bpe ok");
if let Backend::Bpe(t) = tok.backend() {
let precomposed = t.encode("é"); let combining = t.encode("e\u{301}"); assert_eq!(precomposed, combining);
} else {
panic!("expected BPE backend");
}
}
#[test]
fn rejects_non_byte_level_keys_when_byte_level() {
let json = r#"{
"pre_tokenizer": {"type": "ByteLevel"},
"model": {"type": "BPE", "vocab": {"你好": 0}, "merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(matches!(err, Err(HfJsonError::InvalidByteLevel(_))));
}
#[test]
fn bpe_honors_merge_order_independent_of_ids() {
let json = r#"{
"pre_tokenizer": {"type": "ByteLevel"},
"model": {"type": "BPE",
"vocab": {"a":0,"b":1,"c":2,"bc":4,"ab":5},
"merges": [["a","b"], ["b","c"]]}
}"#;
let Backend::Bpe(t) = from_json_bytes(json.as_bytes()).unwrap().into_backend() else {
panic!("bpe");
};
assert_eq!(t.encode("abc"), vec![5, 2]);
}
#[test]
fn unigram_uses_viterbi_not_greedy() {
let json = r#"{
"pre_tokenizer": {"type": "Metaspace"},
"model": {"type": "Unigram", "unk_id": 0, "vocab": [
["<unk>", 0.0], ["</s>", 0.0],
["▁abc", -1.0], ["▁ab", -5.0], ["c", -1.0],
["▁a", -3.0], ["b", -3.0]
]}
}"#;
let Backend::Unigram(t) = from_json_bytes(json.as_bytes()).unwrap().into_backend() else {
panic!("unigram");
};
assert_eq!(t.encode("abc"), vec![2]);
}
#[test]
fn dispatches_unigram() {
let json = r#"{
"added_tokens": [
{"id": 0, "content": "<unk>", "special": true},
{"id": 1, "content": "</s>", "special": true}
],
"model": {
"type": "Unigram",
"unk_id": 0,
"vocab": [["<unk>", 0.0], ["</s>", 0.0], ["a", -1.0], ["b", -2.0]]
}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("unigram ok");
assert_eq!(tok.family(), "Unigram");
assert!(matches!(tok.backend(), Backend::Unigram(_)));
}
#[test]
fn unigram_without_eos_loads() {
let json = r#"{
"added_tokens": [],
"model": {"type": "Unigram", "unk_id": 0, "vocab": [["<unk>", 0.0], ["a", 0.0]]}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("loads");
assert_eq!(tok.family(), "Unigram");
}
#[test]
fn dispatches_wordpiece_with_lowercasing() {
let json = r###"{
"added_tokens": [],
"normalizer": {"type": "BertNormalizer", "lowercase": true, "strip_accents": null},
"pre_tokenizer": {"type": "BertPreTokenizer"},
"model": {
"type": "WordPiece",
"unk_token": "[UNK]",
"continuing_subword_prefix": "##",
"max_input_chars_per_word": 100,
"vocab": {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "hello": 3, "##world": 4, "world": 5}
}
}"###;
let tok = from_json_bytes(json.as_bytes()).expect("wordpiece ok");
assert_eq!(tok.family(), "WordPiece");
let Backend::WordPiece(t) = tok.backend() else {
panic!("expected wordpiece");
};
assert_eq!(t.encode("HELLO"), vec![3]);
assert_eq!(t.encode("hello"), vec![3]);
}
#[test]
fn unigram_applies_replace_normalizer_in_order() {
let json = r#"{
"normalizer": {"type": "Sequence", "normalizers": [
{"type": "Replace", "pattern": {"String": "''"}, "content": "x"}
]},
"pre_tokenizer": {"type": "Metaspace"},
"added_tokens": [{"id": 1, "content": "</s>", "special": true}],
"model": {"type": "Unigram", "unk_id": 0, "vocab": [
["<unk>", 0.0], ["</s>", 0.0], ["▁x", -1.0], ["▁", -2.0], ["x", -3.0]
]}
}"#;
let Backend::Unigram(t) = from_json_bytes(json.as_bytes()).unwrap().into_backend() else {
panic!("unigram");
};
assert_eq!(t.encode("''"), vec![2]);
}
#[test]
fn wordpiece_custom_continuation_prefix() {
let json = r###"{
"model": {"type": "WordPiece", "unk_token": "[UNK]",
"continuing_subword_prefix": "@@", "max_input_chars_per_word": 100,
"vocab": {"[UNK]": 0, "foo": 1, "@@bar": 2}}
}"###;
let Backend::WordPiece(t) = from_json_bytes(json.as_bytes()).unwrap().into_backend() else {
panic!("wordpiece");
};
assert_eq!(t.encode("foobar"), vec![1, 2]);
}
#[test]
fn bpe_matches_added_tokens_in_encode() {
let json = r#"{
"added_tokens": [{"id": 2, "content": "<sp>", "special": false}],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "b": 1, "<sp>": 2}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).unwrap();
assert_eq!(tok.encode("a<sp>b"), vec![0, 2, 1]);
}
#[test]
fn wordpiece_matches_added_tokens_in_input() {
let json = r###"{
"added_tokens": [{"id": 2, "content": "[SEP]", "special": true}],
"model": {"type": "WordPiece", "unk_token": "[UNK]",
"continuing_subword_prefix": "##", "max_input_chars_per_word": 100,
"vocab": {"[UNK]": 0, "a": 1, "[SEP]": 2, "b": 3}}
}"###;
let tok = from_json_bytes(json.as_bytes()).unwrap();
assert_eq!(tok.encode("a [SEP] b"), vec![1, 2, 3]);
}
#[test]
fn unigram_matches_added_tokens_in_input() {
let json = r#"{
"added_tokens": [{"id": 1, "content": "</s>", "special": true}],
"pre_tokenizer": {"type": "Metaspace"},
"model": {"type": "Unigram", "unk_id": 0, "vocab": [
["<unk>", 0.0], ["</s>", 0.0], ["▁a", -1.0], ["▁b", -1.0]
]}
}"#;
let tok = from_json_bytes(json.as_bytes()).unwrap();
assert_eq!(tok.encode("a</s>b"), vec![2, 1, 3]);
}
#[test]
fn post_processor_wraps_with_special_tokens() {
let tok = from_json_bytes(BERT_PAIR_JSON.as_bytes()).unwrap();
assert_eq!(tok.encode_raw("hi"), vec![3]); assert_eq!(tok.encode("hi"), vec![1, 3, 2]); }
const BERT_PAIR_JSON: &str = r###"{
"post_processor": {"type": "BertProcessing", "cls": ["[CLS]", 1], "sep": ["[SEP]", 2]},
"model": {"type": "WordPiece", "unk_token": "[UNK]",
"continuing_subword_prefix": "##", "max_input_chars_per_word": 100,
"vocab": {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "hi": 3, "yo": 4}}
}"###;
#[test]
fn encode_pair_applies_the_bert_pair_template() {
let tok = from_json_bytes(BERT_PAIR_JSON.as_bytes()).unwrap();
assert_eq!(tok.encode_pair("hi", "yo").unwrap(), vec![1, 3, 2, 4, 2]);
}
#[test]
fn encode_pair_without_a_template_errors() {
let json = r#"{
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "b": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).unwrap();
assert!(matches!(
tok.encode_pair("a", "b"),
Err(PolicyError::NoPairTemplate)
));
}
#[test]
fn policy_exposes_named_specials_and_eos() {
let tok = from_json_bytes(
r###"{
"added_tokens": [
{"id": 1, "content": "[CLS]", "special": true},
{"id": 2, "content": "[SEP]", "special": true}
],
"model": {"type": "WordPiece", "unk_token": "[UNK]",
"continuing_subword_prefix": "##", "max_input_chars_per_word": 100,
"vocab": {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "hi": 3}}
}"###
.as_bytes(),
)
.unwrap();
assert_eq!(tok.special_token_id("[CLS]"), Some(1));
assert_eq!(tok.eos_token_id(), Some(2));
assert!(tok.is_eos(2));
}
#[test]
fn decode_skips_special_keeps_nonspecial_added_tokens() {
let json = r#"{
"added_tokens": [
{"id": 2, "content": "<|end|>", "special": true},
{"id": 3, "content": "<sp>", "special": false}
],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "b": 1, "<|end|>": 2, "<sp>": 3}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).unwrap();
let ids = tok.encode("a<|end|><sp>b"); assert_eq!(ids, vec![0, 2, 3, 1]);
assert_eq!(Tokenize::decode(&tok, &ids).unwrap(), "a<sp>b");
}
#[test]
fn unsupported_model_type_errors() {
let json = r#"{"model": {"type": "Phantom", "vocab": {}}}"#;
let err = from_json_bytes(json.as_bytes());
assert!(matches!(err, Err(HfJsonError::UnsupportedModelType(t)) if t == "Phantom"));
}
#[test]
fn missing_model_errors() {
let err = from_json_bytes(b"{}");
assert!(matches!(err, Err(HfJsonError::MissingField("model"))));
}
#[test]
fn infers_unigram_without_model_type() {
let json = r#"{
"added_tokens": [{"id": 1, "content": "</s>", "special": true}],
"model": {"unk_id": 0, "vocab": [["<unk>", 0.0], ["</s>", 0.0], ["x", -1.0]]}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("inferred unigram");
assert_eq!(tok.family(), "Unigram");
}
#[test]
fn infers_wordpiece_without_model_type() {
let json = r#"{
"model": {
"unk_token": "[UNK]", "continuing_subword_prefix": "@@",
"max_input_chars_per_word": 100, "vocab": {"[UNK]": 0, "hi": 1}
}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("inferred wordpiece");
assert_eq!(tok.family(), "WordPiece");
}
#[test]
fn infers_bpe_without_model_type() {
let json = r#"{
"pre_tokenizer": {"type": "ByteLevel"},
"model": {"vocab": {"a": 0, "Ġa": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("inferred bpe");
assert_eq!(tok.family(), "BPE");
}
#[test]
fn unknown_normalizer_errors_not_silently_dropped() {
let json = r#"{
"normalizer": {"type": "Sequence", "normalizers": [
{"type": "NFC"},
{"type": "SomeFutureNormalizer"}
]},
"pre_tokenizer": {"type": "ByteLevel"},
"model": {"type": "BPE", "vocab": {"a": 0}, "merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(
matches!(&err, Err(HfJsonError::UnsupportedNormalizer(t)) if t.contains("SomeFutureNormalizer"))
);
}
#[test]
fn uncompilable_replace_regex_errors_not_literal() {
let json = r#"{
"normalizer": {"type": "Replace", "pattern": {"Regex": "(?P<"}, "content": "x"},
"pre_tokenizer": {"type": "ByteLevel"},
"model": {"type": "BPE", "vocab": {"a": 0}, "merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(matches!(&err, Err(HfJsonError::InvalidNormalizerRegex(_))));
}
#[test]
fn unknown_pretokenizer_without_recognized_split_errors() {
let json = r#"{
"pre_tokenizer": {"type": "UnicodeScripts"},
"model": {"type": "BPE", "vocab": {"a": 0}, "merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(
matches!(&err, Err(HfJsonError::UnsupportedPreTokenizer(t)) if t.contains("UnicodeScripts"))
);
}
#[test]
fn engine_handled_pretokenizer_without_bytelevel_still_loads() {
let json = r#"{
"pre_tokenizer": {"type": "Digits", "individual_digits": true},
"model": {"type": "BPE", "vocab": {"a": 0, "1": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("digits pre-tokenizer loads");
assert_eq!(tok.family(), "BPE");
}
#[test]
fn unknown_pretokenizer_is_ok_when_split_is_anchored() {
let json = r#"{
"pre_tokenizer": {"type": "Sequence", "pretokenizers": [
{"type": "UnicodeScripts"},
{"type": "ByteLevel", "add_prefix_space": false}
]},
"model": {"type": "BPE", "vocab": {"a": 0, "Ġa": 1}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("loads with anchored split");
assert_eq!(tok.family(), "BPE");
}
#[test]
fn added_token_lstrip_is_read_from_json_and_reaches_the_matcher() {
let json = r#"{
"added_tokens": [
{"id": 10, "content": "<mask>", "special": true, "lstrip": true, "rstrip": false},
{"id": 11, "content": "<pad>", "special": true, "lstrip": false, "rstrip": false}
],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "Ġ": 1, "b": 2}, "merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("bpe ok");
let Backend::Bpe(t) = tok.backend() else {
panic!("expected BPE backend");
};
assert_eq!(t.encode("a <mask>b"), vec![0, 10, 2]);
assert_eq!(t.encode("a <pad>b"), vec![0, 1, 11, 2]);
}
#[test]
fn vocab_entry_that_is_also_an_added_token_is_taken_literally() {
let json = r#"{
"added_tokens": [
{"id": 3, "content": "<|eos|>", "special": true},
{"id": 4, "content": "<|tool|>", "special": false}
],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"decoder": {"type": "ByteLevel"},
"model": {"type": "BPE",
"vocab": {"a": 0, "Ġ": 1, "b": 2, "<|eos|>": 3, "<|tool|>": 4},
"merges": []}
}"#;
let tok = from_json_bytes(json.as_bytes()).expect("loads despite literal vocab entries");
assert_eq!(tok.encode_raw("a<|tool|>b"), vec![0, 4, 2]);
assert_eq!(tok.encode_raw("<|eos|>"), vec![3]);
assert_eq!(
Tokenize::decode(&tok, &[0, 4, 2]).expect("decodes"),
"a<|tool|>b"
);
assert_eq!(Tokenize::decode(&tok, &[3]).expect("decodes"), "");
}
#[test]
fn non_added_token_vocab_entry_that_is_not_byte_level_still_errors() {
let json = r#"{
"added_tokens": [{"id": 3, "content": "<|eos|>", "special": true}],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE",
"vocab": {"a": 0, "<|eos|>": 3, "你好": 5},
"merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(matches!(&err, Err(HfJsonError::InvalidByteLevel(t)) if t == "你好"));
}
#[test]
fn added_token_id_disagreeing_with_the_vocab_id_errors() {
let json = r#"{
"added_tokens": [{"id": 7, "content": "<|eos|>", "special": true}],
"pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false},
"model": {"type": "BPE", "vocab": {"a": 0, "<|eos|>": 3}, "merges": []}
}"#;
let err = from_json_bytes(json.as_bytes());
assert!(matches!(
&err,
Err(HfJsonError::AddedTokenIdConflict {
vocab_id: 3,
added_id: 7,
..
})
));
}