| |
| """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)) |
| from capu import Capu, Tokenizer |
| from common import HERE, PUNCTUATION |
|
|
| 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"] |
| mismatches = [] |
| for word in words: |
| expected = theirs.encode([word], is_pretokenized=True, add_special_tokens=False).ids |
| if word == "$START": |
| expected = [ours.start_id] |
| 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]} |
| 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() |
|
|
