checkngram.py

"""Assert predict/ngram.js and .build/ngram.py are the same model.

Three layers, because each can break independently:

  1. the PRNG stream — a seeded generator is worthless if the two languages
     disagree on the 32-bit arithmetic, and JS coercions vs Python masking is
     exactly where that goes wrong
  2. the distributions, before and after temperature / top-k / top-p
  3. whole generated sequences, which will diverge on the first mismatch of
     either of the above

Probabilities are compared with a tolerance, because pow() may differ in the
last bit between implementations. Chosen tokens are compared exactly: if a
last-bit difference ever flips a choice, that is worth knowing about.
"""

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 ngram  # noqa: E402
from corpora import CORPORA, TALE  # noqa: E402

TOL = 1e-12

SETTINGS = [
    {"temp": 1.0, "top_k": 0, "top_p": 0.0},
    {"temp": 0.0, "top_k": 0, "top_p": 0.0},     # greedy
    {"temp": 0.5, "top_k": 0, "top_p": 0.0},
    {"temp": 2.5, "top_k": 0, "top_p": 0.0},
    {"temp": 1.0, "top_k": 3, "top_p": 0.0},
    {"temp": 1.0, "top_k": 0, "top_p": 0.5},
    {"temp": 1.7, "top_k": 5, "top_p": 0.9},
]


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, "predict", "ngram.js"), encoding="utf-8").read())
    return ctx


def close(a, b):
    return abs(a - b) <= TOL


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

    # 1. PRNG streams must be bit-identical.
    for seed in [0, 1, 7, 12345, 2 ** 31, 4294967295]:
        py = [ngram.Rng(seed).next() for _ in range(1)]
        r = ngram.Rng(seed)
        py = [r.next() for _ in range(2000)]
        ctx.set("_seed", seed)
        js = json.loads(ctx.eval(
            "(function(){var r=new NGram.Rng(_seed),o=[];"
            "for(var i=0;i<2000;i++)o.push(r.next());return JSON.stringify(o)})()"))
        bad = [i for i, (a, b) in enumerate(zip(py, js)) if a != b]
        if bad:
            failures.append(f"PRNG seed {seed}: differs at draw {bad[0]} "
                            f"({py[bad[0]]} vs {js[bad[0]]})")
    if not failures:
        print("PASS  PRNG — 6 seeds x 2000 draws, bit-identical")

    corpora = [("tale", TALE, 2)] + [(n, t, 2) for n, t, _ in CORPORA]

    for name, text, order in corpora:
        for o in (order, 1, 0):
            py_model = ngram.train(text, o)
            ctx.set("_t", text)
            ctx.set("_o", o)
            ctx.eval("var M = NGram.train(_t, _o);")

            if json.loads(ctx.eval("JSON.stringify(M.vocab)")) != py_model["vocab"]:
                failures.append(f"{name}/order{o}: vocabulary differs")
                continue

            # 2. distributions, raw and shaped
            probes = [py_model["tokens"][:i] for i in (0, 1, 2, 3)]
            for probe in probes:
                py_items, py_order = ngram.distribution(py_model, probe)
                ctx.set("_c", json.dumps(probe))
                ctx.eval("var D = NGram.distribution(M, JSON.parse(_c));")
                js_items = json.loads(ctx.eval("JSON.stringify(D.items)"))
                js_order = int(ctx.eval("D.order"))

                if js_order != py_order:
                    failures.append(f"{name}/order{o}: backoff order "
                                    f"{js_order} vs {py_order}")
                    break
                if [t for t, _ in py_items] != [t for t, _ in js_items]:
                    failures.append(f"{name}/order{o}: distribution token order")
                    break
                if any(not close(a[1], b[1]) for a, b in zip(py_items, js_items)):
                    failures.append(f"{name}/order{o}: probabilities differ")
                    break

                for s in SETTINGS:
                    py_shaped = ngram.shape(py_items, s["temp"], s["top_k"],
                                            s["top_p"])
                    ctx.set("_temp", s["temp"])
                    ctx.set("_k", s["top_k"])
                    ctx.set("_p", s["top_p"])
                    js_shaped = json.loads(ctx.eval(
                        "JSON.stringify(NGram.shape(D.items, _temp, _k, _p))"))
                    if [t for t, _ in py_shaped] != [t for t, _ in js_shaped]:
                        failures.append(f"{name}/order{o} {s}: shaped tokens differ")
                        break
                    if any(not close(a[1], b[1])
                           for a, b in zip(py_shaped, js_shaped)):
                        failures.append(f"{name}/order{o} {s}: shaped probs differ")
                        break

            # 3. whole generated sequences
            for s in SETTINGS:
                for seed in (3, 7, 99):
                    py_gen = ngram.generate(py_model, "", 40, s["temp"],
                                            s["top_k"], s["top_p"], seed)
                    ctx.set("_temp", s["temp"])
                    ctx.set("_k", s["top_k"])
                    ctx.set("_p", s["top_p"])
                    ctx.set("_seed", seed)
                    js_text = ctx.eval(
                        "NGram.generate(M, '', 40, _temp, _k, _p, _seed).text")
                    if js_text != py_gen["text"]:
                        failures.append(
                            f"{name}/order{o} {s} seed {seed}: generated text "
                            f"diverges\n      js {js_text[:60]!r}\n      py "
                            f"{py_gen['text'][:60]!r}")
                        break

        print(f"PASS  {name:<14} orders {order}/1/0, {len(SETTINGS)} settings, "
              f"3 seeds")

    print()
    if failures:
        print(f"{len(failures)} disagreement(s):")
        for f in failures[:8]:
            print(f"  - {f}")
        return 1
    print(f"Browser sampler and Python reference agree on {len(corpora)} corpora.")
    return 0


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