Buckets:
| """Predictions and metrics. | |
| python evaluate.py predict --readout base --reason base --latency base | |
| python evaluate.py predict --model checkpoints/joint_merged --readout joint --reason joint --latency joint | |
| python evaluate.py metrics | |
| Trained models are evaluated from merged bf16 checkpoints: vLLM 0.32 silently ignores LoRA adapters on Qwen3.5 | |
| (checked: its outputs with an adapter equal the base model's), so adapters are never served as LoRA here. | |
| predict writes results/predictions/readout_<name>.jsonl (letter log-probs, dataset option order) for eval and | |
| calibration items, and reason_<name>.jsonl (one thinking sample per item, its parsed letter, the readout after | |
| its chain, tokens). Reasoning streams to disk and resumes. Every eval file holds one prediction per item. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import json | |
| import os | |
| import time | |
| import numpy as np | |
| from transformers import AutoTokenizer | |
| from common import read_jsonl, split_chain, write_jsonl | |
| PRED = "results/predictions" | |
| def items_all(): | |
| ev, cal = read_jsonl("data/eval.jsonl"), read_jsonl("data/calib.jsonl") | |
| return ev + cal | |
| async def predict(args): | |
| from vllm_utils import latency_async, make_async_engine, readout_async, think_async | |
| os.makedirs(PRED, exist_ok=True) | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| engine = make_async_engine(args.model, gpu_mem=args.gpu_mem) | |
| lora_of = lambda name: None # names only label output files | |
| items = items_all() | |
| if args.limit: | |
| items = items[:args.limit] | |
| orders = [list(range(len(x["options"]))) for x in items] | |
| for name in filter(None, args.readout.split(",")): | |
| path = f"{PRED}/readout_{name}.jsonl" | |
| if os.path.exists(path) and len(read_jsonl(path)) == len(items): | |
| print(f"readout {name}: exists", flush=True) | |
| continue | |
| t0 = time.time() | |
| lps = await readout_async(engine, tok, items, orders, lora=lora_of(name)) | |
| write_jsonl(path, [{"id": it["id"], "set": it["set"], "lp": [float(v) for v in lp]} for it, lp in zip(items, lps)]) | |
| print(f"readout {name}: {len(items)} items in {time.time() - t0:.0f}s", flush=True) | |
| for name in filter(None, args.reason.split(",")): | |
| raw = f"{PRED}/reason_{name}.raw.jsonl" | |
| done = {r["id"] for r in read_jsonl(raw)} if os.path.exists(raw) else set() | |
| jobs = [(it, o) for it, o in zip(items, orders) if it["id"] not in done] | |
| t0 = time.time() | |
| stats = await think_async(engine, tok, jobs, raw, n=1, deadline=args.deadline or float("inf"), | |
| concurrency=args.concurrency, lora=lora_of(name), post_readout=True) | |
| print(f"reason {name}: {json.dumps(stats)}", flush=True) | |
| recs = {r["id"]: r for r in read_jsonl(raw)} | |
| by_id = {it["id"]: it for it in items} | |
| # readout after the chain: taken right after generation when available, else one prefill pass here | |
| closed = [(by_id[i], r) for i, r in recs.items() if r["samples"][0]["closed"] and r["samples"][0].get("post_lp") is None] | |
| chains = [split_chain(r["samples"][0]["text"])[0] for _, r in closed] | |
| post = await readout_async(engine, tok, [it for it, _ in closed], [r["order"] for _, r in closed], | |
| lora=lora_of(name), think_chains=chains) | |
| post_by = {it["id"]: lp for (it, _), lp in zip(closed, post)} | |
| post_by.update({i: r["samples"][0]["post_lp"] for i, r in recs.items() if r["samples"][0].get("post_lp") is not None}) | |
| print(f"reason {name}: post-chain readout {len(post_by) - len(closed)} taken at generation, {len(closed)} by prefill", flush=True) | |
| rows = [] | |
| for it in items: | |
| r = recs.get(it["id"]) | |
| if r is None: | |
| continue | |
| s = r["samples"][0] | |
| lp = post_by.get(it["id"]) | |
| rows.append({"id": it["id"], "set": it["set"], "answer": s["answer"], "closed": s["closed"], | |
| "looped": s["looped"], "tokens": s["tokens"], "finish": s["finish"], | |
| "post_lp": None if lp is None else [float(v) for v in lp], "text": s["text"]}) | |
| write_jsonl(f"{PRED}/reason_{name}.jsonl", rows) | |
| print(f"reason {name}: {len(rows)} items, post-readout in {time.time() - t0:.0f}s total", flush=True) | |
| for name in filter(None, args.latency.split(",")): | |
| ev = [x for x in items if x["set"] != "calib_proc" and x["set"] != "calib_ts"] | |
| rng = np.random.default_rng(17) | |
| pick = [ev[i] for i in rng.choice(len(ev), size=min(100, len(ev)), replace=False)] | |
| porders = [list(range(len(x["options"]))) for x in pick] | |
| await latency_async(engine, tok, pick[:5], porders[:5], think=False, lora=lora_of(name)) # warm-up | |
| r_lat = await latency_async(engine, tok, pick, porders, think=False, lora=lora_of(name)) | |
| t_lat = await latency_async(engine, tok, pick[:20], porders[:20], think=True, lora=lora_of(name)) | |
| out = {"readout_ms": {"median": 1000 * float(np.median(r_lat)), "p90": 1000 * float(np.percentile(r_lat, 90)), "n": len(r_lat)}, | |
| "reasoning_ms": {"median": 1000 * float(np.median(t_lat)), "p90": 1000 * float(np.percentile(t_lat, 90)), "n": len(t_lat)}} | |
| with open(f"{PRED}/latency_{name}.json", "w") as f: | |
| json.dump(out, f, indent=1) | |
| print(f"latency {name}: {json.dumps(out)}", flush=True) | |
| engine.shutdown() | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| sub = ap.add_subparsers(dest="cmd", required=True) | |
| p = sub.add_parser("predict") | |
| p.add_argument("--model", default="Qwen/Qwen3.5-4B") | |
| p.add_argument("--readout", default="") | |
| p.add_argument("--reason", default="") | |
| p.add_argument("--latency", default="") | |
| p.add_argument("--deadline", type=float, default=None) | |
| p.add_argument("--concurrency", type=int, default=256) | |
| p.add_argument("--gpu-mem", type=float, default=0.90) | |
| p.add_argument("--limit", type=int, default=None) | |
| m = sub.add_parser("metrics") | |
| m.add_argument("--out", default="results/metrics.json") | |
| args = ap.parse_args() | |
| if args.cmd == "predict": | |
| asyncio.run(predict(args)) | |
| else: | |
| from metrics import compute_all | |
| compute_all(args.out) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.38 kB
- Xet hash:
- d3eed0a91019d4670c669e0b340d163eebca49732971f5982362f6f406cb80d7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.