A 3B was supposed to rival models 9× larger. On math, it actually does.

#!/usr/bin/env python3
"""Convert VibeThinker repo math parquets -> flat jsonl the harness eats.

Source rows (verl format):
  prompt:        [{"role":"user","content": "<problem ... \\boxed{}>"}]
  reward_model:  {"ground_truth": "<answer>", "style": "rule"}

Output rows:
  {"id","category","question","answer"}

The prompt already carries the "output the final answer within \\boxed{}"
instruction, so we send it verbatim with NO extra system prompt -- identical
input for every model under test (the controlled-input invariant).
"""
import os, json, glob, argparse
import pandas as pd

DEFAULT_SRC = os.path.expanduser(
    "~/Documents/AI/VibeThinker/eval/math/data")  # adjust on the host

def convert(src, out):
    os.makedirs(out, exist_ok=True)
    total = 0
    for f in sorted(glob.glob(os.path.join(src, "*.parquet"))):
        cat = os.path.splitext(os.path.basename(f))[0]   # aime / aime25 / hmmt25 / gpqa
        df = pd.read_parquet(f)
        rows = []
        for i, r in df.iterrows():
            prompt = r["prompt"]
            # prompt is a numpy array / list of {role,content}
            content = prompt[0]["content"] if not isinstance(prompt, str) else prompt
            gt = r["reward_model"]["ground_truth"]
            rows.append({
                "id": f"{cat}-{i:03d}",
                "category": cat,
                "question": content,
                "answer": str(gt).strip(),
            })
        with open(os.path.join(out, f"{cat}.jsonl"), "w") as fh:
            for row in rows:
                fh.write(json.dumps(row, ensure_ascii=False) + "\n")
        print(f"{cat:<8} {len(rows):>4} items -> {cat}.jsonl")
        total += len(rows)
    print(f"TOTAL {total} items")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=DEFAULT_SRC)
    ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "data"))
    a = ap.parse_args()
    convert(a.src, a.out)

#!/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)

#!/usr/bin/env python3
"""Grade the Haiku sub-agent battery against the withheld gold key, reusing the
EXACT grader from vibe_eval (math-verify symbolic equiv for non-int golds, strict
int for AIME, letter-match for GPQA). Each agent returned a JSON object mapping
problem id -> final answer in its last assistant message; we extract that JSON
from the agent transcript files without dumping them into anyone's context.
"""
import json, re, sys, glob, os
from collections import defaultdict
from vibe_eval import grade  # reuse identical grading

GOLD = json.load(open("/tmp/haiku_gold.json"))   # id -> [category, answer]
TASKDIR = sys.argv[1]                              # dir holding <agentId>.output files
AGENTS = sys.argv[2:]                              # agentIds to scan

ID_RE = re.compile(r'^(aime|aime25|hmmt25|gpqa)-\d+$')

def extract_preds(path):
    """Find every JSON object in the transcript whose keys look like problem ids;
    merge them (last write wins). Robust to prose around the JSON."""
    raw = open(path, encoding="utf-8", errors="ignore").read()
    preds = {}
    # find balanced-ish {...} blobs; greedy per-line and whole-text candidates
    for m in re.finditer(r'\{[^{}]*\}', raw):
        blob = m.group(0)
        if '"aime' not in blob and '"hmmt' not in blob and '"gpqa' not in blob:
            continue
        try:
            d = json.loads(blob)
        except Exception:
            continue
        for k, v in d.items():
            if ID_RE.match(k) and isinstance(v, str):
                preds[k] = v
    return preds

allpreds = {}
for aid in AGENTS:
    p = os.path.join(TASKDIR, f"{aid}.output")
    if not os.path.exists(p):
        print(f"!! missing {aid}"); continue
    allpreds.update(extract_preds(p))

# grade
cats = defaultdict(lambda: {"n": 0, "ok": 0, "miss": 0})
missing = []
for iid, (cat, gold) in GOLD.items():
    cats[cat]["n"] += 1
    if iid not in allpreds:
        cats[cat]["miss"] += 1; missing.append(iid); continue
    raw = "\\boxed{%s}" % allpreds[iid]      # same wrapping the harness uses for cons grading
    ok, _ = grade(cat, raw, gold)
    cats[cat]["ok"] += int(ok)

order = ["aime", "aime25", "hmmt25", "gpqa"]
print("\n=== HAIKU 4.5 (closed-book, no-tools, sub-agent path) ===")
print(f"{'category':<10}{'pass@1':>9}{'correct':>9}{'n':>5}{'missing':>9}")
for c in order:
    if c not in cats: continue
    v = cats[c]; n = v["n"]
    print(f"{c:<10}{100*v['ok']/n:>8.1f}%{v['ok']:>9}{n:>5}{v['miss']:>9}")
if missing:
    print("\nmissing ids (no answer returned):", missing)
print(f"\ntotal answers parsed: {len(allpreds)}/140")

#!/usr/bin/env bash
# Drive the eval against the llama-swap endpoint on the Ubuntu box.
# llama-swap auto-loads/swaps the backend by the request's "model" field, so you
# do NOT serve models manually -- just run the three tags below in sequence.
#
# Runtime reality (llama.cpp, 5090, --parallel 4, ~500 tok/s aggregate):
#   llama.cpp splits -c across slots, so concurrency is capped -- this is MUCH
#   slower than vLLM for long reasoning traces. Ballpark:
#     math  avg@16 (90 items)  : ~8h vibe + ~10h qwen3-think + ~1h base  ~= 19h
#     gpqa  avg@16 (198 items) : ~9h each reasoner                       ~= 20h
#   => avg@16 on everything is ~1.5-2 DAYS. So: math stays avg@16 (headline,
#   needs tight CI); GPQA drops to avg@4 (it's the control, expected weak).
#   Bump N_GPQA if you want, or move to the vLLM path for ~5-10x throughput.
set -euo pipefail
BASE="${BASE:-http://LLAMA_SWAP_HOST:8000}"
N_MATH=16
N_GPQA=4
MATH="aime,aime25,hmmt25"
COMMON="--base $BASE --concurrency 4 --max-tokens 40960"

case "${1:?usage: BASE=... ./run_all.sh vibe3b|base3b|qwen3think}" in
  vibe3b)
    python3 vibe_eval.py $COMMON --model vibethinker-3b --tag vibe3b      --n $N_MATH --temp 1.0 --cats $MATH
    python3 vibe_eval.py $COMMON --model vibethinker-3b --tag vibe3b_gpqa --n $N_GPQA --temp 1.0 --cats gpqa ;;
  base3b)
    python3 vibe_eval.py $COMMON --model base-3b --tag base3b      --n $N_MATH --temp 0.7 --cats $MATH
    python3 vibe_eval.py $COMMON --model base-3b --tag base3b_gpqa --n $N_GPQA --temp 0.7 --cats gpqa ;;
  qwen3think)
    python3 vibe_eval.py $COMMON --model qwen3-4b-think --tag qwen3think      --n $N_MATH --temp 0.6 --cats $MATH
    python3 vibe_eval.py $COMMON --model qwen3-4b-think --tag qwen3think_gpqa --n $N_GPQA --temp 0.6 --cats gpqa ;;
esac