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