use std::sync::OnceLock;
use tokenizers::Tokenizer;
use super::{LazyTable, TokenIndex};
use crate::embeddings::granite::{measuring_tokenizer_from_bytes, test_artifact};
fn measuring_tok() -> Tokenizer {
measuring_tokenizer_from_bytes(test_artifact::tokenizer_bytes()).expect("measuring tokenizer")
}
fn oracle(tok: &Tokenizer, s: &str) -> usize {
tok
.encode(s, true)
.expect("encode substring")
.get_ids()
.len()
}
struct Rng(u64);
impl Rng {
fn next_u64(&mut self) -> u64 {
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut z = self.0;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
fn below(&mut self, n: usize) -> usize {
(self.next_u64() % (n as u64)) as usize
}
}
fn char_boundaries(text: &str) -> Vec<usize> {
(0..=text.len())
.filter(|&i| text.is_char_boundary(i))
.collect()
}
#[track_caller]
fn check(index: &TokenIndex, tok: &Tokenizer, text: &str, a: usize, b: usize) {
let got = index
.measure_range(tok, text, a, b)
.expect("measure_range must not fail on the granite tokenizer");
let want = oracle(tok, &text[a..b]);
assert_eq!(
got,
want,
"measure_range({a}, {b}) = {got} but encode({:?}) = {want} (full text {:?})",
&text[a..b],
text
);
}
fn differential_over(tok: &Tokenizer, text: &str, rng: &mut Rng) {
if text.is_empty() {
return;
}
let index = TokenIndex::build(tok, text).expect("build index");
let bounds = char_boundaries(text);
let m = bounds.len();
if m <= 96 {
for i in 0..m {
for j in (i + 1)..m {
check(&index, tok, text, bounds[i], bounds[j]);
}
}
return;
}
for _ in 0..400 {
let b = bounds[rng.below(m)];
if b > 0 {
check(&index, tok, text, 0, b);
}
let a = bounds[rng.below(m)];
if a < text.len() {
check(&index, tok, text, a, text.len());
}
}
for _ in 0..1_500 {
let mut a = bounds[rng.below(m)];
let mut b = bounds[rng.below(m)];
if a == b {
continue;
}
if a > b {
std::mem::swap(&mut a, &mut b);
}
check(&index, tok, text, a, b);
}
}
const ADVERSARIAL: &[&str] = &[
"1234567890",
"a1234567890b",
"192.168.100.254",
"Order #A1234-99 shipped 2026-07-18 at 09:41:59",
"1000000 2000000 3000000 phone 555-0142 zip 94107-1234",
"a b c",
"a\t\t\tb\tc",
"a\r\nb\r\n\r\nc",
"word\n\n\nword\n\nword",
"a\u{00A0}\u{00A0}\u{00A0}b",
"a\u{2009}\u{2009}\u{2009}b",
"mix \u{00A0}\t \u{2003}\u{2009} run x",
"half\u{3000}width\u{3000}\u{3000}ideographic",
" leading and trailing ",
" ",
"\n\n",
"can't won't I'll we've they're it's",
" spaced words with double gaps ",
"cafe\u{0301} na\u{0308}ive re\u{0301}sume\u{0301}",
"best 🍕🍅🧀 crust 🌿 done",
"👍🏽 family 👨\u{200D}👩\u{200D}👧\u{200D}👦 flag 🏳️\u{200D}🌈 end",
"x½ + ¼ = ¾ y² z₃",
"使用 transformers 库 for retrieval 检索 テスト 테스트",
"Café — naïve façade; ½ + ¼ = ¾. \"Quotes\" & <tags> and\ttabs.",
];
fn golden_texts() -> Vec<String> {
const CORPUS: &str = include_str!("../../../../tests/granite/fixtures/goldens/corpus.json");
let v: serde_json::Value = serde_json::from_str(CORPUS).expect("parse corpus.json");
v["entries"]
.as_array()
.expect("entries array")
.iter()
.map(|e| e["text"].as_str().expect("entry text").to_string())
.collect()
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_range_matches_encode_over_adversarial_corpus() {
let tok = measuring_tok();
let mut rng = Rng(0x0DDB_1A5E_5EED_1234);
for &text in ADVERSARIAL {
differential_over(&tok, text, &mut rng);
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_range_matches_encode_over_goldens() {
let tok = measuring_tok();
let mut rng = Rng(0xF00D_CAFE_1357_9BDF);
for text in golden_texts() {
differential_over(&tok, &text, &mut rng);
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_range_matches_encode_over_seeded_random_document() {
let tok = measuring_tok();
let mut rng = Rng(0xABCD_1234_5678_9EF0);
let mut doc = String::new();
for p in 0..40u32 {
for w in 0..12u32 {
doc.push_str(&format!("para{p}word{w} "));
}
doc.push_str(&format!(
"{}.{}.{}.{} ",
p,
p * 7,
p * 13 % 256,
p * 251 % 256
));
doc.push_str("can't 你好 café\u{0301} 🍕 ");
doc.push_str("\n\n");
}
differential_over(&tok, &doc, &mut rng);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn differential_is_load_bearing_against_a_shifted_boundary() {
let tok = measuring_tok();
let witnesses: &[(&str, usize, usize)] = &[
("a b c", 0, 1), ("1234567890", 0, 6), ("hello world", 0, 5), ("return (a, b)", 0, 6), ];
for &(text, a, b) in witnesses {
let index = TokenIndex::build(&tok, text).expect("build");
let truth = oracle(&tok, &text[a..b]);
assert_eq!(index.measure_range(&tok, text, a, b).unwrap(), truth);
let b2 = (b + 1..=text.len())
.find(|&i| text.is_char_boundary(i))
.expect("a right neighbour boundary exists");
let shifted = oracle(&tok, &text[a..b2]);
assert_ne!(
truth, shifted,
"witness {text:?} [{a},{b}) vs [{a},{b2}) must differ for the differential to bite"
);
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn direct_only_fallback_is_still_exact() {
let tok = measuring_tok();
let mut rng = Rng(0x1122_3344_5566_7788);
for &text in &[
"hello world foo bar",
"1234567 abc\u{00A0}\u{00A0}def",
"a\n\nb",
] {
let real = TokenIndex::build(&tok, text).expect("build");
let direct = TokenIndex {
pretoken_ends: Vec::new(),
count_prefix: vec![0],
digit: Vec::new(),
direct_only: true,
};
let bounds = char_boundaries(text);
for _ in 0..200 {
let m = bounds.len();
let mut a = bounds[rng.below(m)];
let mut b = bounds[rng.below(m)];
if a == b {
continue;
}
if a > b {
std::mem::swap(&mut a, &mut b);
}
let want = oracle(&tok, &text[a..b]);
assert_eq!(real.measure_range(&tok, text, a, b).unwrap(), want);
assert_eq!(direct.measure_range(&tok, text, a, b).unwrap(), want);
}
}
}
const KILLERS: &[&str] = &[
"456 1",
"a 1",
"12 34 56",
"x 9",
"a\t\t9",
"a\u{00A0}\u{00A0}9",
"a\u{0085}\u{0085}9",
"a\u{2009}\u{2009}9",
"a !",
"a\t\t.",
"a\u{00A0}\u{00A0}#",
"a 🍕",
"b\u{2009}\u{2009}🌿",
"a b",
"a b c",
"a\t\tb",
"a\u{2003}\u{2003}b",
"a\u{3000}\u{3000}b",
"1 2 3 4",
"n\u{00A0}9\u{00A0}\u{00A0}8",
"a!\r\n\r\n Next t",
"x.\r\n\r\ny",
"p!\r\nq",
"u/\r\n/v",
"end.\r\n\r\n\r\nStart here",
"a?!\r\n\r\nB",
"a!\r\n1",
"a!\r\n b",
"a!\r\n\r\nb",
" it'station end",
"can'ther",
"we'reunited now",
"I'lloop back",
"he'daily",
"you'venue",
"she'small",
"It'STELLAR",
"cafe\u{0301}s here",
"a\u{0345}b",
"क\u{093E}ख",
"ন\u{09BE}দ",
"a\u{05B4}b",
"a\u{064B}c",
"re\u{0301}sume\u{0301} now",
"1234567890",
"٠١٢٣٤٥٦٧٨٩",
"0123456789",
"a1234567890b",
"192.168.100.254",
"v1234\u{0660}\u{0661}\u{0662}z",
"👨\u{200D}👩\u{200D}👧\u{200D}👦x",
"🏳️\u{200D}🌈y",
"a👍🏽b",
"!!!???...",
"a,,,b",
"(()){}[]",
"http://a/b/c/d",
"/usr/local/bin",
"a//b//c",
"x.\r\n/y/z",
"tag\u{E0067}\u{E0067}\nff\u{5b4}7", ];
const DROPPABLE_KILLERS: &[&str] = &[
"\u{c}\n\u{2003}\u{202f}\nffi\u{5b4}\u{64b}23", "\u{b}\n\u{2003}\u{202f}\nffi\u{5b4}\u{64b}23", "a\u{c}\n\u{2003}word 12",
"x\u{b}\ny\u{2009}\u{2009}9",
"p\u{0}q\nffi\u{5b4}z8", "m\u{7}\n\u{2003}n\u{64b}5", "\u{c}\u{c}\n\u{3000}\u{3000}ffi\u{93e}9",
"list:\u{1f}item\n\u{2003}\u{5b4}4", "a\u{40000}b 12", "q\u{80000}\n\u{2003}r 3", "t\u{100000}u\u{2009}\u{2009}8", ];
fn droppable_fragment_soup(seed: u64, target_chars: usize) -> String {
const FRAGS: &[&str] = &[
"a",
"b",
"c",
"Z",
"it",
"café",
"no",
" ",
" ",
"\t",
"\t\t",
"\u{00A0}",
"\u{00A0}\u{00A0}",
"\u{2009}",
"\u{2003}",
"\u{3000}",
"\r\n",
"\r\n\r\n",
"\n",
"\n\n",
"1",
"12",
"123",
"1234",
"٧",
"٨٩",
"5",
"'s",
"'t",
"'re",
"'ll",
"!",
"!!",
".",
"...",
"/",
"//",
"#",
",",
"(",
")",
"🍕",
"🌿",
"\u{0301}",
"\u{093E}",
"\u{064B}",
" ",
"x",
"y",
"\u{000B}", "\u{000C}", "\u{0000}", "\u{0007}", "\u{40000}", "\u{100000}", ];
let mut rng = Rng(seed);
let mut s = String::new();
while s.chars().count() < target_chars {
s.push_str(FRAGS[rng.below(FRAGS.len())]);
}
s
}
fn count_divergences(
tok: &Tokenizer,
text: &str,
cap: usize,
out: &mut Vec<(String, usize, usize, usize, usize)>,
) -> usize {
if text.is_empty() {
return 0;
}
let index = TokenIndex::build(tok, text).expect("build index");
let bounds = char_boundaries(text);
let m = bounds.len();
let mut tested = 0usize;
let one = |a: usize, b: usize, out: &mut Vec<_>| {
let got = index.measure_range(tok, text, a, b).expect("measure_range");
let want = oracle(tok, &text[a..b]);
if got != want {
out.push((text.to_string(), a, b, got, want));
}
};
if m <= cap {
for i in 0..m {
for j in (i + 1)..m {
one(bounds[i], bounds[j], out);
tested += 1;
}
}
} else {
let mut rng = Rng(0xC0DE_F00D ^ text.len() as u64);
for _ in 0..(cap * cap) {
let mut a = bounds[rng.below(m)];
let mut b = bounds[rng.below(m)];
if a == b {
continue;
}
if a > b {
std::mem::swap(&mut a, &mut b);
}
one(a, b, out);
tested += 1;
}
}
tested
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_range_zero_divergence_over_killers() {
let tok = measuring_tok();
let mut out: Vec<(String, usize, usize, usize, usize)> = Vec::new();
let mut pairs = 0usize;
for &text in KILLERS {
pairs += count_divergences(&tok, text, 96, &mut out);
}
for &text in DROPPABLE_KILLERS {
assert!(
TokenIndex::build(&tok, text)
.expect("build droppable")
.direct_only,
"droppable killer must trip a build-time guard (direct_only): {text:?}"
);
pairs += count_divergences(&tok, text, 96, &mut out);
}
let big: String = KILLERS.join(" | ");
assert!(
!TokenIndex::build(&tok, &big)
.expect("build big join")
.direct_only,
"the big cross-class join must exercise the INDEX path, not the direct_only arm"
);
pairs += count_divergences(&tok, &big, 200, &mut out);
for seed in 0..6u64 {
let soup = droppable_fragment_soup(0xA5A5_0000 ^ (seed.wrapping_mul(0x9E37_79B9)), 160);
pairs += count_divergences(&tok, &soup, 200, &mut out);
}
eprintln!(
"[killer-sweep] pairs_tested={pairs} divergences={}",
out.len()
);
for (t, a, b, got, want) in out.iter().take(25) {
eprintln!(
" DIVERGE measure_range({a},{b})={got} != encode({:?})={want} in {:.48?}",
&t[*a..*b],
t
);
}
assert!(
out.is_empty(),
"{} divergences over {pairs} killer-class pairs (see the list above) — the single-pass \
measure is not exact",
out.len()
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn witness_f1_ws_run_split_by_following_digit() {
let tok = measuring_tok();
let text = "456 1";
let index = TokenIndex::build(&tok, text).expect("build");
assert_eq!(
index.measure_range(&tok, text, 3, 5).unwrap(),
oracle(&tok, " "),
"F1: [3,5) must be encode(\" \"), not the split-run overcount"
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn witness_f2_scan_back_snaps_out_of_punct_crlf_tail() {
let tok = measuring_tok();
let text = "a!\r\n\r\n Next t";
let index = TokenIndex::build(&tok, text).expect("build");
assert_eq!(
index.measure_range(&tok, text, 0, 8).unwrap(),
oracle(&tok, &text[0..8]),
"F2: the `!` head of the punct+CRLF pre-token must not be dropped"
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn witness_f3_contraction_suffix_letter_adjacency() {
let tok = measuring_tok();
let text = " it'station end";
let index = TokenIndex::build(&tok, text).expect("build");
assert_eq!(
index.measure_range(&tok, text, 4, 11).unwrap(),
oracle(&tok, "station"),
"F3: [4,11) must be encode(\"station\"), not \"s\"+\"tation\""
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn witness_f8_vocab_dropped_byte_corrupts_index_offsets() {
let tok = measuring_tok();
let text = "\u{c}\n\u{2003}\u{202f}\nffi\u{5b4}\u{64b}23";
let survived: usize = tok
.encode(text, false)
.expect("encode")
.get_tokens()
.iter()
.map(|t| t.chars().count())
.sum();
assert!(
survived < text.len(),
"the form feed's byte-level char must be dropped: {survived} chars vs {} bytes",
text.len()
);
let index = TokenIndex::build(&tok, text).expect("build");
for &(a, b) in &[(8usize, 16usize), (8, 17), (8, 18)] {
assert_eq!(
index.measure_range(&tok, text, a, b).unwrap(),
oracle(&tok, &text[a..b]),
"F8: [{a},{b}) must equal the direct encode, not the dropped-byte undercount"
);
}
assert!(
index.direct_only,
"the byte-coverage guard must reject the corrupted reconstruction (direct_only)"
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn witness_f9_added_token_lookaround_guard() {
let tok = measuring_tok();
let text = "a<|reserved_200020|>b 12 3";
assert_eq!(&text[1..20], "<|reserved_200020|>");
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
index.direct_only,
"the added-token guard must reject the reconstruction (direct_only)"
);
assert_eq!(
index.measure_range(&tok, text, 1, 20).unwrap(),
oracle(&tok, "<|reserved_200020|>"),
"F9: [1,20) must equal encode(\"<|reserved_200020|>\") = 3, not the index over-count"
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn added_literals_guard_fires_and_measures_exactly() {
let tok = measuring_tok();
for text in [
"pre <|endoftext|> post",
"<|startoftext|>x<|return|>",
"a<|reserved_200020|>b 12 3",
"no<|end|>break\u{2009}\u{2009}9",
"[MASK] in text",
"mix <|reserved_200063|> and <|call|> here",
] {
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
index.direct_only,
"the added-token guard must fire (direct_only) for {text:?}"
);
let bounds = char_boundaries(text);
for i in 0..bounds.len() {
for j in (i + 1)..bounds.len() {
check(&index, &tok, text, bounds[i], bounds[j]);
}
}
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn real_nonbmp_drops_fall_back_and_stay_exact() {
fn lead(cp: char) -> u8 {
let mut buf = [0u8; 4];
cp.encode_utf8(&mut buf).as_bytes()[0]
}
assert_eq!(lead('\u{40000}'), 0xF1);
assert_eq!(lead('\u{80000}'), 0xF2);
assert_eq!(lead('\u{100000}'), 0xF4);
assert_eq!(lead('\u{7FFFF}'), 0xF1);
assert_eq!(lead('\u{10FFFD}'), 0xF4);
let tok = measuring_tok();
for text in [
"a\u{40000}b 12", "a\u{40000}b 12 3", "q\u{80000}\n\u{2003}r 3", "q\u{80000}\n\u{2003}\u{FB03}\u{5b4}r 3", "t\u{100000}u\u{2009}\u{2009}8", "x \u{7FFFF}\u{7FFFF} y9\u{2009}\u{2009}8", "\u{10FFFD}\u{10FFFD} end.\r\n\r\nNext", ] {
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
index.direct_only,
"non-BMP drop must fall back to direct_only for {text:?}"
);
let bounds = char_boundaries(text);
for i in 0..bounds.len() {
for j in (i + 1)..bounds.len() {
check(&index, &tok, text, bounds[i], bounds[j]);
}
}
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn tag_and_pua_do_not_drop_and_index_stays_exact() {
fn lead(cp: char) -> u8 {
let mut buf = [0u8; 4];
cp.encode_utf8(&mut buf).as_bytes()[0]
}
assert_eq!(lead('\u{E0067}'), 0xF3); assert_eq!(lead('\u{F0000}'), 0xF3);
let tok = measuring_tok();
for text in [
"tag\u{E0067}\u{E0067}\nff\u{5b4}7",
"lang\u{E0001}tag 12 3",
"pua\u{F0000}\u{FFFFD}z\u{2009}\u{2009}9",
] {
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
!index.direct_only,
"0xF3-lead chars do not drop; the index path must stay live for {text:?}"
);
let bounds = char_boundaries(text);
for i in 0..bounds.len() {
for j in (i + 1)..bounds.len() {
check(&index, &tok, text, bounds[i], bounds[j]);
}
}
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn k5_free_soup_exercises_index_path() {
fn clean_soup(seed: u64, target_chars: usize) -> String {
const FRAGS: &[&str] = &[
"a",
"b",
"c",
"Z",
"it",
"café",
"no",
" ",
" ",
"\t",
"\t\t",
"\u{00A0}",
"\u{00A0}\u{00A0}",
"\u{2009}",
"\u{2003}",
"\u{3000}",
"\r\n",
"\r\n\r\n",
"\n",
"\n\n",
"1",
"12",
"123",
"1234",
"٧",
"٨٩",
"5",
"'s",
"'t",
"'re",
"'ll",
"!",
"!!",
".",
"...",
"/",
"//",
"#",
",",
"(",
")",
"🍕",
"🌿",
"\u{0301}",
"\u{093E}",
"\u{064B}",
" ",
"x",
"y",
];
let mut rng = Rng(seed);
let mut s = String::new();
while s.chars().count() < target_chars {
s.push_str(FRAGS[rng.below(FRAGS.len())]);
}
s
}
let tok = measuring_tok();
let mut out: Vec<(String, usize, usize, usize, usize)> = Vec::new();
let mut pairs = 0usize;
for seed in 100..106u64 {
let soup = clean_soup(0x5EED_0000 ^ (seed.wrapping_mul(0x9E37_79B9)), 160);
let index = TokenIndex::build(&tok, &soup).expect("build");
assert!(
!index.direct_only,
"clean soup must take the index path (not direct_only): {soup:?}"
);
pairs += count_divergences(&tok, &soup, 200, &mut out);
}
eprintln!("[k5-free-soup] pairs={pairs} divergences={}", out.len());
for (t, a, b, got, want) in out.iter().take(10) {
eprintln!(" DIVERGE ({a},{b}) got={got} want={want} in {t:?}");
}
assert!(
out.is_empty(),
"{} divergences over {pairs} pairs",
out.len()
);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn oversized_single_pretoken_encodes_each_probe_once() {
let tok = measuring_tok();
let n = 4 * 1024 * 1024;
let text = "a".repeat(n);
let index = TokenIndex::build(&tok, &text).expect("build");
assert!(
!index.direct_only,
"'a'*n stays on the live index path (no dropped byte, no added literal)"
);
assert_eq!(
index.pretoken_ends.len(),
1,
"the oversized-single case requires the whole input to be ONE pre-token"
);
let probes = [
(1_000usize, 1_001usize), (1_000, 1_010),
(1_000, 11_000),
(1_000, 211_000), (2_000_000, 2_050_000), (n - 50, n - 1), ];
for &(a, b) in &probes {
super::encode_meter::reset();
let got = index
.measure_range(&tok, &text, a, b)
.expect("measure_range must not fail on the granite tokenizer");
let calls = super::encode_meter::calls();
let bytes = super::encode_meter::get();
assert_eq!(
calls, 1,
"probe [{a},{b}) must encode ONCE, not twice: metered {calls} encode calls"
);
assert_eq!(
bytes,
b - a,
"probe [{a},{b}) must encode exactly its own {} bytes once, not {} (the pre-fix \
double-encode)",
b - a,
2 * (b - a)
);
assert_eq!(
got,
oracle(&tok, &text[a..b]),
"reused whole-query count must equal encode(&text[{a}..{b}], true).len()"
);
}
}
#[track_caller]
fn check_fast(
index: &TokenIndex,
tok: &Tokenizer,
lane: &mut super::FastLane<'_>,
text: &str,
a: usize,
b: usize,
) {
let got = index
.measure_range_fast(tok, text, a, b, lane)
.expect("measure_range_fast must not fail on the granite tokenizer");
let want = oracle(tok, &text[a..b]);
assert_eq!(
got,
want,
"measure_range_fast({a}, {b}) = {got} but encode({:?}) = {want}",
&text[a..b]
);
}
fn merge_table(tok: &Tokenizer) -> super::MergeTable {
super::MergeTable::from_tokenizer(tok).expect("the artifact's BPE is mirrorable")
}
fn built(tok: &Tokenizer) -> OnceLock<Option<super::MergeTable>> {
OnceLock::from(Some(merge_table(tok)))
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn fast_lane_matches_encode_over_packer_shaped_probes() {
let tok = measuring_tok();
let cell = built(&tok);
let cjk = "你好世界模型推理文本嵌入检索".repeat(60);
let xs = "x".repeat(1500);
let mixed = "你好hello世界world".repeat(40);
for text in [cjk.as_str(), xs.as_str(), mixed.as_str()] {
let index = TokenIndex::build(&tok, text).expect("build");
let mut lane = super::FastLane::engaged(512, LazyTable::new(&cell, &tok));
let bounds = char_boundaries(text);
for &start_i in &[1usize, 2, 5, 17, 60, 200] {
if start_i >= bounds.len() - 1 {
continue;
}
let a = bounds[start_i];
for &b in bounds[start_i + 1..].iter().take(400) {
check_fast(&index, &tok, &mut lane, text, a, b);
}
}
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn fast_lane_matches_encode_over_random_ranges_of_prose() {
let tok = measuring_tok();
let cell = built(&tok);
let mut rng = Rng(0x5EED_5EED_0072_0072);
let mut doc = String::new();
for p in 0..30u32 {
doc.push_str(&format!(
"internationalization{p} 你好世界模型推理文本嵌入检索 can't 12345 \
supercalifragilisticexpialidocious résumé ",
));
if p % 5 == 4 {
doc.push_str("\n\n");
}
}
let index = TokenIndex::build(&tok, &doc).expect("build");
let mut lane = super::FastLane::engaged(512, LazyTable::new(&cell, &tok));
let bounds = char_boundaries(&doc);
for _ in 0..3000 {
let i = rng.below(bounds.len() - 1);
let span = 1 + rng.below(40.min(bounds.len() - 1 - i));
let (a, b) = (bounds[i], bounds[i + span]);
check_fast(&index, &tok, &mut lane, &doc, a, b);
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn fast_lane_falls_back_when_the_suffix_is_a_vocabulary_entry() {
let tok = measuring_tok();
let cell = built(&tok);
let text = "模型推理你好";
let index = TokenIndex::build(&tok, text).expect("build");
let mut lane = super::FastLane::engaged(512, LazyTable::new(&cell, &tok));
let a = "模型推理".len();
for b in ["模型推理你".len(), text.len()] {
check_fast(&index, &tok, &mut lane, text, a, b);
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_within_floor_agrees_with_measure_at_every_limit() {
use windit::split::MeasureText;
let tok = measuring_tok();
let cell = built(&tok);
let text = "你好世界模型推理文本嵌入检索 hello world, it's 12345 x".repeat(20);
let index = TokenIndex::build(&tok, &text).expect("build");
let m = super::IndexMeasure::new(&text, &index, &tok, LazyTable::new(&cell, &tok), 512);
let bounds = char_boundaries(&text);
let mut rng = Rng(0x0072_0072);
for _ in 0..400 {
let i = rng.below(bounds.len() - 1);
let span = 1 + rng.below((bounds.len() - 1 - i).min(60));
let s = &text[bounds[i]..bounds[i + span]];
let truth = oracle(&tok, s);
for limit in [
0usize,
1,
2,
3,
truth.saturating_sub(1),
truth,
truth + 1,
10_000,
] {
let got = m.measure_within(s, limit);
assert_eq!(
got,
(truth <= limit).then_some(truth),
"{s:?} at limit {limit}"
);
}
}
super::encode_meter::reset();
let (lo, hi) = (bounds[1], bounds[100]);
assert_eq!(m.measure_within(&text[lo..hi], 0), None);
assert_eq!(super::encode_meter::calls(), 0);
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn measure_within_never_applies_the_floor_without_byte_coverage() {
use windit::split::MeasureText;
let tok = measuring_tok();
let cell = built(&tok);
let text = "\0\0";
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
index.is_direct_only(),
"premise: NUL has no symbol, so the index is direct_only"
);
let m = super::IndexMeasure::new(text, &index, &tok, LazyTable::new(&cell, &tok), 2);
assert_eq!(oracle(&tok, text), 2);
assert_eq!(m.measure_within(text, 2), Some(2));
assert_eq!(m.measure_within(text, 1), None);
let covered = TokenIndex::build(&tok, "abc").expect("build");
let mc = super::IndexMeasure::new("abc", &covered, &tok, LazyTable::new(&cell, &tok), 2);
assert!(!covered.is_direct_only());
let foreign = String::from("\0\0");
assert_eq!(mc.measure_within(&foreign, 2), Some(2));
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn the_lane_refuses_a_pre_token_with_an_uppercase_run_after_its_first_char() {
let tok = measuring_tok();
let cell = built(&tok);
for text in [
"中UCCESSa",
"中中中UCCESSa中中",
"中WARRANTIESa",
"检索文本SUCCESS返回",
"检索文本WARRANTY条款",
"检索文本MERCHANTABILITY条款",
"检索文本Windows系统",
"检索文本API接口",
"检索文本Dža",
"检索文本ᾈ返回",
"aUCCESSa",
"Abc中文",
] {
let index = TokenIndex::build(&tok, text).expect("build");
assert!(
!index.is_direct_only(),
"{text:?} is a covered, added-token-free text"
);
let mut lane = super::FastLane::engaged(512, LazyTable::new(&cell, &tok));
let bounds = char_boundaries(text);
for (i, &a) in bounds.iter().enumerate() {
for &b in &bounds[i + 1..] {
check_fast(&index, &tok, &mut lane, text, a, b);
}
}
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn the_tail_class_agrees_with_the_pre_tokenizer_over_the_vocabularys_chars() {
use tokenizers::{OffsetReferential, OffsetType, PreTokenizedString, PreTokenizer};
let tok = measuring_tok();
let tail = super::bpe_mirror::TailClass::new().expect("compiles");
let pre = tok
.get_pre_tokenizer()
.expect("the artifact has a pre-tokenizer");
let glued = |c: char| {
let mut p = PreTokenizedString::from(format!("a{c}").as_str());
pre.pre_tokenize(&mut p).expect("pre-tokenize");
p.get_splits(OffsetReferential::Original, OffsetType::Byte)
.len()
== 1
};
let mut chars: std::collections::BTreeSet<char> = (0..tok.get_vocab_size(true))
.filter_map(|id| tok.decode(&[u32::try_from(id).expect("id")], false).ok())
.flat_map(|s| s.chars().collect::<Vec<_>>())
.collect();
chars.extend("DžLjNjDzᾈᾘᾨᾼῌῼAΑАaαаʰ々中א가\u{301}\u{903}\u{20dd}0 \t\r\n'.。!".chars());
let (mut tails, mut total) = (0usize, 0usize);
for c in chars {
if c == 'a' {
continue;
}
assert_eq!(tail.contains(c), glued(c), "U+{:04X} {c:?}", c as u32);
total += 1;
tails += usize::from(tail.contains(c));
}
assert!(
tails > 1_000 && total - tails > 200,
"{tails} tail chars of {total}"
);
for c in [
'a', 'α', 'а', 'ʰ', '々', '中', 'א', '가', '\u{301}', '\u{903}', '\u{20dd}',
] {
assert!(tail.contains(c), "{c:?} is a tail char");
}
for c in ['A', 'Α', 'А', 'Dž', 'ᾈ', 'ῼ', '0', ' ', '\'', '。', '\n'] {
assert!(!tail.contains(c), "{c:?} is not a tail char");
}
}
#[test]
#[ignore = "requires the granite tokenizer.json staged beside the model bundle (EMBEDKIT_TEST_MODELS)"]
fn an_added_token_spelled_like_a_word_does_not_fool_the_whole_word_shortcut() {
let base = measuring_tok();
let spelled = merge_table(&base).spell(b" zzqxjkw");
assert_eq!(spelled, "\u{120}zzqxjkw");
assert!(
base.token_to_id(&spelled).is_none(),
"premise: not a token yet"
);
let mut value: serde_json::Value =
serde_json::from_slice(test_artifact::tokenizer_bytes()).expect("parse");
let next_id = u64::try_from(base.get_vocab_size(true)).expect("id");
value["added_tokens"]
.as_array_mut()
.expect("array")
.push(serde_json::json!({
"id": next_id, "content": spelled, "single_word": false, "lstrip": false,
"rstrip": false, "normalized": false, "special": false
}));
let hacked = Tokenizer::from_bytes(serde_json::to_vec(&value).expect("serialize")).expect("load");
assert_eq!(
hacked.token_to_id(&spelled),
Some(u32::try_from(next_id).expect("id"))
);
let cell = OnceLock::from(Some(
super::MergeTable::from_tokenizer(&hacked).expect("premise: the pins accept it"),
));
for text in ["1 zzqxjkwabc1", "1 zzqxjkw1", "中文 zzqxjkw中文"] {
let index = TokenIndex::build(&hacked, text).expect("build");
assert!(
!index.is_direct_only(),
"premise: the literal {spelled:?} is absent from {text:?}, so the guard does not fire"
);
let mut lane = super::FastLane::engaged(512, LazyTable::new(&cell, &hacked));
assert_eq!(
oracle(&hacked, " zzqxjkw"),
6,
"premise: the crate runs the merges"
);
let bounds = char_boundaries(text);
for (i, &a) in bounds.iter().enumerate() {
for &b in &bounds[i + 1..] {
check_fast(&index, &hacked, &mut lane, text, a, b);
}
}
}
}