checkbpe.py

"""Assert the browser's BPE trainer and the Python reference agree exactly.

The page claims the algorithm you run in your browser is the one that produced
the static figures. That claim is only worth making if it is checked, so this
compares merge lists, step counts and encodings across a range of corpora —
including the ones designed to break naive implementations: astral-plane
characters, combining marks, and text with no repetition at all.
"""

import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "pylib"))
ROOT = os.path.abspath(os.path.join(HERE, ".."))

import quickjs  # noqa: E402

import bpe  # noqa: E402
from corpora import CORPORA  # noqa: E402

PROBES = [" strawberry", "strawberry", " raspberry", "strawberries", " kiwi",
          "", " ", "\n", "🍓🍓", "café", "the the the"]


def js_context():
    ctx = quickjs.Context()
    ctx.set_memory_limit(1 << 29)
    ctx.set_max_stack_size(1 << 22)
    ctx.eval("var globalThis = this;")
    ctx.eval(open(os.path.join(ROOT, "vocabulary", "bpe.js"), encoding="utf-8").read())
    return ctx


def main():
    ctx = js_context()
    failures = []

    for name, text, n_merges in CORPORA:
        want = bpe.train(text, n_merges)
        ctx.set("_t", text)
        ctx.set("_n", n_merges)
        got = json.loads(ctx.eval("JSON.stringify(BPE.train(_t, _n))"))

        w_merges = [list(m) for m in want["merges"]]
        if got["merges"] != w_merges:
            # Report the first divergence rather than dumping both lists.
            for i, (a, b) in enumerate(zip(got["merges"], w_merges)):
                if a != b:
                    failures.append(f"{name}: merge {i + 1} js={a} py={b}")
                    break
            else:
                failures.append(f"{name}: merge count js={len(got['merges'])} "
                                f"py={len(w_merges)}")
            continue

        if got["vocabSize"] != want["vocab_size"]:
            failures.append(f"{name}: vocab js={got['vocabSize']} "
                            f"py={want['vocab_size']}")

        if [s["count"] for s in got["steps"]] != [s["count"] for s in want["steps"]]:
            failures.append(f"{name}: step counts differ")

        for probe in PROBES:
            py = bpe.encode(probe, want["merges"])
            ctx.set("_p", probe)
            js = json.loads(ctx.eval(
                "JSON.stringify(BPE.encode(_p, BPE.train(_t, _n).merges))"))
            if js != py:
                failures.append(f"{name}: encode({probe!r}) js={js} py={py}")
                break

        print(f"PASS  {name:<22} {len(w_merges):>3} merges, "
              f"vocab {want['vocab_size']:>3}, {len(PROBES)} probes")

    # Pre-tokenization must agree too — it is the rule that welds on spaces.
    for name, text, _ in CORPORA:
        ctx.set("_t", text)
        js = json.loads(ctx.eval("JSON.stringify(BPE.pretokenize(_t))"))
        py = bpe.pretokenize(text)
        if js != py:
            failures.append(f"{name}: pretokenize differs "
                            f"({len(js)} vs {len(py)} chunks)")

    print()
    if failures:
        print(f"{len(failures)} disagreement(s) between trainer.js and bpe.py:")
        for f in failures[:10]:
            print(f"  - {f}")
        return 1
    print(f"Browser trainer and Python reference agree exactly on "
          f"{len(CORPORA)} corpora.")
    return 0


if __name__ == "__main__":
    sys.exit(main())