use memra_gguf::GgufFile;
use memra_tokenizer::Tokenizer;
fn unhex(s: &str) -> Vec<u8> {
assert!(s.len() % 2 == 0, "odd hex length");
(0..s.len() / 2)
.map(|i| u8::from_str_radix(&s[2 * i..2 * i + 2], 16).expect("bad hex"))
.collect()
}
fn parse_ids(s: &str) -> Vec<u32> {
if s.is_empty() {
return Vec::new();
}
s.split(',').map(|x| x.parse().expect("bad id")).collect()
}
fn first_div(a: &[u32], b: &[u32]) -> usize {
let mut i = 0;
while i < a.len() && i < b.len() && a[i] == b[i] {
i += 1;
}
i
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut args = std::env::args().skip(1);
let usage = "usage: tok-parity <model.gguf> <corpus.tsv> <ref-ids.tsv>";
let model = args.next().expect(usage);
let corpus_path = args.next().expect(usage);
let ref_path = args.next().expect(usage);
let g = GgufFile::open(&model)?;
let tok = Tokenizer::from_gguf(&g)?;
println!(
"tok-parity: model={model} pre={} vocab={} bos={:?} add_bos matches encode(add_special)",
tok.pre(),
tok.vocab_size(),
tok.bos_id()
);
let mut refs = std::collections::HashMap::new();
for line in std::fs::read_to_string(&ref_path)?.lines() {
let mut f = line.splitn(3, '\t');
let name = f.next().expect("ref name").to_string();
let spec = parse_ids(f.next().expect("ref special ids"));
let plain = parse_ids(f.next().expect("ref plain ids"));
refs.insert(name, (spec, plain));
}
let corpus = std::fs::read_to_string(&corpus_path)?;
let (mut n_cases, mut n_fail) = (0usize, 0usize);
for line in corpus.lines() {
let mut f = line.splitn(2, '\t');
let name = f.next().expect("corpus name");
let text = String::from_utf8(unhex(f.next().expect("corpus hex")))?;
let (ref_spec, ref_plain) = refs
.get(name)
.unwrap_or_else(|| panic!("no reference ids for case {name:?}"));
n_cases += 1;
let mut case_ok = true;
for (mode, add_special, want) in
[("special", true, ref_spec), ("plain", false, ref_plain)]
{
let got = tok.encode(&text, add_special);
if &got != want {
case_ok = false;
let d = first_div(&got, want);
println!("MISMATCH {name} [{mode}] first_div={d} memra_len={} ref_len={}",
got.len(), want.len());
println!(" text bytes: {:?}", text);
println!(" memra ids: {got:?}");
println!(" ref ids: {want:?}");
let ctx = d.saturating_sub(2);
println!(" memra tail decode from {ctx}: {:?}",
tok.decode_special(&got[ctx..], true));
println!(" ref tail decode from {ctx}: {:?}",
tok.decode_special(&want[ctx..], true));
}
}
if case_ok {
println!("OK {name}");
} else {
n_fail += 1;
}
}
println!("tok-parity: {}/{n_cases} cases identical in BOTH modes", n_cases - n_fail);
if n_fail > 0 {
println!("tok-parity: FAIL ({n_fail} mismatching case(s))");
std::process::exit(1);
}
println!("tok-parity: PASS");
Ok(())
}