verify.py

"""Prove this package's cached decode matches a full recompute.

Greedy is deterministic, so cached and uncached decoding MUST produce identical
tokens. This is the test that caught the bug these packages fix.

The one-step check passes the SAME engram context the real decode path uses
(package.json: engram_kwarg). Omitting it measures a gap the generator does not
actually have -- an earlier version of this script did exactly that and reported a
spurious 1.3e-03.

    python verify.py
"""
import os, sys
os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
os.environ.setdefault("USE_TF", "0")
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from generate import load, generate, MANIFEST     # noqa: E402


@torch.no_grad()
def main():
    torch.set_num_threads(4)
    model, tok = load("cpu")
    ids = tok.encode("The capital of France is Paris, and the city is known for")
    ids = (ids.tolist() if hasattr(ids, "tolist") else list(ids))[:20]

    def reset():
        try:
            from model_v2 import reset_memory_cache
            reset_memory_cache(model)
        except Exception:
            pass

    cfg = model.config
    kw = {}
    if MANIFEST.get("engram_kwarg") and getattr(cfg, "use_engram", False):
        n = max(1, int(getattr(cfg, "engram_max_ngram", 3)) - 1)
        kw["engram_context_ids"] = torch.tensor([ids[-(n + 1):-1]])

    reset()
    full = model(torch.tensor([ids])).logits[0, -1].float()
    reset()
    out = model(torch.tensor([ids[:-1]]), use_cache=True)
    step = model(torch.tensor([[ids[-1]]]),
                 past_key_values=out.past_key_values, use_cache=True, **kw)
    d = float((full - step.logits[0, -1].float()).abs().max())

    p = "The Industrial Revolution began in Britain and transformed"
    a = generate(model, tok, p, max_new=32, temp=0.0, use_cache=True)
    b = generate(model, tok, p, max_new=32, temp=0.0, use_cache=False)

    print(f"{MANIFEST['name']}")
    print(f"  engram kwarg used in one-step test : {bool(kw)}")
    print(f"  one-step max|dlogit| : {d:.4e}   (float32 floor is ~1e-5)")
    print(f"  32-tok greedy identical cached vs uncached : {a == b}")
    ok = d < 1e-3 and a == b
    print("  RESULT:", "PASS - cache path is exact" if ok else "FAIL - cached path diverges")
    if not ok:
        print(f"    cached : {a[:90]!r}")
        print(f"    nocache: {b[:90]!r}")
    return 0 if ok else 1


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