"""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()