push artifacts/scripts/eval_seq.py

artifacts/scripts/eval_seq.py

@@ -0,0 +1,120 @@
+"""Sequential-batch eval harness for a single-GPU box: batches of samples go teacher (concurrent
+vLLM) -> king reads -> challenger reads -> compare, keeping states files within the TTL and the
+shm budget. Same scoring as the eval service.
+
+    .venv/bin/python eval_seq.py --dataset pick20.jsonl --teacher DIR --king DIR --challenger DIR
+        [--samples N] [--batch 8] [--workers 8] [--out results.jsonl]
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from eval_service import dataset, metrics, run, servers
+from eval_service.run import Distance, _render, generate, load_tokenizer
+
+
+def aggregate(sums: list, role: str) -> Distance:
+    rows = [by_role[role] for by_role in sums if role in by_role]
+    if not rows:
+        return Distance(kl=float("inf"), dist=float("inf"))
+    return Distance(
+        kl=sum(s.kl for s in rows) / sum(s.predictions for s in rows),
+        dist=sum(s.layer_mean_dist for s in rows) / max(sum(s.activations for s in rows), 1),
+    )
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", type=Path, required=True)
+    ap.add_argument("--teacher", type=Path, required=True)
+    ap.add_argument("--king", type=Path, required=True)
+    ap.add_argument("--challenger", type=Path, required=True)
+    ap.add_argument("--samples", type=int, default=20)
+    ap.add_argument("--batch", type=int, default=8)
+    ap.add_argument("--workers", type=int, default=8)
+    ap.add_argument("--out", type=Path, default=Path("eval_seq.jsonl"))
+    args = ap.parse_args()
+
+    samples = dataset.load(args.dataset).samples[: args.samples]
+    tokenizer = load_tokenizer(args.teacher)
+    models = {"teacher": args.teacher, "king": args.king, "challenger": args.challenger}
+    device = "cuda:0"
+    heads = metrics.heads(models, [device])
+    all_sums = []
+    records_out = []
+
+    try:
+        for b0 in range(0, len(samples), args.batch):
+            batch = samples[b0 : b0 + args.batch]
+            records = []
+            groups = servers.serve({"teacher": args.teacher})
+            with ThreadPoolExecutor(args.workers) as pool:
+                def one(i, sample, prompt):
+                    answer, tpath = generate(groups["teacher"], prompt, i, sample.id)
+                    for _ in range(120):  # the connector writes states asynchronously
+                        if Path(tpath).exists():
+                            break
+                        time.sleep(0.5)
+                    keep = Path("/root/states") / Path(tpath).name
+                    keep.write_bytes(Path(tpath).read_bytes())
+                    Path(str(keep) + ".lock").touch()
+                    return answer, str(keep)
+
+                futures = {i: (s, p, pool.submit(one, b0 + i, s, p)) for i, (s, p) in enumerate((sample, _render(tokenizer, sample.messages, sample.tools)) for sample in batch)}
+                for i, (sample, prompt, fut) in futures.items():
+                    answer, tpath = fut.result()
+                    records.append({"id": sample.id, "prompt": prompt, "answer": answer, "tpath": tpath})
+                    print(f"teacher {b0+i+1}/{len(samples)}: {sample.id} p{len(prompt)} a{len(answer)}", flush=True)
+            servers.release("teacher")
+
+            for role, dir_ in (("king", args.king), ("challenger", args.challenger)):
+                groups = servers.serve({role: dir_})
+                group = groups[role]
+                for i, rec in enumerate(records):
+                    read, path = group.complete(rec["prompt"], len(rec["answer"]), b0 + i, rec["prompt"] + rec["answer"])
+                    if read != rec["answer"]:
+                        raise SystemExit(f"{role} left the teacher's answer on {rec['id']}")
+                    rec[role] = path
+                    print(f"{role} {b0+i+1}/{len(samples)}: {rec['id']}", flush=True)
+                servers.release(role)
+
+            for rec in records:
+                read = rec["prompt"] + rec["answer"][:-1]
+                states = {
+                    role: metrics.load_states(rec[role] if role != "teacher" else rec["tpath"], read, len(rec["prompt"]) - 1)
+                    for role in ("teacher", "king", "challenger")
+                }
+                by_role = metrics.compare(states, heads[device], rec["answer"], device)
+                all_sums.append(by_role)
+                records_out.append({"id": rec["id"], **{role: s.summary() for role, s in by_role.items()}})
+                print(
+                    f"scored {rec['id']}: king KL {by_role['king'].kl / by_role['king'].predictions:.4f}"
+                    f" challenger KL {by_role['challenger'].kl / by_role['challenger'].predictions:.4f}",
+                    flush=True,
+                )
+                for p in (rec["tpath"], rec["king"], rec["challenger"]):
+                    servers.discard(p)
+
+        agg = {role: aggregate(all_sums, role) for role in ("teacher", "king", "challenger")}
+        for role, d in agg.items():
+            print(f"{role:12s} KL {d.kl:.6f}  distance {d.dist:.5f}")
+        if agg["king"].kl < float("inf"):
+            need = Distance(kl=0.015 * agg["king"].kl, dist=0.0084 * agg["king"].dist)
+            print(
+                f"margin needed: KL {need.kl:.6f}  dist {need.dist:.6f} ;"
+                f" challenger leads by KL {agg['king'].kl - agg['challenger'].kl:.6f}"
+                f" ({100 * (agg['king'].kl - agg['challenger'].kl) / agg['king'].kl:.3f}%)"
+                f" dist {agg['king'].dist - agg['challenger'].dist:.6f}"
+                f" ({100 * (agg['king'].dist - agg['challenger'].dist) / agg['king'].dist:.3f}%)"
+            )
+        args.out.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records_out))
+    finally:
+        servers.stop_all()
+
+
+if __name__ == "__main__":
+    main()