reproduce/01_benchmark.py

#!/usr/bin/env python3
"""Reproduce claim 01 (results/claim_01_benchmark.md) — vanilla-vs-Lite at param-fair.

NOTE: This is the documentation-only entry point. The actual val-bpc
benchmark requires:
  1. The FineWeb-Edu training pipeline (not bundled here).
  2. A clean 3-seed vanilla replication run (~$2.60 on an A40 SXM —
     queued, not run; we ran out of budget on RunPod first).

What you can verify FROM THE KIT alone is the architecture itself:
the same `TilelliLiteLM` class that produced the val-bpc numbers loads
cleanly from `checkpoints/tilelli_chat_v4.pt`, with 10.18 M parameters,
3-pathway routing, and FP32 weights. This script confirms that load
and prints the shape + param count so the architecture audit is
non-empty.

If you want the full vanilla-vs-Lite re-run, the training launchers live
in the private working repo. Reach out if you want them; the budget to
run them yourself is ~$15 of GPU community pricing.
"""
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

import torch
from tilelli.eval.metacog_probe import load_bridge


def main():
    ckpt_path = ROOT / "checkpoints" / "tilelli_chat_v4.pt"
    print(f"[reproduce] loading {ckpt_path.name}")
    model, _abstain, tok = load_bridge(str(ckpt_path))
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[reproduce] architecture: {type(model).__name__}")
    print(f"[reproduce]   params: {n_params:,}  ({n_params / 1e6:.2f} M)")
    print(f"[reproduce]   pathways: 3 (local conv k=5 + sparse top-k attention + dense FFN)")
    print(f"[reproduce]   weights: FP32 (the deployed v4 ckpt does not exercise the ternary path)")
    print(f"[reproduce]   max_seq_len: {getattr(model, 'max_seq_len', 'unknown')}")
    expected = 10_000_000
    tolerance = 0.05
    lo, hi = int(expected * (1 - tolerance)), int(expected * (1 + tolerance))
    if not (lo <= n_params <= hi):
        print(f"[reproduce] FAIL — param count {n_params} not within 5% of expected {expected}")
        sys.exit(1)
    print(f"[reproduce] PASS — architecture loads cleanly, within ±5% of 10M params")
    print()
    print("[reproduce] For the val-bpc vs vanilla number (0.5686 vs 0.5707):")
    print("           see results/claim_01_benchmark.md. That number was produced")
    print("           by training the same architecture from scratch on FineWeb-Edu.")
    print("           This kit ships an inference-only contract; the full")
    print("           train-from-scratch reproducer is not bundled.")


if __name__ == "__main__":
    main()