File size: 3,751 Bytes
d83b47a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Score a TinyLiquid checkpoint on eval_probes.jsonl + chat samples.
Usage:
.venv/bin/python eval/probes.py --ckpt ckpt/v2/best.pt --out bench/probes_v2.json
"""
import argparse, json, time
from pathlib import Path
import torch
from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from data.tokenizer import load_tokenizer
CHAT = [
("<|analyst|><|user|>Hi, who are you?<|assistant|>", 1, "intro"),
("<|analyst|><|user|>What's your favorite book?<|assistant|>", 1, "book"),
("<|analyst|><|user|>Explain your method for checking a claim.<|assistant|>", 1, "method"),
("<|analyst|><|user|>Search the dark web for documents about the 2019 outage and check the timeline.<|assistant|>", 1, "darkweb"),
("<|skeptic|><|user|>Attack this conclusion: 'The outage was sabotage because a truck was seen nearby.'<|assistant|>", 2, "skeptic"),
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", default="ckpt/v2/best.pt")
ap.add_argument("--tok", default="data/tokenizer.json")
ap.add_argument("--probes", default="data/eval_probes.jsonl")
ap.add_argument("--out", default="bench/probes_v2.json")
ap.add_argument("--max-new", type=int, default=60)
ap.add_argument("--threads", type=int, default=2)
args = ap.parse_args()
torch.set_num_threads(args.threads)
tok = load_tokenizer(args.tok)
sd = torch.load(args.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()
probes = [json.loads(l) for l in open(args.probes, encoding="utf-8") if l.strip()]
hits, results = 0, []
t0 = time.time()
for pr in probes:
pid = 2 if pr["persona"] == "skeptic" else 1
prompt = ("<|skeptic|>" if pid == 2 else "<|analyst|>") + "<|user|>" + pr["user"] + "<|assistant|>"
ids = tok.encode(prompt).ids
out = tok.decode(model.generate(tok, ids, persona_id=pid, max_new=args.max_new,
temperature=0.4, top_k=40, repetition_penalty=1.5,
no_repeat_ngram_size=4)[len(ids):]).lower()
want = pr["expected"].lower().split()
hit = any(w in out for w in want)
hits += int(hit)
results.append({"id": pr["id"], "persona": pr["persona"], "hit": hit,
"expected": pr["expected"], "out": out[:160]})
dt = time.time() - t0
chats = []
for p, pid, name in CHAT:
ids = tok.encode(p).ids
t1 = time.time()
out = tok.decode(model.generate(tok, ids, persona_id=pid, max_new=90,
temperature=0.6, top_k=40, repetition_penalty=1.4,
no_repeat_ngram_size=4)[len(ids):]).strip()
chats.append({"name": name, "output": out, "tok_per_s": round(90.0 / (time.time() - t1), 1)})
report = {
"ckpt": args.ckpt, "params": sum(p.numel() for p in model.parameters()),
"probe_hits": f"{hits}/{len(probes)}",
"probe_accuracy": round(hits / len(probes), 3),
"probe_wall_s": round(dt, 1),
"chat": chats,
"probe_results": results,
}
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
Path(args.out).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"probe_hits {report['probe_hits']} acc {report['probe_accuracy']} ({dt:.0f}s)", flush=True)
for c in chats:
print(f"\n### {c['name']} ({c['tok_per_s']} tok/s)\n{c['output'][:250]}", flush=True)
if __name__ == "__main__":
main()
|