#!/usr/bin/env python3 """avg@k eval against an OpenAI-compatible vLLM endpoint, for VibeThinker-3B and its controls (Qwen2.5-Coder-3B-Instruct, Qwen3-4B-Thinking-2507). Faithful to the paper's protocol: NO system prompt (the \\boxed instruction is in the question), temp=1.0, top_p=0.95, top_k=-1, k independent samples/item. Reports the same metrics the report uses: pass@1 = mean correctness over all k samples (the headline avg@k number) pass@k = any of the k correct cons@k = majority-vote (self-consistency) correctness Crucially logs finish_reason per sample so truncated traces (finish_reason= "length") are visible -- a truncated long CoT scores 0 and would otherwise read as a real miss (the false-zero trap). Usage (one resident model at a time): python3 vibe_eval.py --base http://localhost:8000 --model vibethinker-3b \ --tag vibe3b --n 16 python3 vibe_eval.py --compare results/vibe3b.jsonl results/base3b.jsonl """ import json, os, re, time, argparse, urllib.request, glob from collections import Counter from concurrent.futures import ThreadPoolExecutor HERE = os.path.dirname(os.path.abspath(__file__)) DATA = os.path.join(HERE, "data") RES = os.path.join(HERE, "results"); os.makedirs(RES, exist_ok=True) # ---------- io ---------- def load_items(cats=None): items = [] for f in sorted(glob.glob(os.path.join(DATA, "*.jsonl"))): cat = os.path.splitext(os.path.basename(f))[0] if cats and cat not in cats: continue with open(f) as fh: items += [json.loads(l) for l in fh if l.strip()] return items def call_one(base, model, key, question, max_tokens, temp, top_p, timeout, think=False): # ONE sample per request. llama.cpp's OpenAI endpoint ignores n>1, so we fire # k independent requests instead (portable across llama.cpp and vLLM). # No system prompt: identical input for every model. payload = {"model": model, "temperature": temp, "top_p": top_p, "max_tokens": max_tokens, "messages": [{"role": "user", "content": question}]} if think: # override server's enable_thinking:false payload["chat_template_kwargs"] = {"enable_thinking": True} body = json.dumps(payload).encode() req = urllib.request.Request(base.rstrip("/") + "/v1/chat/completions", data=body, headers={"Content-Type": "application/json", **({"Authorization": f"Bearer {key}"} if key else {})}) with urllib.request.urlopen(req, timeout=timeout) as r: d = json.load(r) ch = d["choices"][0] msg = ch.get("message", {}) text = msg.get("content") or msg.get("reasoning_content") or "" return {"text": text, "finish": ch.get("finish_reason")}, d.get("model") # ---------- grading ---------- def extract(text): """Return (pred, from_box). from_box=True means it came from \\boxed{}, which lets math grading apply a STRICT integer parse (reject fractions).""" b = list(re.finditer(r"\\boxed\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}", text)) if b: return b[-1].group(1).strip(), True m = list(re.finditer(r"ANSWER:\s*(.+)", text, re.I)) if m: return m[-1].group(1).strip(), False return (text.strip().splitlines()[-1].strip() if text.strip() else ""), False def strict_int(s): # boxed content must BE an integer s = s.replace(",", "").replace("$", "").replace(" ", "").strip() m = re.fullmatch(r"-?\d+", s) or re.fullmatch(r"(-?\d+)\.0*", s) return int(m.group(1) if m and m.lastindex else s) if m else None def loose_int(s): # only for the no-box fallback path s = s.replace(",", "").replace("$", "").replace("\\", "").strip() m = re.search(r"-?\d+", s) return int(m.group(0)) if m else None try: # symbolic LaTeX equivalence (HMMT non-int answers) from math_verify import parse as _mv_parse, verify as _mv_verify _HAS_MV = True except Exception: _HAS_MV = False def math_equiv(pred, gold): """Symbolic equivalence for non-integer golds (fractions, radicals, sets). e.g. \\dfrac{1}{576}==\\frac{1}{576}, 8\\sqrt{10}==\\sqrt{640}. Falls back to exact-string match if math_verify is unavailable or a parse throws.""" if not _HAS_MV: return pred.strip() == gold.strip() try: return bool(_mv_verify(_mv_parse(f"${gold}$"), _mv_parse(f"${pred}$"))) except Exception: return pred.strip() == gold.strip() def grade(category, raw, gold): pred, from_box = extract(raw) if category == "gpqa": # multiple-choice A-D m = re.search(r"\b([A-D])\b", pred.upper()) or re.search(r"([A-D])", pred.upper()) return bool(m) and m.group(1) == gold.strip().upper(), pred g = strict_int(gold) # aime / hmmt -> integer answers p = strict_int(pred) if from_box else loose_int(pred) if g is not None and p is not None: return p == g, pred if g is None: # non-integer gold (hmmt fractions/radicals) return math_equiv(pred, gold), pred return pred.strip() == gold.strip(), pred # ---------- run ---------- def run(args): items = load_items(set(args.cats.split(",")) if args.cats else None) if args.limit: items = items[:args.limit] out = os.path.join(RES, f"{args.tag}.jsonl") print(f"[{args.tag}] {len(items)} items x avg@{args.n} -> {args.model} @ {args.base}") # Flatten to (item, sample) tasks -> k independent requests/item (portable). tasks = [(it, s) for it in items for s in range(args.n)] print(f"[{args.tag}] {len(items)} items x avg@{args.n} = {len(tasks)} gens " f"-> {args.model} @ {args.base}") def do_task(task): it, _ = task try: s, served = call_one(args.base, args.model, args.key, it["question"], args.max_tokens, args.temp, args.top_p, args.timeout, think=args.think) ok, pred = grade(it["category"], s["text"], it["answer"]) return {"id": it["id"], "ok": int(ok), "pred": pred, "trunc": int(s["finish"] == "length"), "served": served, "err": None} except Exception as e: return {"id": it["id"], "ok": 0, "pred": None, "trunc": 0, "served": None, "err": str(e)} by_id = {it["id"]: it for it in items} order = [it["id"] for it in items] agg = {it["id"]: {"ok": [], "pred": [], "trunc": 0, "err": 0, "served": set()} for it in items} def finalize(iid): it = by_id[iid]; a = agg[iid]; scores = a["ok"]; n = len(scores) preds = [p for p in a["pred"] if p] cons_pred = Counter(preds).most_common(1)[0][0] if preds else "" cons_ok = grade(it["category"], "\\boxed{%s}" % cons_pred, it["answer"])[0] if cons_pred else False return {"id": iid, "category": it["category"], "answer": it["answer"], "pass1": (sum(scores) / n if n else 0.0), "passk": int(any(scores)), "cons": int(cons_ok), "n": n, "trunc": a["trunc"], "err": a["err"], "pred": cons_pred, "preds": a["pred"], # persist model answer(s) for offline re-grading "served_model": sorted(a["served"])} rows, done, total_trunc = [], 0, 0 fout = open(out, "w") with ThreadPoolExecutor(max_workers=args.concurrency) as ex: for r in ex.map(do_task, tasks): # yields in task order -> grouped by item a = agg[r["id"]] a["ok"].append(r["ok"]); a["pred"].append(r["pred"]); a["trunc"] += r["trunc"] if r["err"]: a["err"] += 1 if r["served"]: a["served"].add(r["served"]) done += 1 if done % args.n == 0: # one item's k samples are all in row = finalize(order[done // args.n - 1]) rows.append(row); fout.write(json.dumps(row, ensure_ascii=False) + "\n"); fout.flush() total_trunc += row["trunc"] tr = f" *** TRUNC={row['trunc']} ***" if row["trunc"] else "" er = f" err={row['err']}" if row["err"] else "" print(f" [{done//args.n:>3}/{len(items)}] {row['pass1']*100:3.0f}% " f"{row['id']:<12} gold={row['answer'][:8]:<8}{er}{tr}", flush=True) fout.close() if total_trunc: print(f"\n!!! {total_trunc} truncated samples (finish_reason=length) — traces exceeded " f"the per-slot ctx. Drop llama-server to --parallel 1 and re-run. !!!") served = sorted({s for r in rows for s in r["served_model"]}) print(f"\nSERVED MODEL(S): {served}") summarize(rows, args.tag) print(f"saved -> {out}") def summarize(rows, tag): cats = {} for r in rows: c = cats.setdefault(r["category"], {"n": 0, "p1": 0.0, "pk": 0, "cons": 0, "tr": 0}) c["n"] += 1; c["p1"] += r["pass1"]; c["pk"] += r["passk"] c["cons"] += r["cons"]; c["tr"] += r.get("trunc", 0) print(f"\n=== {tag} ===") print(f"{'category':<10}{'pass@1':>9}{'pass@k':>9}{'cons@k':>9}{'trunc':>8}") for c, v in sorted(cats.items()): n = v["n"] print(f"{c:<10}{100*v['p1']/n:>8.1f}%{100*v['pk']/n:>8.0f}%{100*v['cons']/n:>8.0f}%{v['tr']:>8}") def compare(a, b): def load(p): return {r["id"]: r for r in (json.loads(l) for l in open(p))} A, B = load(a), load(b) ta, tb = os.path.basename(a).split(".")[0], os.path.basename(b).split(".")[0] cats = {} for i in A: if i not in B: continue c = cats.setdefault(A[i]["category"], [0, 0.0, 0.0]) c[0] += 1; c[1] += A[i]["pass1"]; c[2] += B[i]["pass1"] print(f"{'category':<10}{ta:>10}{tb:>10}{'delta':>9}") for c, (n, x, y) in sorted(cats.items()): print(f"{c:<10}{100*x/n:>9.1f}%{100*y/n:>9.1f}%{100*(x-y)/n:>+8.1f}") if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--base"); ap.add_argument("--model"); ap.add_argument("--key") ap.add_argument("--tag", default="run") ap.add_argument("--n", type=int, default=16, help="samples per item (avg@k)") ap.add_argument("--limit", type=int, help="cap items per run (e.g. subsample gpqa)") ap.add_argument("--cats", help="comma list e.g. aime,aime25,hmmt25 (default: all)") ap.add_argument("--max-tokens", type=int, default=40960, dest="max_tokens") ap.add_argument("--temp", type=float, default=1.0) ap.add_argument("--top-p", type=float, default=0.95, dest="top_p") ap.add_argument("--timeout", type=int, default=1800) ap.add_argument("--think", action="store_true", help="send chat_template_kwargs enable_thinking=true (override server default)") ap.add_argument("--concurrency", type=int, default=4, help="match llama-server --parallel (4); bump to 16+ for vLLM") ap.add_argument("--compare", nargs=2) a = ap.parse_args() if a.compare: compare(*a.compare) else: run(a)