mod corpus;
use anyhow::Context;
use corpus::{Case, CORPUS};
use frink_gguf::ShardedGguf;
use frink_models::tokenizer::{
should_add_bos_token, GgufBpeTokenizer, GgufPlamo2Tokenizer, GgufSpmTokenizer,
GgufUnigramTokenizer, GgufWordPieceTokenizer, SpecialTokens,
};
use std::path::Path;
use std::process::Command;
const EXIT_MODEL_UNSUPPORTED: i32 = 3;
#[derive(Debug)]
struct Reference {
add_bos: bool,
add_eos: bool,
n_vocab: usize,
cases: Vec<[Vec<u32>; SETTINGS.len()]>,
}
const SETTINGS: [SpecialTokens; 2] = [SpecialTokens::AsText, SpecialTokens::Parse];
fn setting_name(mode: SpecialTokens) -> &'static str {
match mode {
SpecialTokens::AsText => "parse_special=false",
SpecialTokens::Parse => "parse_special=true",
}
}
#[derive(Debug, PartialEq, Eq)]
pub(super) struct Divergence {
pub index: usize,
pub byte_offset: usize,
pub llama: Vec<(u32, String)>,
pub frink: Vec<(u32, String)>,
pub llama_len: usize,
pub frink_len: usize,
}
pub(super) struct CaseOutcome {
pub name: &'static str,
pub why: &'static str,
pub text: &'static str,
pub mode: SpecialTokens,
pub n_tokens: usize,
pub divergence: Option<Divergence>,
}
pub(super) struct Report {
pub cases: Vec<CaseOutcome>,
pub bos_llama: bool,
pub bos_frink: bool,
pub eos_llama: bool,
pub vocab_llama: usize,
pub vocab_frink: usize,
pub pre: String,
pub model: String,
}
impl Report {
pub fn diverged(&self) -> bool {
self.bos_llama != self.bos_frink
|| self.vocab_llama != self.vocab_frink
|| self.cases.iter().any(|c| c.divergence.is_some())
}
fn n_bad(&self) -> usize {
self.cases.iter().filter(|c| c.divergence.is_some()).count()
}
fn n_tokens(&self) -> usize {
self.cases.iter().map(|c| c.n_tokens).sum()
}
}
enum Encoder {
Bpe(Box<GgufBpeTokenizer>),
Spm(GgufSpmTokenizer),
Unigram(GgufUnigramTokenizer),
WordPiece(GgufWordPieceTokenizer),
Plamo2(Box<GgufPlamo2Tokenizer>),
}
impl Encoder {
fn from_gguf(file: &ShardedGguf) -> anyhow::Result<Self> {
Ok(match file.metadata_str("tokenizer.ggml.model") {
Some("gpt2" | "gemma4") => Encoder::Bpe(Box::new(GgufBpeTokenizer::from_gguf(file)?)),
Some("llama") => Encoder::Spm(GgufSpmTokenizer::from_gguf(file)?),
Some("t5") => Encoder::Unigram(GgufUnigramTokenizer::from_gguf(file)?),
Some("bert") => Encoder::WordPiece(GgufWordPieceTokenizer::from_gguf(file)?),
Some("plamo2") => Encoder::Plamo2(Box::new(GgufPlamo2Tokenizer::from_gguf(file)?)),
other => anyhow::bail!(
"parity does not cover tokenizer {other:?} — it builds gpt2/gemma4 (BPE), \
llama (SPM), t5 (unigram), bert (WordPiece) and plamo2 vocabularies only"
),
})
}
fn encode(&self, text: &str, mode: SpecialTokens) -> Vec<u32> {
match self {
Encoder::Bpe(t) => t.encode(text, mode),
Encoder::Spm(t) => t.encode(text, mode),
Encoder::Unigram(t) => t.encode(text, mode),
Encoder::WordPiece(t) => t.encode(text, mode),
Encoder::Plamo2(t) => t.encode(text, mode),
}
}
fn piece(&self, id: u32) -> String {
match self {
Encoder::Bpe(t) => t.decode(&[id]),
Encoder::Spm(t) => t.decode(&[id]),
Encoder::Unigram(t) => t.decode(&[id]),
Encoder::WordPiece(t) => t.decode(&[id]),
Encoder::Plamo2(t) => t.decode(&[id]),
}
}
fn vocab_size(&self) -> usize {
match self {
Encoder::Bpe(t) => t.vocab_size(),
Encoder::Spm(t) => t.vocab_size(),
Encoder::Unigram(t) => t.vocab_size(),
Encoder::WordPiece(t) => t.vocab_size(),
Encoder::Plamo2(t) => t.vocab_size(),
}
}
}
pub(super) fn run(dumper: &Path, model: &Path) -> anyhow::Result<Option<Report>> {
let file = ShardedGguf::open(model)?;
let encoder = Encoder::from_gguf(&file)?;
let pre = file
.metadata_str("tokenizer.ggml.pre")
.unwrap_or("(unset)")
.to_string();
let Some(reference) = reference_tokenization(dumper, model, CORPUS)? else {
return Ok(None);
};
if reference.cases.len() != CORPUS.len() {
anyhow::bail!(
"reference tokenizer returned {} cases for {} sent",
reference.cases.len(),
CORPUS.len()
);
}
let encoder = &encoder;
let cases = CORPUS
.iter()
.zip(&reference.cases)
.flat_map(|(case, llama_runs)| {
SETTINGS
.iter()
.zip(llama_runs)
.map(move |(&mode, llama_ids)| {
let frink_ids = encoder.encode(case.text, mode);
CaseOutcome {
name: case.name,
why: case.why,
text: case.text,
mode,
n_tokens: llama_ids.len(),
divergence: first_divergence(llama_ids, &frink_ids, &|id| {
encoder.piece(id)
}),
}
})
})
.collect();
Ok(Some(Report {
cases,
bos_llama: reference.add_bos,
bos_frink: should_add_bos_token(&file),
eos_llama: reference.add_eos,
vocab_llama: reference.n_vocab,
vocab_frink: encoder.vocab_size(),
pre,
model: model.display().to_string(),
}))
}
const WINDOW: usize = 3;
pub(super) fn first_divergence(
llama: &[u32],
frink: &[u32],
piece: &dyn Fn(u32) -> String,
) -> Option<Divergence> {
let common = llama
.iter()
.zip(frink)
.position(|(a, b)| a != b)
.unwrap_or_else(|| llama.len().min(frink.len()));
if common == llama.len() && common == frink.len() {
return None;
}
let byte_offset = llama[..common].iter().map(|&id| piece(id).len()).sum();
let lo = common.saturating_sub(WINDOW);
let hi = common + WINDOW + 1;
let window = |ids: &[u32]| -> Vec<(u32, String)> {
ids[lo.min(ids.len())..hi.min(ids.len())]
.iter()
.map(|&id| (id, piece(id)))
.collect()
};
Some(Divergence {
index: common,
byte_offset,
llama: window(llama),
frink: window(frink),
llama_len: llama.len(),
frink_len: frink.len(),
})
}
pub(super) fn print_report(r: &Report) {
let name = Path::new(&r.model)
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.unwrap_or_else(|| r.model.clone());
let verdict = if r.diverged() { "DIVERGES" } else { "MATCH" };
println!(
"tokenizer {name}: {verdict} ({} cases x {} parse_special settings / {} tokens, pre={}, \
frink vs llama.cpp)",
r.cases.len() / SETTINGS.len(),
SETTINGS.len(),
r.n_tokens(),
r.pre
);
println!(
" vocab llama {} / frink {} add_bos llama {} / frink {} (llama add_eos {})",
r.vocab_llama, r.vocab_frink, r.bos_llama, r.bos_frink, r.eos_llama
);
if r.vocab_llama != r.vocab_frink {
println!(
" vocab sizes differ: the two id spaces are not the same, so every case below \
compares numbers that do not mean the same thing."
);
}
if r.bos_llama != r.bos_frink {
println!(
" add-BOS policy differs. Frink's rule is `should_add_bos_token`; llama's is \
`llama_vocab_get_add_bos`. Every prompt on this checkpoint is off by one token."
);
}
if r.n_bad() == 0 {
println!(
" all {} cases tokenize identically under both parse_special settings (digit \
runs, multi-space, indents, blank lines, CJK, emoji, contractions, special-token \
markers as prose).",
r.cases.len() / SETTINGS.len()
);
return;
}
println!(" {}/{} case runs diverge:", r.n_bad(), r.cases.len());
for c in r.cases.iter().filter(|c| c.divergence.is_some()) {
let d = c.divergence.as_ref().expect("filtered on is_some");
println!(
"\n [{} @ {}] token {} of {} (llama) / {} (frink), byte ~{} of {}",
c.name,
setting_name(c.mode),
d.index,
d.llama_len,
d.frink_len,
d.byte_offset,
c.text.len()
);
println!(" why this case: {}", c.why);
println!(" input around it: {}", around(c.text, d.byte_offset));
println!(" llama {}", render_window(&d.llama, d.index));
println!(" frink {}", render_window(&d.frink, d.index));
}
println!(
"\n A divergence here means frink and llama.cpp disagree about the PROMPT, so a \
logit comparison below is measuring two different inputs. Fix this first."
);
}
fn around(text: &str, offset: usize) -> String {
let mid = floor_boundary(text, offset.min(text.len()));
let lo = floor_boundary(text, mid.saturating_sub(24));
let hi = ceil_boundary(text, (mid + 24).min(text.len()));
format!("{:?} >|< {:?}", &text[lo..mid], &text[mid..hi])
}
fn floor_boundary(s: &str, mut i: usize) -> usize {
while i > 0 && !s.is_char_boundary(i) {
i -= 1;
}
i
}
fn ceil_boundary(s: &str, mut i: usize) -> usize {
while i < s.len() && !s.is_char_boundary(i) {
i += 1;
}
i
}
fn render_window(window: &[(u32, String)], index: usize) -> String {
let lo = index.saturating_sub(WINDOW);
window
.iter()
.enumerate()
.map(|(k, (id, piece))| {
let mark = if lo + k == index { "*" } else { "" };
format!("{mark}{id}:{piece:?}")
})
.collect::<Vec<_>>()
.join(" ")
}
fn encode_cases(cases: &[Case]) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(b"FXTK");
out.extend_from_slice(&(cases.len() as u32).to_le_bytes());
for c in cases {
out.extend_from_slice(&(c.text.len() as u32).to_le_bytes());
out.extend_from_slice(c.text.as_bytes());
}
out
}
fn parse_result(bytes: &[u8]) -> anyhow::Result<Reference> {
let u32_at = |at: &mut usize| -> anyhow::Result<u32> {
let end = *at + 4;
let slice = bytes
.get(*at..end)
.context("reference tokenization is truncated")?;
*at = end;
Ok(u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]))
};
if bytes.len() < 4 || &bytes[..4] != b"FXTK" {
anyhow::bail!(
"reference tokenization is not an FXTK file. The dumper at hand is probably older \
than the --tokenize mode; rebuild it with ./tools/build_llama_logits.sh"
);
}
let mut at = 4usize;
let version = u32_at(&mut at)?;
if version != 2 {
anyhow::bail!(
"reference tokenization is FXTK v{version}, this build reads v2 (one run per \
parse_special setting). Rebuild the dumper with ./tools/build_llama_logits.sh"
);
}
let flags = u32_at(&mut at)?;
let n_vocab = u32_at(&mut at)? as usize;
let n_cases = u32_at(&mut at)? as usize;
let mut cases = Vec::with_capacity(n_cases);
for i in 0..n_cases {
let mut runs: [Vec<u32>; SETTINGS.len()] = Default::default();
for run in runs.iter_mut() {
let n = u32_at(&mut at)? as usize;
let remaining = bytes.len().saturating_sub(at) / 4;
if n > remaining {
anyhow::bail!(
"reference tokenization case {i} declares {n} ids with {remaining} left"
);
}
let mut ids = Vec::with_capacity(n);
for _ in 0..n {
let raw = u32_at(&mut at)? as i32;
if raw < 0 {
anyhow::bail!(
"reference tokenization case {i} holds a negative token id {raw}"
);
}
ids.push(raw as u32);
}
*run = ids;
}
cases.push(runs);
}
if at != bytes.len() {
anyhow::bail!(
"reference tokenization has {} trailing bytes after {n_cases} cases",
bytes.len() - at
);
}
Ok(Reference {
add_bos: flags & 1 != 0,
add_eos: flags & 2 != 0,
n_vocab,
cases,
})
}
fn reference_tokenization(
dumper: &Path,
model: &Path,
cases: &[Case],
) -> anyhow::Result<Option<Reference>> {
let dir = std::env::temp_dir();
let pid = std::process::id();
let in_path = dir.join(format!("frink-parity-cases-{pid}.bin"));
let out_path = dir.join(format!("frink-parity-toks-{pid}.bin"));
std::fs::write(&in_path, encode_cases(cases))
.with_context(|| format!("writing corpus to {}", in_path.display()))?;
let out = Command::new(dumper)
.arg("--tokenize")
.arg(model)
.arg(&in_path)
.arg(&out_path)
.output()
.context("running the reference tokenizer")?;
let _ = std::fs::remove_file(&in_path);
if !out.status.success() {
let _ = std::fs::remove_file(&out_path);
if out.status.code() == Some(EXIT_MODEL_UNSUPPORTED) {
return Ok(None);
}
anyhow::bail!(
"reference tokenizer failed: {}\n(if it says \"failed to load --tokenize\", the \
dumper predates this mode — rebuild it with ./tools/build_llama_logits.sh)",
String::from_utf8_lossy(&out.stderr)
.lines()
.last()
.unwrap_or("(no stderr)")
);
}
let bytes = std::fs::read(&out_path)
.with_context(|| format!("reading reference tokenization from {}", out_path.display()))?;
let _ = std::fs::remove_file(&out_path);
parse_result(&bytes).map(Some)
}
#[cfg(test)]
mod tests {
use super::*;
fn fake_piece(id: u32) -> String {
char::from_u32(id).unwrap_or('?').to_string()
}
#[test]
fn identical_sequences_do_not_diverge() {
let ids = vec![97u32, 98, 99, 100];
assert_eq!(first_divergence(&ids, &ids, &fake_piece), None);
}
#[test]
fn a_mismatched_pair_is_reported_as_a_divergence() {
let llama = vec![9000u32, 123, 456, 7, 46];
let frink = vec![9000u32, 1234567, 46];
let d = first_divergence(&llama, &frink, &fake_piece)
.expect("a digit-run split difference must be reported");
assert_eq!(d.index, 1, "they agree only on the leading token");
assert_eq!(d.llama_len, 5);
assert_eq!(d.frink_len, 3);
assert_eq!(d.llama.first().map(|(id, _)| *id), Some(9000));
assert!(d.llama.iter().any(|(id, _)| *id == 123));
assert!(d.frink.iter().any(|(id, _)| *id == 1234567));
}
#[test]
fn an_equal_prefix_with_a_longer_tail_still_diverges() {
let llama = vec![97u32, 98, 99];
let frink = vec![97u32, 98, 99, 100];
let d = first_divergence(&llama, &frink, &fake_piece)
.expect("a trailing extra token is a divergence");
assert_eq!(d.index, 3);
assert_eq!((d.llama_len, d.frink_len), (3, 4));
}
#[test]
fn the_byte_offset_counts_the_agreed_prefix() {
let llama = vec![97u32, 98, 99, 100];
let frink = vec![97u32, 98, 99, 101];
let d = first_divergence(&llama, &frink, &fake_piece).expect("differs at index 3");
assert_eq!(d.index, 3);
assert_eq!(d.byte_offset, 3);
}
#[test]
fn an_empty_pair_agrees_and_a_half_empty_pair_does_not() {
assert_eq!(first_divergence(&[], &[], &fake_piece), None);
let d = first_divergence(&[], &[97], &fake_piece).expect("empty vs one token differs");
assert_eq!(d.index, 0);
assert_eq!(d.byte_offset, 0);
}
#[test]
fn the_window_is_clamped_at_both_ends() {
let llama: Vec<u32> = (97..107).collect();
let frink = vec![200u32];
let d = first_divergence(&llama, &frink, &fake_piece).expect("differs at 0");
assert_eq!(d.index, 0);
assert_eq!(d.llama.len(), WINDOW + 1);
assert_eq!(d.frink.len(), 1);
}
#[test]
fn the_case_file_is_length_prefixed_so_newlines_survive() {
let cases = [
Case {
name: "a",
why: "",
text: "hi\n\n",
},
Case {
name: "b",
why: "",
text: "",
},
];
let blob = encode_cases(&cases);
assert_eq!(&blob[..4], b"FXTK");
assert_eq!(u32::from_le_bytes(blob[4..8].try_into().unwrap()), 2);
assert_eq!(u32::from_le_bytes(blob[8..12].try_into().unwrap()), 4);
assert_eq!(&blob[12..16], b"hi\n\n");
assert_eq!(u32::from_le_bytes(blob[16..20].try_into().unwrap()), 0);
assert_eq!(blob.len(), 20);
}
fn result_blob(flags: u32, n_vocab: u32, cases: &[[&[i32]; 2]]) -> Vec<u8> {
let mut b = Vec::from(*b"FXTK");
b.extend_from_slice(&2u32.to_le_bytes());
b.extend_from_slice(&flags.to_le_bytes());
b.extend_from_slice(&n_vocab.to_le_bytes());
b.extend_from_slice(&(cases.len() as u32).to_le_bytes());
for runs in cases {
for c in runs {
b.extend_from_slice(&(c.len() as u32).to_le_bytes());
for id in *c {
b.extend_from_slice(&id.to_le_bytes());
}
}
}
b
}
#[test]
fn a_result_file_round_trips_flags_vocab_and_both_runs_of_every_case() {
let blob = result_blob(0b11, 128_256, &[[&[1, 2, 3], &[9]], [&[], &[]]]);
let r = parse_result(&blob).expect("well-formed result must parse");
assert!(r.add_bos && r.add_eos);
assert_eq!(r.n_vocab, 128_256);
assert_eq!(r.cases, vec![[vec![1u32, 2, 3], vec![9]], [vec![], vec![]]]);
let none = parse_result(&result_blob(0, 32, &[[&[7], &[7]]])).unwrap();
assert!(!none.add_bos && !none.add_eos);
}
#[test]
fn a_version_one_result_is_refused_and_names_the_rebuild() {
let mut blob = result_blob(0, 32, &[[&[1], &[1]]]);
blob[4..8].copy_from_slice(&1u32.to_le_bytes());
let err = parse_result(&blob).unwrap_err().to_string();
assert!(
err.contains("FXTK v1") && err.contains("build_llama_logits.sh"),
"got {err}"
);
}
#[test]
fn a_malformed_result_is_refused_rather_than_half_read() {
let err = parse_result(b"not fxtk at all").unwrap_err().to_string();
assert!(err.contains("build_llama_logits.sh"), "got {err}");
let mut blob = result_blob(0, 32, &[[&[1, 2, 3], &[1, 2, 3]]]);
blob.truncate(blob.len() - 5);
assert!(parse_result(&blob).is_err());
let mut blob = result_blob(0, 32, &[[&[1], &[1]]]);
blob.push(0);
assert!(parse_result(&blob).is_err());
assert!(parse_result(&result_blob(0, 32, &[[&[-3], &[3]]])).is_err());
let mut blob = result_blob(0, 32, &[[&[1], &[1]]]);
let at = blob.len() - 8;
blob[at..at + 4].copy_from_slice(&u32::MAX.to_le_bytes());
let err = parse_result(&blob).unwrap_err().to_string();
assert!(err.contains("declares"), "got {err}");
}
#[test]
fn the_context_slice_never_splits_a_character() {
for case in CORPUS {
for off in 0..=case.text.len() {
let _ = around(case.text, off);
}
}
assert!(around("日本", 1).contains(">|<"));
}
const REPO: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../..");
#[test]
fn the_unsupported_exit_code_matches_the_c_side() {
let c = std::fs::read_to_string(std::path::Path::new(REPO).join("tools/llama_logits.c"))
.expect("the dumper source is tracked next to this crate");
assert!(
c.contains(&format!(
"#define EXIT_MODEL_UNSUPPORTED {EXIT_MODEL_UNSUPPORTED}"
)),
"tools/llama_logits.c no longer defines EXIT_MODEL_UNSUPPORTED as \
{EXIT_MODEL_UNSUPPORTED}"
);
assert!(
c.matches("return EXIT_MODEL_UNSUPPORTED;").count() >= 2,
"both dumper modes must report an unloadable checkpoint with that code"
);
}
const SWEEP: &[&str] = &[
"Llama-3.2-1B-Instruct-Q4_K_M.gguf", "Qwen2.5-1.5B-Instruct-Q4_K_M.gguf", "DeepSeek-R1-Distill-Qwen-1.5B-Q4_K_M.gguf", "olmoe-1b-7b-0924-q4_0.gguf", "tinyllama-1.1b-chat-v1.0.Q8_0.gguf", "gemma-2-2b-it-Q4_K_M.gguf", "gemma-4-E2B-it-Q4_K_M.gguf", "Phi-4-mini-instruct-Q4_K_M.gguf", "Yi-1.5-6B-Chat-Q4_K_M.gguf",
"Mistral-7B-Instruct-v0.2-Q4_K_M.gguf",
];
#[test]
#[ignore = "needs ./tools/build_llama_logits.sh and checkpoints under models/"]
fn frink_and_llama_cpp_tokenize_the_corpus_identically() {
let dumper = std::path::PathBuf::from(REPO).join("target/llama_logits");
assert!(
dumper.exists(),
"run ./tools/build_llama_logits.sh first ({} is missing)",
dumper.display()
);
let mut checked = 0usize;
let mut bad_ids: Vec<String> = Vec::new();
let mut bad_policy: Vec<String> = Vec::new();
for name in SWEEP {
let model = std::path::PathBuf::from(REPO).join("models").join(name);
if !model.exists() {
println!("skip {name}: not downloaded");
continue;
}
let Some(report) = run(&dumper, &model).expect("tokenizer parity must run") else {
println!("skip {name}: the installed libllama cannot load it");
continue;
};
checked += 1;
print_report(&report);
println!();
if report.n_bad() > 0 || report.vocab_llama != report.vocab_frink {
bad_ids.push(format!("{name} ({} cases)", report.n_bad()));
}
if report.bos_llama != report.bos_frink {
bad_policy.push(format!(
"{name} (pre={}, llama {} / frink {})",
report.pre, report.bos_llama, report.bos_frink
));
}
}
assert!(
checked > 0,
"no checkpoint in SWEEP is present under models/"
);
assert!(
bad_ids.is_empty(),
"these checkpoints tokenize the corpus differently from llama.cpp: {bad_ids:?}"
);
assert!(
bad_policy.is_empty(),
"these checkpoints disagree about add-BOS, so every raw completion prompt is off \
by one token — see `should_add_bos_token` against llama.cpp's per-pre arms: \
{bad_policy:?}"
);
}
}