pkg/mmlu_pro_scorer.py

#!/usr/bin/env python
"""Self-contained MMLU-Pro downstream-quality scorer for a local OpenAI server.

Measures the REAL downstream-reasoning accuracy of a vLLM /v1/chat/completions
endpoint. Used to compare the lffn TREATMENT (LFFN_LINEAR=1, layer-29 affine
surrogate) against the CONTROL (LFFN_LINEAR=0, exact FFN). No inspect_evals
dependency; talks only to 127.0.0.1.

Review fixes folded in:
  * SAMPLING VERIFY (blocker): a greedy server-side OVERRIDE_GENERATION_CONFIG
    can silently override per-request temperature. Before scoring, send the SAME
    prompt twice at temp 1.0 with two DIFFERENT seeds; if the completions are
    identical the server is running greedy -> ABORT (the comparison would be
    invalid). The entrypoint also neutralizes the override, but we verify, not
    trust.
  * TRUNCATION (serious): capture finish_reason. finish_reason=="length" means
    thinking ran out of budget -> bucket it separately, never silently "wrong".
    Report accuracy three ways: over-all, over-parsed, over-finished.
  * POWER (serious): k samples per question (EVAL_K, default 3) with a
    per-(question,sample) seed (42 + idx*100 + k). Report mean accuracy and a
    bootstrap 95% CI on the control-treatment delta is computed offline from the
    two summary JSONs; here we emit the per-sample JSONL + a clean ACCURACY line.
  * PARSING: explicit "answer is (X)" form is tried over the FULL text first
    (final-answer line), only then a guarded bare-letter fallback on the tail.
    Raw completion tail is logged for unparsed rows so they can be audited.

Env / CLI knobs (CLI overrides env):
  --base-url        local server (default http://127.0.0.1:8000)
  --model           served model name (default gemma-4-e4b-it)
  --n               number of questions (default 300)            [EVAL_N]
  --k               samples per question (default 3)             [EVAL_K]
  --max-tokens      generation budget incl. thinking (default 8192) [EVAL_MAX_TOKENS]
  --seed            base seed for shuffle + sampling (default 42)
  --category        optional MMLU-Pro category filter
  --no-thinking     disable enable_thinking (default: thinking on)
  --skip-sampling-check  do not abort if greedy is detected (NOT recommended)
  --out / --summary writable result paths (default under /state)
"""
from __future__ import annotations

import argparse
import json
import os
import re
import string
import sys
import time

import requests

LETTERS = string.ascii_uppercase  # A..Z; MMLU-Pro uses up to J (10 options)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--base-url", default=os.environ.get("EVAL_BASE_URL", "http://127.0.0.1:8000"))
    p.add_argument("--model", default=os.environ.get("EVAL_MODEL", "gemma-4-e4b-it"))
    p.add_argument("--n", type=int, default=int(os.environ.get("EVAL_N", "300")))
    p.add_argument("--k", type=int, default=int(os.environ.get("EVAL_K", "3")))
    p.add_argument("--max-tokens", type=int, default=int(os.environ.get("EVAL_MAX_TOKENS", "8192")))
    p.add_argument("--temperature", type=float, default=float(os.environ.get("EVAL_TEMPERATURE", "1.0")))
    p.add_argument("--top-p", type=float, default=float(os.environ.get("EVAL_TOP_P", "0.95")))
    p.add_argument("--top-k", type=int, default=int(os.environ.get("EVAL_TOP_K", "64")))
    p.add_argument("--seed", type=int, default=int(os.environ.get("EVAL_SEED", "42")))
    p.add_argument("--category", default=os.environ.get("EVAL_CATEGORY", "") or None)
    p.add_argument("--no-thinking", action="store_true",
                   default=os.environ.get("EVAL_NO_THINKING", "0") == "1")
    p.add_argument("--skip-sampling-check", action="store_true",
                   default=os.environ.get("EVAL_SKIP_SAMPLING_CHECK", "0") == "1")
    p.add_argument("--request-timeout", type=float, default=float(os.environ.get("EVAL_REQUEST_TIMEOUT", "1200")))
    p.add_argument("--out", default=os.environ.get("EVAL_OUT", "/state/mmlu_pro_eval.jsonl"))
    p.add_argument("--summary", default=os.environ.get("EVAL_SUMMARY", "/state/mmlu_pro_summary.json"))
    return p.parse_args()


def load_questions(n: int, seed: int, category: str | None) -> list[dict]:
    from datasets import load_dataset

    ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test")
    if category:
        ds = ds.filter(lambda r: r["category"] == category)
    ds = ds.shuffle(seed=seed)
    rows = []
    for r in ds:
        rows.append(r)
        if len(rows) >= n:
            break
    return rows


def build_prompt(row: dict) -> str:
    opts = row["options"]
    lines = [f"{LETTERS[i]}. {opt}" for i, opt in enumerate(opts)]
    choices = "\n".join(lines)
    last = LETTERS[len(opts) - 1]
    return (
        "Answer the following multiple choice question. Think step by step, "
        "then give your final answer.\n\n"
        f"Question: {row['question']}\n\n"
        f"Options:\n{choices}\n\n"
        f"Reason carefully, then end your response with a line of EXACTLY this "
        f"form:\nThe answer is (X)\n"
        f"where X is one of the option letters A through {last}."
    )


# Ordered most-specific -> least-specific. "answer is (X)" wins.
_EXPLICIT_PATTERNS = [
    re.compile(r"answer\s*is\s*:?\s*\(?\s*([A-J])\s*\)?", re.IGNORECASE),
    re.compile(r"\banswer\s*[:\-]\s*\(?\s*([A-J])\s*\)?", re.IGNORECASE),
    re.compile(r"\boption\s*\(?\s*([A-J])\s*\)?\b", re.IGNORECASE),
    re.compile(r"\(([A-J])\)"),
]
# Guarded bare-letter fallback: a standalone letter on its own line.
_BARE_LINE = re.compile(r"(?m)^\s*\(?\s*([A-J])\s*\)?\s*[\.\)]?\s*$")


def parse_answer(text: str, n_options: int) -> str | None:
    """Extract the chosen letter. Explicit forms over the full text win; only
    then a guarded standalone-letter fallback. Always take the LAST match."""
    valid = set(LETTERS[:n_options])
    # 1. Explicit "answer is (X)" etc., over the FULL text (final answer is last).
    for pat in _EXPLICIT_PATTERNS:
        matches = [m.group(1).upper() for m in pat.finditer(text)]
        matches = [m for m in matches if m in valid]
        if matches:
            return matches[-1]
    # 2. Guarded fallback: a letter alone on its own line, near the end.
    tail = text[-600:] if len(text) > 600 else text
    matches = [m.group(1).upper() for m in _BARE_LINE.finditer(tail)]
    matches = [m for m in matches if m in valid]
    if matches:
        return matches[-1]
    return None


def chat_once(args: argparse.Namespace, prompt: str, seed: int) -> dict:
    body = {
        "model": args.model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": args.max_tokens,
        "temperature": args.temperature,
        "top_p": args.top_p,
        "top_k": args.top_k,
        "seed": seed,
        "stream": False,
    }
    if not args.no_thinking:
        body["chat_template_kwargs"] = {"enable_thinking": True}
    r = requests.post(
        f"{args.base_url.rstrip('/')}/v1/chat/completions",
        json=body,
        timeout=args.request_timeout,
    )
    r.raise_for_status()
    data = r.json()
    choice = data["choices"][0]
    return {
        "text": choice["message"].get("content") or "",
        "finish_reason": choice.get("finish_reason"),
    }


def assert_sampling_is_stochastic(args: argparse.Namespace) -> None:
    """BLOCKER guard: prove the server honours temperature>0. Send the same
    prompt twice at temp 1.0 with two different seeds; greedy decoding would make
    the two completions identical. If so, OVERRIDE_GENERATION_CONFIG (greedy)
    leaked through and the whole control/treatment delta is invalid -> abort."""
    probe = (
        "Write one short, creative sentence about the ocean. "
        "Be imaginative and avoid clichés."
    )
    a = chat_once(args, probe, seed=args.seed)["text"].strip()
    b = chat_once(args, probe, seed=args.seed + 9999)["text"].strip()
    print(f"[sampling-check] seedA_len={len(a)} seedB_len={len(b)} "
          f"identical={a == b}", flush=True)
    if a and a == b:
        msg = ("[sampling-check] FATAL: identical completions for two different "
               "seeds at temperature=1.0 -> server is running GREEDY "
               "(OVERRIDE_GENERATION_CONFIG leaked). The control/treatment delta "
               "would be invalid. Aborting.")
        print(msg, flush=True)
        if not args.skip_sampling_check:
            raise SystemExit(2)
    elif not a:
        print("[sampling-check] WARNING: empty probe completion; cannot verify "
              "sampling. Proceeding (use --skip-sampling-check to silence).",
              flush=True)


def main() -> int:
    args = parse_args()
    print(
        f"[eval] base_url={args.base_url} model={args.model} n={args.n} k={args.k} "
        f"temp={args.temperature} top_p={args.top_p} top_k={args.top_k} "
        f"seed={args.seed} thinking={not args.no_thinking} "
        f"max_tokens={args.max_tokens} category={args.category or 'ALL'} "
        f"LFFN_LINEAR={os.environ.get('LFFN_LINEAR')}",
        flush=True,
    )

    assert_sampling_is_stochastic(args)

    rows = load_questions(args.n, args.seed, args.category)
    print(f"[eval] loaded {len(rows)} questions x {args.k} samples", flush=True)

    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    n_samples = 0
    correct = 0
    parsed = 0
    finished = 0          # finish_reason == "stop"
    truncated = 0         # finish_reason == "length"
    correct_finished = 0
    errors = 0
    started = time.time()

    with open(args.out, "w", encoding="utf-8") as fh:
        for i, row in enumerate(rows):
            gold = row["answer"].strip().upper()
            n_opts = len(row["options"])
            prompt = build_prompt(row)
            for kk in range(args.k):
                sample_seed = args.seed + i * 100 + kk
                pred = None
                text = ""
                finish = None
                err = None
                try:
                    res = chat_once(args, prompt, seed=sample_seed)
                    text = res["text"]
                    finish = res["finish_reason"]
                    pred = parse_answer(text, n_opts)
                except Exception as exc:  # noqa: BLE001
                    err = repr(exc)
                    errors += 1
                n_samples += 1
                is_trunc = finish == "length"
                if is_trunc:
                    truncated += 1
                if finish == "stop":
                    finished += 1
                if err is None and pred is not None:
                    parsed += 1
                is_correct = pred is not None and pred == gold
                correct += int(is_correct)
                if finish == "stop":
                    correct_finished += int(is_correct)
                rec = {
                    "idx": i,
                    "k": kk,
                    "seed": sample_seed,
                    "question_id": row.get("question_id"),
                    "category": row.get("category"),
                    "gold": gold,
                    "pred": pred,
                    "correct": is_correct,
                    "n_options": n_opts,
                    "finish_reason": finish,
                    "error": err,
                    "completion_len": len(text),
                    # audit trail for unparsed/truncated rows only (keep log small)
                    "tail": (text[-300:] if (pred is None or is_trunc) else None),
                }
                fh.write(json.dumps(rec) + "\n")
                fh.flush()
                running_acc = correct / n_samples
                flag = "OK" if is_correct else ("LEN" if is_trunc else "XX")
                print(
                    f"[{i + 1}/{len(rows)} k{kk}] gold={gold} pred={pred} "
                    f"{flag} fin={finish} running_acc={running_acc:.4f}"
                    + (f" ERR={err}" if err else ""),
                    flush=True,
                )

    acc = correct / n_samples if n_samples else 0.0
    acc_parsed = correct / parsed if parsed else 0.0
    acc_finished = correct_finished / finished if finished else 0.0
    elapsed = time.time() - started
    summary = {
        "dataset": "TIGER-Lab/MMLU-Pro",
        "split": "test",
        "category": args.category or "ALL",
        "n_questions": len(rows),
        "k": args.k,
        "n_samples": n_samples,
        "correct": correct,
        "accuracy": acc,
        "accuracy_parsed_only": acc_parsed,
        "accuracy_finished_only": acc_finished,
        "parsed": parsed,
        "unparsed": n_samples - parsed - errors,
        "finished": finished,
        "truncated": truncated,
        "errors": errors,
        "model": args.model,
        "sampling": {
            "temperature": args.temperature, "top_p": args.top_p,
            "top_k": args.top_k, "base_seed": args.seed,
            "enable_thinking": not args.no_thinking, "max_tokens": args.max_tokens,
        },
        "lffn_linear": os.environ.get("LFFN_LINEAR"),
        "elapsed_s": elapsed,
        "per_question_file": args.out,
    }
    os.makedirs(os.path.dirname(args.summary) or ".", exist_ok=True)
    with open(args.summary, "w", encoding="utf-8") as fh:
        json.dump(summary, fh, indent=2)

    print("\n===== MMLU-Pro RESULT =====", flush=True)
    print(f"ACCURACY={acc:.4f}  ({correct}/{n_samples} samples, "
          f"{len(rows)}q x k{args.k})", flush=True)
    print(f"ACCURACY_PARSED_ONLY={acc_parsed:.4f}  "
          f"ACCURACY_FINISHED_ONLY={acc_finished:.4f}", flush=True)
    print(f"parsed={parsed} truncated={truncated} errors={errors} "
          f"elapsed_s={elapsed:.1f}", flush=True)
    print(f"LFFN_LINEAR={os.environ.get('LFFN_LINEAR')}", flush=True)
    print(f"summary_file={args.summary}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())