#!/usr/bin/env python3 """Checks the exported repo against the original PyTorch pipeline and against the FLEURS ground truth. 1. Tokenizer: capu.Tokenizer (the algorithm the Kotlin and Swift ports implement) against tokenizers' Unigram on every word of the test texts, as given and as corrected. 2. Pipeline: capu.Capu with the fp32 and the int8 model against out/reference.json (reference.py): share of texts with exactly the original's output. 3. Quality: punctuation (. , : ?) and capitalization F1 against the original FLEURS transcriptions, for the original, fp32 and int8. Writes out/validation.json (read by publish.py), and fixtures for the Kotlin and Swift checks: out/fixtures/cases.tsv (text, original, fp32 and int8 output) and out/fixtures/predictions.tsv (every int8 model call: token ids, word offsets, label ids). Fails when int8 drifts too far from the original. Usage: .venv/bin/python -I validate.py [--min-int8-f1-ratio 0.97] """ import argparse import difflib import json import re import sys import time import unicodedata from pathlib import Path from tokenizers import Tokenizer as HfTokenizer sys.path.insert(0, str(Path(__file__).resolve().parent)) # -I drops the script's own directory from capu import Capu, Tokenizer # noqa: E402 from common import HERE, PUNCTUATION # noqa: E402 REPO = HERE / "out" / "repo" def check_tokenizer(rows: list[dict]) -> dict: ours = Tokenizer(REPO / "unigram.tsv") theirs = HfTokenizer.from_file(str(REPO / "tokenizer.json")) words = sorted({w for r in rows for key in ("text", "truth", "original") for w in r[key].split()}) words += ["$START", "😀", "abc😀😀def", "fi", "Ｖｉệｔ", "x²", "́a"] # unk, NFKC, combining mark mismatches = [] for word in words: expected = theirs.encode([word], is_pretokenized=True, add_special_tokens=False).ids if word == "$START": expected = [ours.start_id] # an added token, matched before the unigram model got = [ours.start_id] else: got = ours.encode_word(word) if got != expected: mismatches.append((word, got, expected)) for word, got, expected in mismatches[:10]: print(f" tokenizer mismatch {word!r}: {got} != {expected}") print(f"tokenizer: {len(words) - len(mismatches)}/{len(words)} words identical") return {"words": len(words), "mismatches": len(mismatches)} def word_labels(text: str) -> list[tuple[str, bool, str]]: """(lower-cased word, capitalized, following . , : ? or "") for each word of a text.""" text = unicodedata.normalize("NFC", text) labels = [] for word, after in re.findall(r"(\w+)([^\w\s]*)", text): mark = next((c for c in after if c in PUNCTUATION), "") labels.append((word.lower(), word[0].isupper(), mark)) return labels def scores(outputs: list[str], truths: list[str]) -> dict: """Micro F1 of punctuation marks and of capitalized words, over the words output and truth have in common.""" counts = {"punct": [0, 0, 0], "upper": [0, 0, 0]} # true positives, predicted, actual for output, truth in zip(outputs, truths): out, ref = word_labels(output), word_labels(truth) matcher = difflib.SequenceMatcher(None, [w for w, _, _ in out], [w for w, _, _ in ref], autojunk=False) for i, j, size in matcher.get_matching_blocks(): for (_, o_up, o_mark), (_, r_up, r_mark) in zip(out[i:i + size], ref[j:j + size]): for key, o, r in (("punct", o_mark, r_mark), ("upper", o_up, r_up)): counts[key][0] += bool(o) and o == r counts[key][1] += bool(o) counts[key][2] += bool(r) result = {} for key, (tp, predicted, actual) in counts.items(): p, r = tp / max(predicted, 1), tp / max(actual, 1) result[f"{key}_f1"] = round(2 * p * r / max(p + r, 1e-9), 4) return result class Recorder: """Wraps an InferenceSession and keeps every call's inputs and argmax labels.""" def __init__(self, session): self.session, self.calls = session, {} def run(self, names, feed): result = self.session.run(names, feed) key = (",".join(map(str, feed["input_ids"][0])), ",".join(map(str, feed["input_offsets"][0]))) self.calls[key] = ",".join(map(str, result[0][0].argmax(-1))) return result def write_fixtures(rows: list[dict], outputs: dict[str, list[str]], calls: dict) -> None: fixtures = HERE / "out" / "fixtures" fixtures.mkdir(exist_ok=True) lines = ["\t".join((r["text"], r["original"], fp32, int8)) for r, fp32, int8 in zip(rows, outputs["fp32"], outputs["int8"])] (fixtures / "cases.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8") lines = ["\t".join((ids, offsets, labels)) for (ids, offsets), labels in calls.items()] (fixtures / "predictions.tsv").write_text("\n".join(lines) + "\n", encoding="utf-8") print(f"wrote {fixtures} ({len(rows)} cases, {len(calls)} model calls)") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--min-int8-f1-ratio", type=float, default=0.97) args = parser.parse_args() rows = json.loads((HERE / "out" / "reference.json").read_text(encoding="utf-8")) graded = [r for r in rows if r["truth"]] long = sum(len(r["text"].split()) > 48 for r in rows) print(f"{len(rows)} texts, {long} longer than one 48-word chunk") report = {"texts": len(rows), "long_texts": long, "tokenizer": check_tokenizer(rows)} report["original"] = scores([r["original"] for r in graded], [r["truth"] for r in graded]) print(f"original: {report['original']}") all_outputs = {} for quantized in (False, True): label = "int8" if quantized else "fp32" capu = Capu(REPO, quantized=quantized, threads=8) start = time.time() outputs = [capu(r["text"]) for r in rows] seconds = time.time() - start all_outputs[label] = outputs same = [o == r["original"] for o, r in zip(outputs, rows)] diffs = [(o, r["original"]) for o, r in zip(outputs, rows) if o != r["original"]] report[label] = { "identical": sum(same), "identical_long": sum(s for s, r in zip(same, rows) if len(r["text"].split()) > 48), **scores([o for o, r in zip(outputs, rows) if r["truth"]], [r["truth"] for r in graded]), "seconds": round(seconds, 1), } print(f"{label}: {report[label]}") for got, want in diffs[:3]: print(f" got: {got}\n want: {want}") (HERE / "out" / "validation.json").write_text(json.dumps(report, indent=1) + "\n") capu.session = Recorder(capu.session) for r in rows: capu(r["text"]) write_fixtures(rows, all_outputs, capu.session.calls) failed = [] if report["tokenizer"]["mismatches"]: failed.append("tokenizer differs from tokenizers' Unigram") if report["fp32"]["identical"] < len(rows) * 0.99: failed.append("fp32 output differs from the original on more than 1% of texts") for key in ("punct_f1", "upper_f1"): if report["int8"][key] < report["original"][key] * args.min_int8_f1_ratio: failed.append(f"int8 {key} {report['int8'][key]} too far below the original's {report['original'][key]}") if failed: sys.exit("error: " + "; ".join(failed)) if __name__ == "__main__": main()