eval/bench2.py

"""Resumable, low-memory benchmark for TinyLiquid.

Usage:
  # accumulate val batches (safe to re-run; resumes from state)
  .venv/bin/python eval/bench2.py --mode val   --max-iters 6
  .venv/bin/python eval/bench2.py --mode probe --idx 0      # one probe
  .venv/bin/python eval/bench2.py --mode speed
  .venv/bin/python eval/bench2.py --mode sample --idx 0     # one sample
  .venv/bin/python eval/bench2.py --mode finish             # write metrics.json
"""
import argparse, json, time
from pathlib import Path
import numpy as np, torch

from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from data.tokenizer import load_tokenizer

PROBES = [
    ("Discrepancy", "Two accounts describe the same event. Account A: 'The meeting started at 9am and ended at 11am.' Account B: 'The meeting started at 9am and ran until noon.' Find discrepancies.", "end time"),
    ("Contradiction", "Account A: 'No officials were present.' Account B: 'An official arrived later.' What can you conclude?", "not necessarily a contradiction"),
    ("Overclaim", "Evaluate: 'Crime doubled last year because of the new policy.' Incidents rose 1,000 to 2,000 while reporting methods changed.", "overclaim"),
    ("Pattern", "What are the weak links in a theory claiming one actor caused three unrelated disasters?", "no direct evidence"),
    ("Verification", "Verify: 'The building was built in 2015 and never renovated.' Assessor record: built 2015, renovation permit 2019.", "false"),
    ("Timeline", "Witness saw a van at 3:30pm; outage reported 3:40pm; inspection 6pm; no report before 5pm. Reconstruct.", "gap"),
    ("History", "2019: 'officials rule out sabotage.' 2021 filing: a second investigation opened the same week. What changed?", "parallel investigation hidden"),
    ("Politics", "Statement: 'We balanced the budget without raising taxes.' Budget includes a reassessment raising collections 9%.", "misleading"),
    ("Source chain", "A claim rests on: company blog, a wire story repeating it, an analyst note quoting the wire. Rate the evidence.", "single chain"),
]
SAMPLE_PROMPTS = [
    "<|analyst|><|user|>Verify: 'The bridge was painted in 2019 and never repainted.' Records show a 2022 repaint permit.<|assistant|>",
    "<|analyst|><|user|>What's the most common mistake you see in research?<|assistant|>",
    "<|skeptic|><|user|>Attack this conclusion: 'Three failures in one week with vans nearby is deliberate.'<|assistant|>",
]

def load(tok_path, ckpt, threads):
    torch.set_num_threads(threads)
    tok = load_tokenizer(tok_path)
    sd = torch.load(ckpt, map_location="cpu")
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                           **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"])
    model.eval()
    return tok, model

def read_state(p):
    if p.exists():
        return json.loads(p.read_text())
    return {"val_total": 0.0, "val_count": 0, "val_iters": 0,
            "probes": {}, "speed": None, "samples": []}

def write_state(p, st):
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(st), encoding="utf-8")

def val_batch(model, tok, arr_t, batch, seq, seed):
    n = (len(arr_t) - 1) // seq
    rng = np.random.RandomState(seed)
    s = int(rng.randint(0, n - batch))
    idx = torch.arange(s * seq, (s + batch) * seq, dtype=torch.long)
    buf = torch.stack([torch.from_numpy(arr_t[int(i): int(i) + seq].astype(np.int64)) for i in idx])
    x, y = buf[:, :-1], buf[:, 1:]
    logits = model(x)
    loss = torch.nn.functional.cross_entropy(
        logits.view(-1, logits.size(-1)), y.reshape(-1))
    return float(loss.item()) * y.numel(), y.numel()

@torch.no_grad()
def generate(model, tok, ids, persona_id, max_new, temperature, top_k,
             repetition_penalty, no_repeat_ngram_size):
    return model.generate(tok, ids, persona_id=persona_id, max_new=max_new,
                          temperature=temperature, top_k=top_k,
                          repetition_penalty=repetition_penalty,
                          no_repeat_ngram_size=no_repeat_ngram_size)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", required=True)
    ap.add_argument("--ckpt", default="ckpt/dpo/model_final.pt")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--val", default="data/valid.bin")
    ap.add_argument("--out", default="bench/metrics.json")
    ap.add_argument("--state", default="bench/state.json")
    ap.add_argument("--val-batches", type=int, default=16)
    ap.add_argument("--max-iters", type=int, default=6)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--seq", type=int, default=64)
    ap.add_argument("--idx", type=int, default=-1)
    ap.add_argument("--threads", type=int, default=1)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    tok, model = load(args.tok, args.ckpt, args.threads)
    print('model loaded', flush=True)
    st = read_state(Path(args.state))

    if args.mode == "val":
        mm = np.memmap(args.val, dtype=np.uint16, mode='r')
        print('val data loaded', flush=True)
        t = mm
        n_iters = min(args.max_iters, args.val_batches - st["val_iters"])
        for i in range(n_iters):
            seed = args.seed * 1000 + st["val_iters"]
            tot, cnt = val_batch(model, tok, t, args.batch, args.seq, seed)
            st["val_total"] += tot; st["val_count"] += cnt; st["val_iters"] += 1
            write_state(Path(args.state), st)
            print(f"val iter {st['val_iters']}/{args.val_batches} "
                  f"partial_ppl={np.exp(st['val_total']/st['val_count']):.3f} "
                  f"tokens={st['val_count']}", flush=True)
        print("val done", flush=True)

    elif args.mode == "probe":
        idx = int(args.idx)
        q = PROBES[idx][1]; want = PROBES[idx][2]
        prompt = "<|analyst|><|user|>" + q + "<|assistant|>"
        ids = tok.encode(prompt).ids
        out = tok.decode(generate(model, tok, ids, persona_id=1, max_new=60,
                                  temperature=0.4, top_k=40,
                                  repetition_penalty=1.5,
                                  no_repeat_ngram_size=4)[len(ids):]).lower()
        hit = any(w in out for w in want.split())
        st["probes"][str(idx)] = {"hit": hit, "out": out[:200]}
        write_state(Path(args.state), st)
        print(f"probe {idx} {PROBES[idx][0]}: hit={hit}", flush=True)

    elif args.mode == "speed":
        ids = tok.encode("<|analyst|><|user|>Evaluate this claim: 'X caused Y.'<|assistant|>").ids
        t0 = time.time()
        generate(model, tok, ids, persona_id=1, max_new=40, temperature=0.6,
                 top_k=40, repetition_penalty=1.4, no_repeat_ngram_size=4)
        dt = time.time() - t0
        st["speed"] = round(40.0 / dt, 1)
        write_state(Path(args.state), st)
        print(f"speed {st['speed']} tok/s", flush=True)

    elif args.mode == "sample":
        idx = int(args.idx)
        p = SAMPLE_PROMPTS[idx]
        persona = 2 if p.startswith("<|skeptic|>") else 1
        ids = tok.encode(p).ids
        out = tok.decode(generate(model, tok, ids, persona_id=persona,
                                  max_new=80, temperature=0.6, top_k=40,
                                  repetition_penalty=1.4,
                                  no_repeat_ngram_size=4)[len(ids):]).strip()
        st["samples"].append({"idx": idx,
                              "prompt": p.split("<|user|>")[1].split("<|assistant|>")[0],
                              "persona": "skeptic" if persona == 2 else "analyst",
                              "output": out})
        write_state(Path(args.state), st)
        print(f"sample {idx} done ({len(out)} chars)", flush=True)

    elif args.mode == "finish":
        assert st["val_iters"] >= args.val_batches, \
            f"val incomplete {st['val_iters']}/{args.val_batches}"
        ppl = np.exp(st["val_total"] / st["val_count"])
        hits = sum(v["hit"] for v in st["probes"].values())
        total = len(PROBES)
        st["samples"] = sorted(st["samples"], key=lambda s: s["idx"])
        metrics = {
            "checkpoint": args.ckpt,
            "params": 7788288,
            "val_loss": round(float(np.log(ppl)), 4),
            "val_ppl": round(float(ppl), 4),
            "val_tokens": int(st["val_count"]),
            "probe_hits": f"{hits}/{total}",
            "probe_accuracy": round(hits / total, 3),
            "gen_speed_tok_per_s": st["speed"],
            "hardware": "8-core ARM, no GPU",
            "samples": st["samples"],
        }
        out = Path(args.out)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8")
        print(json.dumps(metrics, indent=2, ensure_ascii=False), flush=True)
        print("metrics written", flush=True)

if __name__ == "__main__":
    main()