tools/eval-harness/inspect_driver.py

#!/usr/bin/env python
"""Run GPQA-Diamond + AIME + MMLU-Pro via inspect_evals against a served
OpenAI-compatible endpoint (our cleanstack serve.py). Emits one scores JSON to
--out and prints it between sentinels so it can be recovered from job logs
(personal-namespace jobs cannot write-mount org buckets).

Capability eval for the gemma-challenge evals taskforce. Sampling + thinking are
the disputed protocol knobs — defaults follow @human-lewtun's recommended-sampling
instruction (generation_config.json: temp 1.0 / top_p 0.95 / top_k 64) with a
fixed seed, enable_thinking=True via vLLM chat_template_kwargs.
"""
from __future__ import annotations
import argparse, json, os, traceback, urllib.request, urllib.error

from inspect_ai import eval as ieval

THINK_TOKEN = "<|think|>"   # gemma-4 template injects this when enable_thinking is true


def render_prompt(base_url: str, model: str, enable_thinking: bool) -> str:
    """Ask vLLM to render a chat request to its prompt string so we can confirm
    the enable_thinking chat_template_kwarg actually reaches the template."""
    body = {
        "model": model,
        "messages": [{"role": "system", "content": "You are a helpful assistant."},
                     {"role": "user", "content": "What is 2+2?"}],
        "max_tokens": 1,
        "chat_template_kwargs": {"enable_thinking": enable_thinking},
    }
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions/render",
        data=json.dumps(body).encode(),
        headers={"content-type": "application/json", "authorization": "Bearer EMPTY"},
        method="POST")
    with urllib.request.urlopen(req, timeout=30) as r:
        return r.read().decode("utf-8", "replace")


def verify_thinking(base_url: str, model: str) -> dict:
    """Hard gate: confirm the enable_thinking chat_template_kwarg actually reaches
    the template. The render output is the ONLY thing that varies between the two
    calls (identical messages, no date/strftime in the gemma-4 template), so
    on != off proves the kwarg is plumbed. The literal think-token may not appear
    if /render returns token-ids, so we gate on `differ` (+ on being longer, since
    the template only ADDS '<|think|>\\n'); we also log raw samples to eyeball.
    Never raises (caller decides)."""
    out: dict = {"ok": False}
    try:
        on = render_prompt(base_url, model, True)
        off = render_prompt(base_url, model, False)
        out["differ"] = on != off
        out["think_token_literal_in_on"] = THINK_TOKEN in on
        out["len_on"], out["len_off"] = len(on), len(off)
        out["sample_on"] = on[:300]
        out["sample_off"] = off[:300]
        # plumbed iff the only differing input (enable_thinking) changed the render,
        # and it added content (think injection only grows the prompt).
        out["ok"] = bool(out["differ"]) and len(on) >= len(off)
    except Exception as e:  # noqa: BLE001
        out["error"] = repr(e)
    return out


def extract(log) -> dict:
    out = {"status": log.status, "n": None, "accuracy": None, "stderr": None}
    if getattr(log, "error", None):
        out["error"] = str(log.error)
    if log.results:
        out["n"] = log.results.completed_samples
        for sc in log.results.scores:
            for mname, metric in sc.metrics.items():
                if mname == "accuracy" or (out["accuracy"] is None and "accuracy" in mname):
                    out["accuracy"] = metric.value
                if mname in ("stderr", "std_err") or (out["stderr"] is None and "stderr" in mname):
                    out["stderr"] = metric.value
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url", default="http://127.0.0.1:8000/v1")
    ap.add_argument("--model-name", default="gemma-4-e4b-it")
    ap.add_argument("--limit-gpqa", type=int, default=50)
    ap.add_argument("--limit-aime", type=int, default=30)
    ap.add_argument("--limit-mmlu", type=int, default=60)
    ap.add_argument("--max-tokens", type=int, default=8192)
    ap.add_argument("--temperature", type=float, default=1.0)
    ap.add_argument("--top-p", type=float, default=0.95)
    ap.add_argument("--top-k", type=int, default=64)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--max-connections", type=int, default=1)
    ap.add_argument("--limit-math", type=int, default=0)        # high-power reasoning confirm (MATH); 0=skip
    ap.add_argument("--limit-aime2025", type=int, default=0)     # AIME 2025 (combine w/ 2024 for n=60); 0=skip
    ap.add_argument("--enable-thinking", type=int, default=1)
    ap.add_argument("--log-dir", default="/tmp/inspect_logs")
    ap.add_argument("--out", default="/tmp/eval_scores.json")
    args = ap.parse_args()

    os.environ.setdefault("OPENAI_API_KEY", "EMPTY")

    # Hard gate: confirm enable_thinking actually reaches the template BEFORE any
    # paid generation. Getting the protocol wrong (e.g. silent no-think, or the
    # @3000 truncation artifact) is exactly what polluted our prior number.
    think_check = verify_thinking(args.base_url, args.model_name)
    print(f"[driver] thinking-gate: {json.dumps(think_check)}", flush=True)
    if bool(args.enable_thinking) and not think_check.get("ok"):
        print("[driver] ABORT: enable_thinking did not take effect; not spending on generation.",
              flush=True)
        with open(args.out, "w") as f:
            json.dump({"_aborted": "thinking_gate_failed", "thinking_check": think_check}, f, indent=2)
        print("=====EVAL_SCORES_JSON=====", flush=True)
        print(json.dumps({"_aborted": "thinking_gate_failed", "thinking_check": think_check}), flush=True)
        print("=====END_EVAL_SCORES_JSON=====", flush=True)
        return

    from inspect_evals.gpqa import gpqa_diamond
    from inspect_evals.aime2024 import aime2024
    from inspect_evals.mmlu_pro import mmlu_pro
    from inspect_evals.math import math as math_task
    from inspect_evals.aime2025 import aime2025

    # top_k + enable_thinking are vLLM extensions -> route through extra_body so
    # they survive inspect's OpenAI provider (which only maps standard fields).
    extra_body = {"top_k": args.top_k,
                  "chat_template_kwargs": {"enable_thinking": bool(args.enable_thinking)}}

    common = dict(
        model="openai/" + args.model_name,
        model_base_url=args.base_url,
        model_args={"responses_api": False},   # force chat-completions: keeps seed + extra_body honored
        log_dir=args.log_dir,
        temperature=args.temperature,
        top_p=args.top_p,
        max_tokens=args.max_tokens,
        seed=args.seed,
        extra_body=extra_body,
        max_connections=args.max_connections,
        epochs=args.epochs,
        fail_on_error=0.25,   # tolerate a few API hiccups per task
    )

    benches = [
        ("gpqa_diamond", gpqa_diamond, args.limit_gpqa),
        ("aime2024", aime2024, args.limit_aime),
        ("mmlu_pro", mmlu_pro, args.limit_mmlu),
        ("math", math_task, args.limit_math),
        ("aime2025", aime2025, args.limit_aime2025),
    ]

    scores: dict = {}
    for key, taskfn, lim in benches:
        if not lim or lim <= 0:
            continue
        try:
            logs = ieval(taskfn(), limit=lim, **common)
            scores[key] = extract(logs[0])
        except Exception as e:  # noqa: BLE001 — one bench failing must not kill the rest
            scores[key] = {"status": "exception", "error": repr(e),
                           "trace": traceback.format_exc()[-1500:]}
        print(f"[driver] {key}: {json.dumps(scores[key])[:400]}", flush=True)

    scores["_meta"] = {
        "thinking_check": think_check,
        "enable_thinking": bool(args.enable_thinking),
        "temperature": args.temperature, "top_p": args.top_p, "top_k": args.top_k,
        "seed": args.seed, "max_tokens": args.max_tokens, "api": "chat_completions",
        "limits": {"gpqa": args.limit_gpqa, "aime": args.limit_aime, "mmlu_pro": args.limit_mmlu},
    }
    with open(args.out, "w") as f:
        json.dump(scores, f, indent=2)
    print("=====EVAL_SCORES_JSON=====", flush=True)
    print(json.dumps(scores), flush=True)
    print("=====END_EVAL_SCORES_JSON=====", flush=True)


if __name__ == "__main__":
    main()