initial commit

eval/.gitignore

@@ -0,0 +1,2 @@
+outputs/
+data/

eval/README.md

@@ -0,0 +1,221 @@
+# Benchmark evaluation
+
+Offline batch evaluation for AIME 2026, HLE, and GPQA Diamond, built on the
+DeepInfer offline batch API.
+
+Generation and scoring are separate stages. The generation stage submits every
+request to the engine in a single batch and writes the raw model outputs to disk;
+the scoring stage reads those outputs and produces the final metrics. Scoring can
+therefore be repeated (with a different judge, for example) without regenerating.
+
+The reported metric is **Pass@1**.
+
+| Benchmark | Tasks | Samples per task | How Pass@1 is computed |
+| :--- | ---: | ---: | :--- |
+| [AIME 2026](aime/run.sh) | 30 | 32 | mean accuracy over 32 samples per task |
+| [HLE](hle/run.sh) | 2500 (2158 text, 342 image) | 1 | single-sample accuracy |
+| [GPQA Diamond](gpqa/run.sh) | 198 | 8 | mean accuracy over 8 samples per task |
+
+---
+
+## 1. Layout
+
+```
+eval/
+|-- run_eval.py          # generation entry point (offline batch)
+|-- score.py             # scoring entry point
+|-- common/              # shared utilities: paths, JSONL I/O, Pass@k, loader, judge, encoder
+|-- aime/                # one directory per benchmark
+|   |-- run.sh           #   two-stage driver: inference, then evaluation
+|   `-- bench.py         #   prompt construction, answer extraction, scoring
+|-- hle/
+|-- gpqa/
+|-- data/                # datasets (or point DATA_FILE elsewhere)
+`-- outputs/             # results
+```
+
+---
+
+## 2. Requirements
+
+Complete the backend-specific installation described in the top-level [README.md](../README.md):
+`./build.sh` installs the requirements and builds the vendored operators and the in-tree
+DeepInfer extension. Keep `DEEPINFER_ACCELERATOR` exported when running evaluation, because it
+selects the platform build of the extension. No additional build step is needed for evaluation.
+
+Chat prompts are rendered by the encoder distributed with the checkpoint
+(`<checkpoint>/encoding/encoding.py`), which the scripts load at runtime. The
+checkpoint directory must contain that file.
+
+---
+
+## 3. Data preparation
+
+The evaluation scripts do not download datasets. Each dataset must be obtained
+from its upstream source and converted to JSONL.
+
+| Benchmark | Source | Fields |
+| :--- | :--- | :--- |
+| AIME 2026 | [MathArena/aime_2026](https://huggingface.co/datasets/MathArena/aime_2026), or the official AIME 2026 I/II problems | `task_id`, `question`, `answer` (an integer string in the range 0-999) |
+| HLE | [cais/hle](https://huggingface.co/datasets/cais/hle) (gated; accept the terms on Hugging Face first) | `task_id`, `question`, `answer`, optional `category`; image questions are described below |
+| GPQA Diamond | [Idavidrein/gpqa](https://huggingface.co/datasets/Idavidrein/gpqa) (gated) or the public mirror [hendrydong/gpqa_diamond_mc](https://huggingface.co/datasets/hendrydong/gpqa_diamond_mc) | `task_id`, `question`, `answer` (a single letter, A-D) |
+
+### 3.1 Record format
+
+Each dataset is a JSONL file with the fields shown above. Place it at the default path
+given in Section 3.3, or point `DATA_FILE` at it.
+
+```jsonc
+// AIME 2026
+{"task_id": "aime_2026_0", "question": "Patrick started walking ...", "answer": "277"}
+
+// HLE, text question
+{"task_id": "hle_0001", "question": "...", "answer": "...", "category": "Math"}
+
+// GPQA Diamond: The question must contain the four choices inline, labelled (A)-(D).
+{"task_id": "gpqa_diamond_0", "question": "Two quantum states ...\n\n(A) 10^-11 eV\n(B) 10^-8 eV\n(C) 10^-9 eV\n(D) 10^-4 eV", "answer": "D", "category": "Physics"}
+```
+
+### 3.2 Image questions
+
+An HLE image question may reference its image in any of the following ways:
+
+```jsonc
+// 1. Path to a local file
+{"task_id": "hle_0500", "question": "...", "answer": "...", "category": "Biology/Medicine",
+ "image": "/path/to/500.png"}
+
+// 2. Inline base64
+{"task_id": "hle_0501", "question": "...", "answer": "...", "category": "Physics",
+ "image": "<base64-string>"}
+
+// 3. multi_modal_data, as consumed by the engine (optionally with image dimensions)
+{"task_id": "hle_0502", "question": "...<image>", "answer": "...", "category": "Chemistry",
+ "multi_modal_data": {"image_urls": ["/path/to/502.jpg"], "image_sizes": [[1080, 1078]]}}
+```
+
+The question must contain an `<image>` placeholder wherever the image belongs;
+generation replaces it with the image token the engine understands.
+
+### 3.3 Locating the data file
+
+Each driver script defines a default `DATA_FILE`:
+
+| Benchmark | Default `DATA_FILE` |
+| :--- | :--- |
+| aime | `eval/data/aime_2026.jsonl` |
+| hle | `eval/data/hle_test.jsonl` |
+| gpqa | `eval/data/gpqa_diamond.jsonl` |
+
+Place the dataset at the corresponding path, or point `DATA_FILE` at it:
+
+```bash
+cd <repo-root>
+
+bash eval/aime/run.sh
+bash eval/gpqa/run.sh
+
+DATA_FILE=/data/my_aime.jsonl bash eval/aime/run.sh
+DATA_FILE=/data/my_hle.jsonl  bash eval/hle/run.sh
+DATA_FILE=/data/my_gpqa.jsonl bash eval/gpqa/run.sh
+```
+
+---
+
+## 4. Running an evaluation
+
+Each driver script runs two stages:
+
+```
+[1/2] inference  ->  [2/2] evaluate
+```
+
+Defaults can be overridden with environment variables:
+
+```bash
+EP=8 N_SAMPLE=4 TAG=smoke bash eval/aime/run.sh            # smoke test
+bash eval/aime/run.sh                                      # full run with the default model
+MODEL=/path/to/DeepSeek-V4.1-Flash-1017 bash eval/aime/run.sh   # local HF directory
+```
+
+Available overrides: `MODEL`, `EP`, `N_SAMPLE`, `REASONING_EFFORT`, `TAG`,
+`DATA_FILE`.
+
+For HLE the sampling configuration is not adjustable through environment
+variables: `n_sample=1` is set in the driver script, and `temperature=1.0` is the
+per-benchmark default in `run_eval.py`. To run multiple samples, edit `--n-sample`
+in the script. For AIME and GPQA the sample count is configurable through
+`N_SAMPLE`, defaulting to 32 and 8 respectively.
+
+Scoring uses an LLM judge service by default; see the next section. A
+deterministic rule-based judge is also provided.
+
+---
+
+## 5. Scoring
+
+Two scoring modes are available. The default is the LLM judge, which is also the
+protocol used by the reference evaluation.
+
+| Mode | Method | Dependencies | Determinism |
+| :--- | :--- | :--- | :--- |
+| `llm` (default) | A judge model decides whether the extracted answer is equivalent to the reference answer. | A judge service | Subject to the judge model's own sampling |
+| `rule` | The answer is extracted by fixed rules and compared exactly: the integer inside `\boxed{}` for AIME, the letter following `ANSWER:` for GPQA. | None | Fully deterministic |
+
+HLE supports the LLM judge only. More than half of its answers are free-form
+text, so exact matching would mark equivalent answers as wrong and systematically
+understate the score; `score.py` therefore rejects `--judge rule` for HLE.
+
+To use the LLM judge, provide a judge service:
+
+```bash
+export EVAL_JUDGE_BASE_URL=https://api.deepseek.com
+export EVAL_JUDGE_API_KEY=<key>
+export EVAL_JUDGE_MODEL=deepseek-v4-flash
+bash eval/aime/run.sh
+```
+
+To score without a judge service, uncomment the rule-based command in stage 2/2 of
+the driver script, or re-score existing output directly; no regeneration is
+required:
+
+```bash
+python eval/score.py --out-dir eval/outputs/aime/<TAG> --judge rule
+```
+
+When the LLM judge is used, `summary.json` additionally reports
+`rule_vs_llm_agreement`, the fraction of rollouts on which both modes agree. On
+AIME 2026 the two modes agreed on all 960 rollouts.
+
+---
+
+## 6. Outputs
+
+Results are written to `eval/outputs/<benchmark>/<TAG>/`:
+
+| File | Contents |
+| :--- | :--- |
+| `summary.json` | Aggregate metrics. `Pass@1` is the headline figure; `Pass@k`, `num_tasks`, `num_rollouts` and `coverage_complete` are also reported. |
+| `scored.jsonl` | One row per rollout: raw model output, extracted answer, verdict. |
+| `raw_rank*.jsonl` | Raw model outputs as written by each data-parallel rank. |
+| `meta.json` | The complete configuration of the run, for reproducibility. |
+
+A `Pass@1` figure is only valid when `coverage_complete` in `summary.json` is
+`true`; otherwise one or more ranks or samples are missing.
+
+---
+
+## 7. Reference results
+
+The following Pass@1 values are reference figures for the V4.1 model under the
+default configuration described above.
+
+| Benchmark | Pass@1 |
+| :--- | ---: |
+| AIME 2026 | 99.9% |
+| GPQA Diamond | 93.4% |
+| HLE-Full | 42.6% |
+
+Due to the stochastic nature of LLM generation, repeated runs may differ by a small margin. 
+These figures are therefore indicative rather than exact, and serve as a sanity check
+that a deployment is configured correctly.

eval/aime/bench.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+EVAL_DIR = Path(__file__).resolve().parent.parent  # eval/
+if str(EVAL_DIR) not in sys.path:
+    sys.path.insert(0, str(EVAL_DIR))
+
+from common.io import require_file  # noqa: E402
+
+EN_BOXED_INSTR = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+
+
+@dataclass(frozen=True)
+class Problem:
+    task_id: str
+    question: str
+    answer: str
+
+
+# --------------------------------------------------------------------------- #
+# Data
+# --------------------------------------------------------------------------- #
+
+
+def load(data_file: str | os.PathLike) -> list[Problem]:
+    path = Path(data_file)
+    require_file(path, "AIME 2026 download: https://huggingface.co/datasets/MathArena/aime_2026")
+
+    problems: list[Problem] = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            row = json.loads(line)
+            problems.append(
+                Problem(
+                    task_id=row["task_id"],
+                    question=row["question"],
+                    answer=str(row["answer"]),
+                )
+            )
+    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
+    return problems
+
+
+# --------------------------------------------------------------------------- #
+# prompt
+# --------------------------------------------------------------------------- #
+
+
+def build_question(problem: Problem) -> str:
+    return EN_BOXED_INSTR.format(question=problem.question.strip())
+
+
+def build_messages(problem: Problem) -> list[dict[str, Any]]:
+    """Single-turn user message; rendered into the model chat template by the checkpoint encoder."""
+    return [{"role": "user", "content": build_question(problem)}]
+
+
+# --------------------------------------------------------------------------- #
+# Answer extraction and scoring
+# --------------------------------------------------------------------------- #
+
+
+def extract_boxed(text: str) -> str | None:
+    r"""Return the content of the last \boxed{...}, supporting nested braces (e.g. \boxed{\frac{1}{2}})."""
+    token = "\\boxed"
+    starts = [match.start() for match in re.finditer(re.escape(token), text)]
+    for start in reversed(starts):
+        open_brace = text.find("{", start + len(token))
+        if open_brace < 0:
+            continue
+        depth = 0
+        for pos in range(open_brace, len(text)):
+            if text[pos] == "{":
+                depth += 1
+            elif text[pos] == "}":
+                depth -= 1
+                if depth == 0:
+                    return text[open_brace + 1 : pos]
+    return None
+
+
+_NUMBER_RE = re.compile(r"(?<![\w.])(\d{1,3})(?![\w.])")
+
+
+def normalize_answer(ans: str | None) -> str | None:
+    """Normalize an answer into a comparable string: strip $, braces, whitespace, thousands separators and leading +/0."""
+    if ans is None:
+        return None
+    text = ans.strip()
+    text = text.strip("$").strip()
+    text = text.replace("\\left", "").replace("\\right", "")
+    text = text.replace("\\!", "").replace("\\,", "").replace("\\;", "")
+    text = text.replace("{,}", ",") if "{,}" in text else text  # LaTeX thousands separator: \boxed{1{,}234}
+    text = text.replace(",", "").replace(" ", "")
+    text = text.rstrip(".")
+    if text.startswith("+"):
+        text = text[1:]
+    if re.fullmatch(r"-?\d+", text):
+        return str(int(text))
+    return text
+
+
+# The fallback scans only the tail: reasoning text holds many numbers, so scanning the whole
+# text would easily pick up a value from an intermediate draft.
+_TAIL_CHARS = 400
+
+
+def extract_answer_with_source(text: str) -> tuple[str | None, str]:
+    r"""Extract the answer and report which path produced it.
+
+    1. `boxed`: the content of the last \boxed{...}. This is the format the prompt requests, so it
+       is the most reliable path;
+    2. `tail_number`: the last 0-999 integer after `</think>` (or within the final 400 characters).
+       A fallback for malformed output, **higher risk**; score.py reports its share separately;
+    3. `none`: nothing found.
+    """
+    boxed = extract_boxed(text)
+    if boxed is not None:
+        norm = normalize_answer(boxed)
+        if norm:
+            return norm, "boxed"
+    tail = text.rsplit("</think>", 1)[-1] if "</think>" in text else text[-_TAIL_CHARS:]
+    nums = _NUMBER_RE.findall(tail)
+    if nums:
+        return normalize_answer(nums[-1]), "tail_number"
+    return None, "none"
+
+
+def is_correct(pred: str | None, gold: str) -> bool:
+    """Exact integer match; non-integer answers fall back to case- and whitespace-insensitive string equality."""
+    if pred is None:
+        return False
+    pred_norm, gold_norm = normalize_answer(pred), normalize_answer(gold)
+    if pred_norm is None or gold_norm is None:
+        return False
+    if re.fullmatch(r"-?\d+", pred_norm) and re.fullmatch(r"-?\d+", gold_norm):
+        return int(pred_norm) == int(gold_norm)
+    return pred_norm.lower() == gold_norm.lower()
+
+
+# --------------------------------------------------------------------------- #
+# Entry points used by run_eval.py / score.py
+# --------------------------------------------------------------------------- #
+
+
+def build_records(problems: list[Problem], *, limit: int | None = None) -> list[dict[str, Any]]:
+    records = []
+    for problem in problems:
+        records.append(
+            {
+                "task_id": problem.task_id,
+                "question": problem.question,
+                "answer": problem.answer,
+                "judge_question": build_question(problem),
+                "messages": build_messages(problem),
+            }
+        )
+        if limit is not None and len(records) >= limit:
+            break
+    return records
+
+
+def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
+    """Score one rollout. The rollout must carry raw_text (the model output)."""
+    pred, source = extract_answer_with_source(rollout.get("raw_text", ""))
+    row = dict(rollout)
+    row["pred"] = pred
+    row["extract_source"] = source
+    row["gold"] = record["answer"]
+    row["correct"] = is_correct(pred, record["answer"])
+    return row

eval/aime/run.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+set -euo pipefail
+
+EVAL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO="$(cd "${EVAL_DIR}/.." && pwd)"
+cd "${REPO}"
+
+MODEL="${MODEL:-deepseek-ai/DeepSeek-V4.1-Flash-1017}"
+EP="${EP:-8}"
+N_SAMPLE="${N_SAMPLE:-32}"
+REASONING_EFFORT="${REASONING_EFFORT:-100}"
+TAG="${TAG:-${MODEL//\//_}-n${N_SAMPLE}-j${REASONING_EFFORT}}"
+DATA_FILE="${DATA_FILE:-${EVAL_DIR}/data/aime_2026.jsonl}"
+OUT="${EVAL_DIR}/outputs/aime/${TAG}"
+
+echo "[aime] model=${MODEL} ep=${EP} n_sample=${N_SAMPLE} effort=${REASONING_EFFORT} -> ${OUT}"
+
+# 1) inference
+python eval/run_eval.py --bench aime --model-name "${MODEL}" --ep "${EP}" \
+    --n-sample "${N_SAMPLE}" --reasoning-effort "${REASONING_EFFORT}" \
+    --tag "${TAG}" --data-file "${DATA_FILE}"
+
+# 2) evaluate
+# LLM judging: need EVAL_JUDGE_BASE_URL / EVAL_JUDGE_API_KEY / EVAL_JUDGE_MODEL
+python eval/score.py --out-dir "${OUT}" --judge llm --judge-concurrency 16
+
+# rule judging: use this instead when no LLM judge is available
+# python eval/score.py --out-dir "${OUT}" --judge rule
+
+echo "done: ${OUT}/summary.json"

eval/common/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

eval/common/encoder.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def load_encoder(model_name: str) -> ModuleType:
+    """Import the chat encoder that ships inside the checkpoint directory."""
+    from examples._config import resolve_model_path
+
+    path = Path(resolve_model_path(model_name)) / "encoding" / "encoding.py"
+    if not path.is_file():
+        raise SystemExit(f"checkpoint has no encoding/encoding.py: {path}")
+    mod_name = "ckpt_encoding"
+    if mod_name in sys.modules:
+        return sys.modules[mod_name]
+    spec = importlib.util.spec_from_file_location(mod_name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[mod_name] = module
+    spec.loader.exec_module(module)
+    return module

eval/common/io.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+
+def read_jsonl(path: str | os.PathLike) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                rows.append(json.loads(line))
+    return rows
+
+
+def write_jsonl(path: str | os.PathLike, rows: Iterable[dict[str, Any]]) -> int:
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    count = 0
+    with open(path, "w") as f:
+        for row in rows:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+            count += 1
+    return count
+
+
+def require_file(path: Path, hint: str) -> Path:
+    """Raise an actionable error when the data file is missing (this directory ships no datasets)."""
+    if not path.is_file():
+        raise FileNotFoundError(
+            f"data file not found: {path}\n"
+            "The evaluation scripts do not download datasets. Obtain it and place it at this path, "
+            f"or set DATA_FILE / --data-file.\n{hint}"
+        )
+    return path

eval/common/judge.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+import requests
+
+ANSWER_EXACT_MATCH_PROMPT = """
+# Role: Deterministic Autograder
+
+You are a precise, automated grading system. Your sole function is to evaluate whether a student's final answer matches the reference answer. You must NOT grade the reasoning, steps, or methodology—only the final extracted result.
+
+# Input Data
+
+<|Start of Question|>
+{question}
+<|End of Question|>
+
+<|Start of Reference Answer|>
+{answer}
+<|End of Reference Answer|>
+
+<|Start of Student's Answer|>
+{output}
+<|End of Student's Answer|>
+
+# Grading Guidelines (Equivalence Rules)
+
+Equivalence is mandatory for a `correct` grade. You must rigorously verify if the answers represent the exact same value, meaning, or expression, even if formatted differently. Actively try to prove equivalence before marking an answer as wrong.
+
+1. **Algebraic Equivalence:** Expressions like `n(n+1)/2` and `n^2/2 + n/2` or `x/sqrt(2)` and `x*sqrt(2)/2` are equivalent. You must mathematically simplify expressions to check for a match.
+2. **Numerical Equivalence:** Fractions, decimals, and percentages that hold the same value are equivalent (e.g., `1/2` == `0.5` == `50%`). 
+3. **Set/List Equivalence:** Unless the question specifically asks for an ordered tuple/vector, the order of elements does not matter (e.g., `{{1, 2}}` is equivalent to `{{2, 1}}`).
+4. **Unit Equivalence:** Ignore minor unit formatting variations. If a student includes correct units and the reference does not (or vice versa), do not penalize them.
+5. **Text Equivalence:** Ignore case, minor typos, and punctuation differences unless they completely change the mathematical/logical meaning.
+6. **Multiple-Choice Equivalence**: For multiple-choice questions, the student's answer is correct if it matches either the option letter (e.g., "A") or the content of the correct option (e.g., answering "100" when the correct option is "A: 100").
+7. **No Partial Credit:** If the student's final answer is incomplete, contains multiple contradictory answers, or is partially incorrect, it is `wrong`.
+8. **No Answer:** If no clear, unambiguous final answer can be extracted from the student's text, it is `wrong`.
+
+# Instructions
+
+- Compare the two answers using the Grading Guidelines, and verify their equivalence.
+- You must NOT grade the reasoning or steps, only the final answer.
+- If the student’s final answer is correct according to the reference answer, output `correct`.
+- Be concise in your thinking, do NOT comment on any background to the problem, do NOT attempt to solve the problem, do NOT argue for any answer different than the reference answer, focus only on whether the answers match.
+
+Output exactly one of `correct` and `wrong`, without any additional text.
+""".strip()
+
+
+HLE_JUDGE_PROMPT = r"""
+Judge whether the following [response] to [question] is correct or not based on the precise and unambiguous [correct_answer] below.
+
+[question]: {question}
+
+[response]: {response}
+
+Your judgement must be in the format and criteria specified below:
+
+extracted_final_answer: The final exact answer extracted from the [response]. Put the extracted answer as 'None' if there is no exact, final answer to extract from the response.
+
+[correct_answer]: {correct_answer}
+
+reasoning: Explain why the extracted_final_answer is correct or incorrect based on [correct_answer], focusing only on if there are meaningful differences between [correct_answer] and the extracted_final_answer. Do not comment on any background to the problem, do not attempt to solve the problem, do not argue for any answer different than [correct_answer], focus only on whether the answers match.
+
+correct: Answer 'yes' if extracted_final_answer matches the [correct_answer] given above, or is within a small margin of error for numerical problems. Answer 'no' otherwise, i.e. if there if there is any inconsistency, ambiguity, non-equivalency, or if the extracted answer is incorrect.
+
+
+confidence: The extracted confidence score between 0|\%| and 100|\%| from [response]. Put 100 if there is no confidence score available.
+""".strip()
+
+
+# Two judge protocols: answer equivalence (general / short math answers) and the official HLE protocol
+JUDGE_KINDS = ("equivalence", "hle")
+
+
+# HLE verdict parsing. The judge is a thinking model rather than the official o3-mini with
+# structured output, so the verdict field drifts. Accepting only a strict `correct:` line
+# silently scores drifting answers as wrong, so fall back through:
+#   `correct:` field (tolerating quotes/brackets/markdown, last occurrence wins)
+#   -> line-initial `Answer: yes|no`
+#   -> wording of the explanation (skip lines that carry the graded model's own answer,
+#      neutralise double negatives, then let a negation outrank an affirmation)
+# If none of the three match, the sample counts as wrong and is flagged unparsed.
+_HLE_CORRECT_RE = re.compile(r"(?<![a-z])correct[\"'\]\*_]*\s*:\s*[\"'\*_ ]*(yes|no)\b")
+_HLE_ANSWER_FALLBACK_RE = re.compile(r"^[\*_ ]*answer[\*_ ]*:\s*[\"'\*_ ]*(yes|no)\b", re.MULTILINE)
+# Under the drift these lines carry the graded model's answer (possibly containing words such as
+# "is correct" or "none"), so they must not take part in the wording decision; if an
+# explanation:/reasoning: section follows on the same line, keep that part.
+_HLE_WORDING_SKIP_LINE_RE = re.compile(
+    r"^[\*_ \"]*(?:extracted_final_answer|answer|confidence)[\*_\"]*\s*:(?:(?!explanation:|reasoning:).)*",
+    re.MULTILINE,
+)
+_HLE_WORDING_NEUTRAL_RE = re.compile(
+    # double negatives: no meaningful difference / not different / does not contradict ...
+    r"\bno\s+(?:\w+\s+){0,2}(?:differences?|discrepanc(?:y|ies)|inconsistenc(?:y|ies)|mismatch(?:es)?|conflicts?|deviations?|contradictions?|non-equivalenc(?:e|y))\b(?:\s+or\s+\w+)?"
+    r"|\bnot\s+(?:\w+\s+)?(?:different|inconsistent|incorrect|wrong)\b"
+    r"|\b(?:does|do)\s+not\s+(?:\w+\s+)?(?:contradict|create|constitute|introduce|represent|affect|alter|change|make|render|amount|rise)\b[^\n.;]{0,80}"
+    r"|\bdiffer(?:s|ing)?\s+only\s+in\b|\bonly\s+differ(?:s|ing|ence)?\s+(?:in|is)\b"
+    r"|\bdiffers?\s+(?:from[^\n]{0,80}?)?\bby\s+(?:only|just|a\s+negligible|a\s+small|less\s+than|about|roughly|approximately)\b"
+    r"|\bwithout\s+(?:any\s+)?(?:\w+\s+)?(?:differences?|inconsistenc(?:y|ies)|discrepanc(?:y|ies))\b"
+    r"|\bnot\s+(?:a\s+)?(?:meaningful|significant|substantive)\b"
+    # the judge echoing the instructions it was given
+    r"|\bcorrect\s+or\s+incorrect\b|\bfocusing\s+only\s+on\s+if\s+there\s+are\s+meaningful\s+differences\b"
+    r"|\bif\s+there\s+(?:if\s+there\s+)?is\s+any\s+inconsistency[^.\n]*"
+    r"|\banswer\s+'?no'?\s+otherwise\b|\bor\s+if\s+the\s+extracted\s+answer\s+is\s+incorrect\b"
+    r"|\bwhy\s+the\s+extracted_final_answer\s+is\b"
+)
+_HLE_WORDING_NO_RE = re.compile(
+    r"\b(?:does|do|did)\s+not\s+(?:\w+\s+)?(?:match|equal|correspond|agree|coincide|align)"
+    r"|\bdoesn'?t\s+(?:\w+\s+)?(?:match|equal|correspond|agree)"
+    r"|\bnot\s+(?:\w+\s+){0,2}(?:match(?:ing|ed)?|equivalent|equal|identical|the\s+same|consistent|correct)\b"
+    r"|\bnon-equivalent\b|\bnot\s+within\s+(?:a\s+)?(?:small\s+)?(?:acceptable\s+)?margin\b"
+    r"|\bmismatch|\bdiffers?\b|\bdisagrees?\b|\bcontradict|\binconsistent\s+with\b"
+    r"|\b(?:is|are|was|were)\s+(?:\w+\s+)?(?:different|incorrect|wrong|inconsistent)\b"
+    r"|\bdifferent\s+(?:\w+\s+)?(?:from|answers?|values?|numbers?|choices?|options?|letters?|terms?|compounds?|results?|sequences?|quantit|express|formul|structure|molecule|name)"
+    r"|\b(?:meaningful|significant|substantive|noticeabl\w+)\s+(?:\w+\s+)?(?:discrepancy|difference|mismatch|deviation|inconsistency|non-equivalence)\b"
+    r"|\binconsistency(?:/ambiguity)?\s+with\b"
+    r"|\bdifference\s+is\s+(?:meaningful|significant|substantive)\b"
+    r"|\bno\s+(?:\w+\s+){0,2}match\b|\bfails?\s+to\s+match\b"
+    r"|\bcannot\s+be\s+considered\s+(?:equivalent|correct|a\s+match)"
+)
+_HLE_WORDING_YES_RE = re.compile(
+    r"\bmatch(?:es|ed|ing)?\b"
+    r"|\bidentical\b|\bequivalent\b|\bsame\s+as\b|\bconsistent\s+with\b|\bagrees?\s+with\b"
+    r"|\b(?:is|are)\s+(?:\w+\s+)?correct\b"
+    r"|\bequals?\s+(?:to\s+)?(?:the|\d)"
+    r"|\bcorrectly\b|\bcorresponds?\s+to\b"
+)
+
+
+def _hle_wording_verdict(text: str) -> str | None:
+    text = _HLE_WORDING_SKIP_LINE_RE.sub(" ", text)
+    text = _HLE_WORDING_NEUTRAL_RE.sub(" ", text)
+    if _HLE_WORDING_NO_RE.search(text):
+        return "no"
+    if _HLE_WORDING_YES_RE.search(text):
+        return "yes"
+    return None
+
+
+def parse_hle_judge(response: str) -> tuple[bool, bool]:
+    """Return (is_correct, verdict_unparsed) for an HLE judge reply."""
+    verdict = response.split("</think>")[-1].strip().lower()
+    matches = _HLE_CORRECT_RE.findall(verdict)
+    if not matches:
+        matches = _HLE_ANSWER_FALLBACK_RE.findall(verdict)
+    if matches:
+        return matches[-1] == "yes", False
+    wording = _hle_wording_verdict(verdict)
+    return wording == "yes", wording is None
+
+
+# Per-request timeout and attempt count: a thinking judge occasionally takes a long time, and one
+# slow request must never invalidate a whole batch (2500 items).
+DEFAULT_TIMEOUT = float(os.environ.get("EVAL_JUDGE_TIMEOUT", "600"))
+MAX_REQUEST_ATTEMPTS = int(os.environ.get("EVAL_JUDGE_ATTEMPTS", "3"))
+
+DEFAULT_MAX_TOKENS = 1024 * 64
+
+
+def truncate_output(text: str, max_paragraphs: int = 10, max_chars: int = 4000) -> str:
+    """Truncate the model output before judging: last 10 paragraphs, last 4000 characters."""
+    tail = text.split("</think>")[-1].strip()
+    return "\n\n".join(tail.split("\n\n")[-max_paragraphs:])[-max_chars:]
+
+
+def parse_judge_is_correct(response: str) -> bool:
+    """Parse a judge reply: only the last section after `</think>`."""
+    verdict = response.split("</think>")[-1].strip()
+    return verdict.strip().lower() == "correct"
+
+
+def resolve_model(model: str | None = None) -> str:
+    """The judge model to use: the explicit argument wins, then EVAL_JUDGE_MODEL."""
+    resolved = model or os.environ.get("EVAL_JUDGE_MODEL") or ""
+    if not resolved:
+        raise RuntimeError("missing judge model: set EVAL_JUDGE_MODEL")
+    return resolved
+
+
+class LLMJudge:
+    """Judge client for an OpenAI-compatible /chat/completions endpoint."""
+
+    def __init__(
+        self,
+        *,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        model: str | None = None,
+        temperature: float = 1.0,
+        top_p: float = 0.95,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        timeout: float = DEFAULT_TIMEOUT,
+        concurrency: int = 8,
+    ) -> None:
+        self.base_url = (base_url or os.environ.get("EVAL_JUDGE_BASE_URL") or "").rstrip("/")
+        self.api_key = api_key or os.environ.get("EVAL_JUDGE_API_KEY") or os.environ.get("DEEPSEEK_API_KEY") or ""
+        self.model = resolve_model(model)
+        self.temperature = temperature
+        self.top_p = top_p
+        self.max_tokens = max_tokens
+        self.timeout = timeout
+        self.concurrency = concurrency
+        if not self.base_url:
+            raise RuntimeError("missing judge base URL: set EVAL_JUDGE_BASE_URL")
+        if not self.api_key:
+            raise RuntimeError("missing judge API key: set EVAL_JUDGE_API_KEY (or DEEPSEEK_API_KEY)")
+
+    def _judge_one(
+        self,
+        question: str,
+        answer: str,
+        output: str,
+        kind: str = "equivalence",
+        max_tokens: int | None = None,
+    ) -> dict[str, Any]:
+        if kind == "hle":
+            content = HLE_JUDGE_PROMPT.format(question=question, correct_answer=answer, response=output)
+            parser = None
+        elif kind == "equivalence":
+            content = ANSWER_EXACT_MATCH_PROMPT.format(question=question, answer=answer, output=output)
+            parser = parse_judge_is_correct
+        else:
+            raise ValueError(f"unknown judge protocol: {kind!r}, expected one of {JUDGE_KINDS}")
+        payload = {
+            "model": self.model,
+            "messages": [{"role": "user", "content": content}],
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "max_tokens": max_tokens or self.max_tokens,
+            "stream": False,
+        }
+        last_err = ""
+        for attempt in range(1, MAX_REQUEST_ATTEMPTS + 1):
+            try:
+                response = requests.post(
+                    f"{self.base_url}/chat/completions",
+                    headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+                    data=json.dumps(payload),
+                    timeout=self.timeout,
+                )
+                response.raise_for_status()
+                choice = response.json()["choices"][0]
+                message = choice.get("message") or {}
+                # Read only `content`: a thinking model puts its draft in `reasoning_content`,
+                # and a phrase such as "correct: no" there must not be taken as the verdict.
+                text = (message.get("content") or "").strip()
+                if kind == "hle":
+                    judge_correct, verdict_unparsed = parse_hle_judge(text) if text else (False, True)
+                    return {
+                        "judge_correct": judge_correct,
+                        "judge_raw": text,
+                        "judge_parsed": bool(text),
+                        "judge_verdict_unparsed": verdict_unparsed,
+                        "judge_finish_reason": choice.get("finish_reason"),
+                    }
+                return {
+                    "judge_correct": bool(text) and parser(text),
+                    "judge_raw": text,
+                    "judge_parsed": bool(text),
+                    "judge_finish_reason": choice.get("finish_reason"),
+                }
+            except Exception as exc:  # network / timeout / 5xx land here; must not abort the batch
+                last_err = f"{type(exc).__name__}: {exc}"
+                if attempt < MAX_REQUEST_ATTEMPTS:
+                    time.sleep(min(2**attempt, 10))
+        return {
+            "judge_correct": False,
+            "judge_raw": "",
+            "judge_parsed": False,
+            "judge_error": last_err,
+        }
+
+    def judge_many(
+        self,
+        items: list[tuple[str, str, str]],
+        kind: str = "equivalence",
+        max_tokens: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """items = [(question, gold_answer, model_output), ...], judged concurrently."""
+
+        def _safe(item: tuple[str, str, str]) -> dict[str, Any]:
+            try:
+                return self._judge_one(*item, kind=kind, max_tokens=max_tokens)
+            except Exception as exc:  # last resort: an unexpected failure must not abort the batch
+                return {
+                    "judge_correct": False,
+                    "judge_raw": "",
+                    "judge_parsed": False,
+                    "judge_error": f"{type(exc).__name__}: {exc}",
+                }
+
+        with ThreadPoolExecutor(max_workers=self.concurrency) as pool:
+            return list(pool.map(_safe, items))
+
+
+def judge_answers(
+    items: list[tuple[str, str, str]],
+    *,
+    kind: str = "equivalence",
+    concurrency: int = 8,
+) -> list[dict[str, Any]]:
+    """Judge every (question, gold_answer, model_output) triple; one verdict per item."""
+    judge = LLMJudge(concurrency=concurrency)
+    verdicts = judge.judge_many(items, kind=kind)
+
+    failed = [verdict for verdict in verdicts if not verdict.get("judge_parsed")]
+    if failed:
+        errors = Counter(verdict["judge_error"] for verdict in failed if verdict.get("judge_error"))
+        print(f"[warn] {len(failed)}/{len(items)} judge replies carried no verdict; counted as wrong", flush=True)
+        for err, count in errors.most_common(5):
+            print(f"[warn]   {count} request(s) failed: {err}", flush=True)
+        if not errors:
+            print("[warn]   replies were empty; raise the judge max_tokens", flush=True)
+    return verdicts

eval/common/loader.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import importlib.util
+import sys
+from types import ModuleType
+
+from common.paths import EVAL_DIR
+
+SUPPORTED = ("aime", "hle", "gpqa")
+
+
+def load_bench(name: str) -> ModuleType:
+    if name not in SUPPORTED:
+        raise SystemExit(f"benchmark not wired up: {name} (supported: {' / '.join(SUPPORTED)})")
+    path = EVAL_DIR / name / "bench.py"
+    if not path.is_file():
+        raise SystemExit(f"benchmark implementation not found: {path}")
+    mod_name = f"eval_bench_{name}"
+    if mod_name in sys.modules:
+        return sys.modules[mod_name]
+    spec = importlib.util.spec_from_file_location(mod_name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[mod_name] = module
+    spec.loader.exec_module(module)
+    return module

eval/common/passk.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import math
+from collections import defaultdict
+from typing import Any
+
+
+def estimator(num_samples: int, num_correct: int, k: int) -> float:
+    """Unbiased Pass@k estimator: 1 - C(n-c, k) / C(n, k).
+
+    (The Codex/HumanEval estimator.)
+      * for k=1 it degenerates to c/n, the mean per-sample accuracy, i.e. the usual avg@n / Pass@1;
+      * when n-c < k (too few wrong samples for a k-subset to be all-correct) it returns 1.0.
+    """
+    assert 0 <= num_correct <= num_samples, f"invalid num_correct={num_correct} num_samples={num_samples}"
+    assert 0 <= k <= num_samples, f"invalid k={k} num_samples={num_samples}"
+    if num_samples - num_correct < k:
+        return 1.0
+    return float(1.0 - math.prod(1.0 - k / i for i in range(num_samples - num_correct + 1, num_samples + 1)))
+
+
+def pass_at_k_by_task(
+    results: list[dict[str, Any]],
+    k_list: list[int],
+    *,
+    task_key: str = "task_id",
+    correct_key: str = "correct",
+) -> dict[str, float]:
+    """Mean Pass@k over tasks, with tasks grouped by `task_key`.
+
+    When a task has fewer than k samples (judging failure, truncated sampling), approximate with
+    min(k, n) instead of dropping the task: dropping it would compute Pass@k only over the "clean"
+    tasks, and those are usually the easier ones, which biases the result upward.
+    """
+    grouped: dict[str, list[bool]] = defaultdict(list)
+    for row in results:
+        grouped[row[task_key]].append(bool(row[correct_key]))
+
+    metrics: dict[str, float] = {}
+    for k in k_list:
+        per_task = [estimator(len(flags), sum(flags), min(k, len(flags))) for flags in grouped.values() if flags]
+        metrics[f"Pass@{k}"] = sum(per_task) / len(per_task) if per_task else 0.0
+    return metrics
+
+
+def default_k_list(n_sample: int) -> list[int]:
+    """Default k values: 1, 2, 4, ... (powers of two), with n_sample appended."""
+    k_list = [2**i for i in range(max(math.ceil(math.log2(n_sample)), 1))]
+    if n_sample not in k_list:
+        k_list.append(n_sample)
+    return k_list
+
+
+def summarize(
+    results: list[dict[str, Any]],
+    *,
+    n_sample: int,
+    k_list: list[int] | None = None,
+    task_key: str = "task_id",
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """One rollout per row: {task_id, correct, ...}."""
+    k_list = k_list or default_k_list(n_sample)
+    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for row in results:
+        grouped[row[task_key]].append(row)
+
+    summary: dict[str, Any] = {
+        "num_tasks": len(grouped),
+        "num_rollouts": len(results),
+        "n_sample": n_sample,
+        "score": sum(1 for row in results if row["correct"]) / len(results) if results else 0.0,
+        "task_accuracy": {
+            task_id: sum(1 for row in task_rows if row["correct"]) / len(task_rows) for task_id, task_rows in sorted(grouped.items())
+        },
+    }
+    summary.update(pass_at_k_by_task(results, k_list, task_key=task_key))
+    if extra:
+        summary.update(extra)
+    return summary
+
+
+def print_report(summary: dict[str, Any], *, title: str = "") -> None:
+    print("\n" + "=" * 72)
+    if title:
+        print(title)
+    print(f"tasks={summary['num_tasks']}  rollouts={summary['num_rollouts']}  n_sample={summary['n_sample']}")
+    pass_keys = sorted(
+        (metric for metric in summary if metric.startswith("Pass@")),
+        key=lambda k: int(k.split("@")[1]),
+    )
+    for key in pass_keys:
+        print(f"  {key:<10s} = {summary[key] * 100:6.2f}")
+    note = "" if abs(summary["score"] - summary.get("Pass@1", 0.0)) < 1e-9 else "  (ragged sample counts, differs from Pass@1)"
+    print(
+        f"  {'score':<10s} = {summary['score'] * 100:6.2f}   (= mean per-sample accuracy, equals Pass@1 when sample counts are even){note}"
+    )
+    print("-" * 72)
+    print("per-task accuracy:")
+    for tid, acc in summary["task_accuracy"].items():
+        bar = "#" * int(round(acc * 20))
+        print(f"  {tid:<20s} {acc * 100:6.2f}  {bar}")
+    print("=" * 72)

eval/common/paths.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+COMMON_DIR = Path(__file__).resolve().parent
+EVAL_DIR = COMMON_DIR.parent
+DATA_DIR = EVAL_DIR / "data"
+DEFAULT_OUT_ROOT = EVAL_DIR / "outputs"
+
+
+def _looks_like_repo(path: Path) -> bool:
+    return (path / "examples" / "_config.py").is_file() and (path / "deepinfer").is_dir()
+
+
+def resolve_repo() -> Path:
+    """Locate the repository root.
+
+    By default it is inferred from the location of eval/; set `DEEPINFER_REPO=<repo root>` when the
+    directory layout is unusual.
+    """
+    candidates: list[Path] = []
+    if os.environ.get("DEEPINFER_REPO"):
+        candidates.append(Path(os.environ["DEEPINFER_REPO"]).expanduser())
+    candidates += [EVAL_DIR.parent, Path.cwd()]
+    for cand in candidates:
+        try:
+            if _looks_like_repo(cand):
+                return cand.resolve()
+        except OSError:
+            continue
+    print(
+        f"[eval] warning: could not infer the repository root from {EVAL_DIR} (no examples/ "
+        "and deepinfer/); set DEEPINFER_REPO=<absolute repo root>",
+        file=sys.stderr,
+    )
+    return EVAL_DIR.parent
+
+
+def ensure_sys_path() -> Path:
+    """Put the repository root and the eval directory on sys.path; return the repository root."""
+    repo = resolve_repo()
+    for entry in (str(repo), str(EVAL_DIR)):
+        if entry not in sys.path:
+            sys.path.insert(0, entry)
+    return repo
+
+
+def resolve_out_dir(out_dir: str | None, *, bench: str, tag: str) -> str:
+    return out_dir or str(DEFAULT_OUT_ROOT / bench / tag)

eval/gpqa/bench.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+EVAL_DIR = Path(__file__).resolve().parent.parent  # eval/
+if str(EVAL_DIR) not in sys.path:
+    sys.path.insert(0, str(EVAL_DIR))
+
+from common.io import require_file  # noqa: E402
+
+GPQA_PROMPT = (
+    "Answer the following multiple choice question. The last line of your response should be "
+    "of the following format: 'ANSWER: $LETTER' (without quotes) where LETTER is one of ABCD. "
+    "Think step by step before answering.\n\n{question}"
+)
+
+
+@dataclass(frozen=True)
+class Problem:
+    task_id: str
+    question: str
+    answer: str  # a single letter A-D
+    category: str = "unknown"
+
+
+# --------------------------------------------------------------------------- #
+# Data
+# --------------------------------------------------------------------------- #
+
+
+# The public mirror appends an answer-format instruction to each question, which conflicts with our
+# prompt template, so it is stripped uniformly.
+_TRAILING_INSTR_RE = re.compile(r"\n*\s*Please write your final answer in the form of.*$", re.I | re.S)
+
+
+def _strip_trailing_instruction(question: str) -> str:
+    return _TRAILING_INSTR_RE.sub("", question).strip()
+
+
+def _normalize_row(row: dict[str, Any], index: int) -> Problem | None:
+    """Accept both the parquet-derived format and the public-mirror format."""
+    question = row.get("question") or row.get("problem")
+    raw_answer = row.get("answer") or row.get("solution")
+    if question is None or raw_answer is None:
+        return None
+    # \boxed{D} / ANSWER: D / a bare D are all normalized to a single letter.
+    # Do not search for the first A-D: the "B" in "BOXED" inside \boxed{D} would match.
+    raw = str(raw_answer).strip()
+    match = (
+        re.search(r"\\boxed\{\s*\(?([A-D])\)?\s*\}", raw, re.I)
+        or re.search(r"(?i)answer\s*[:：]\s*\(?([A-D])\b", raw)
+        or re.fullmatch(r"\s*\(?([A-D])\)?\s*", raw)
+    )
+    if not match:
+        return None
+    task_id = str(row.get("task_id") or row.get("id") or f"gpqa_diamond_{index}")
+    category = str(row.get("category") or row.get("domain") or "unknown")
+    return Problem(
+        task_id=task_id,
+        question=_strip_trailing_instruction(str(question)),
+        answer=match.group(1),
+        category=category,
+    )
+
+
+def load(data_file: str | os.PathLike) -> list[Problem]:
+    path = Path(data_file)
+    require_file(
+        path,
+        "GPQA Diamond download: https://huggingface.co/datasets/Idavidrein/gpqa (gated)"
+        " or the public mirror https://huggingface.co/datasets/hendrydong/gpqa_diamond_mc",
+    )
+
+    problems: list[Problem] = []
+    skipped = 0
+    with open(path) as f:
+        for line_no, line in enumerate(f):
+            line = line.strip()
+            if not line:
+                continue
+            problem = _normalize_row(json.loads(line), line_no)
+            if problem is None:
+                skipped += 1
+                continue
+            problems.append(problem)
+    assert problems, f"no problem parsed from {path}"
+    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
+    if skipped:
+        print(f"[gpqa] skipped {skipped} lines (missing fields, or an answer that is not A-D)", flush=True)
+    return problems
+
+
+def build_question(problem: Problem) -> str:
+    return GPQA_PROMPT.format(question=problem.question.strip())
+
+
+def build_messages(problem: Problem) -> list[dict[str, Any]]:
+    return [{"role": "user", "content": build_question(problem)}]
+
+
+# --------------------------------------------------------------------------- #
+# Answer extraction and scoring
+# --------------------------------------------------------------------------- #
+
+_ANSWER_LINE_RE = re.compile(r"(?im)^\s*ANSWER\s*[:：]\s*\(?([A-D])\b")
+_BOXED_RE = re.compile(r"\\boxed\{\s*\(?([A-D])\)?\s*\}")
+_STANDALONE_RE = re.compile(r"(?<![A-Za-z])([A-D])(?![A-Za-z])")
+
+
+def extract_answer_with_source(text: str) -> tuple[str | None, str]:
+    """Extract the answer letter by priority and report which path matched."""
+    if not text:
+        return None, "none"
+    matches = _ANSWER_LINE_RE.findall(text)
+    if matches:
+        return matches[-1].upper(), "answer_line"
+    matches = _BOXED_RE.findall(text)
+    if matches:
+        return matches[-1].upper(), "boxed"
+    tail = text.rsplit("</think>", 1)[-1] if "</think>" in text else text[-400:]
+    matches = _STANDALONE_RE.findall(tail)
+    if matches:
+        return matches[-1].upper(), "tail_letter"
+    return None, "none"
+
+
+def is_correct(pred: str | None, gold: str) -> bool:
+    return pred is not None and pred.upper() == str(gold).strip().upper()
+
+
+# --------------------------------------------------------------------------- #
+# Entry points used by run_eval.py / score.py
+# --------------------------------------------------------------------------- #
+
+
+def build_records(problems: list[Problem], *, limit: int | None = None) -> list[dict[str, Any]]:
+    records = []
+    for problem in problems:
+        records.append(
+            {
+                "task_id": problem.task_id,
+                "question": problem.question,
+                "answer": problem.answer,
+                "category": problem.category,
+                "judge_question": build_question(problem),
+                "messages": build_messages(problem),
+            }
+        )
+        if limit is not None and len(records) >= limit:
+            break
+    return records
+
+
+def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
+    pred, source = extract_answer_with_source(rollout.get("raw_text", ""))
+    row = dict(rollout)
+    row["pred"] = pred
+    row["extract_source"] = source
+    row["gold"] = record["answer"]
+    row["correct"] = is_correct(pred, record["answer"])
+    return row

eval/gpqa/run.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+set -euo pipefail
+
+EVAL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO="$(cd "${EVAL_DIR}/.." && pwd)"
+cd "${REPO}"
+
+MODEL="${MODEL:-deepseek-ai/DeepSeek-V4.1-Flash-1017}"
+EP="${EP:-8}"
+N_SAMPLE="${N_SAMPLE:-8}"
+REASONING_EFFORT="${REASONING_EFFORT:-100}"
+TAG="${TAG:-${MODEL//\//_}-n${N_SAMPLE}-j${REASONING_EFFORT}}"
+DATA_FILE="${DATA_FILE:-${EVAL_DIR}/data/gpqa_diamond.jsonl}"
+OUT="${EVAL_DIR}/outputs/gpqa/${TAG}"
+
+echo "[gpqa] model=${MODEL} ep=${EP} n_sample=${N_SAMPLE} effort=${REASONING_EFFORT} -> ${OUT}"
+
+# 1) inference
+python eval/run_eval.py --bench gpqa --model-name "${MODEL}" --ep "${EP}" \
+    --n-sample "${N_SAMPLE}" --reasoning-effort "${REASONING_EFFORT}" \
+    --tag "${TAG}" --data-file "${DATA_FILE}"
+
+# 2) evaluate
+# LLM judging: need EVAL_JUDGE_BASE_URL / EVAL_JUDGE_API_KEY / EVAL_JUDGE_MODEL
+python eval/score.py --out-dir "${OUT}" --judge llm --judge-concurrency 16
+
+# rule judging: use this instead when no LLM judge is available
+# python eval/score.py --out-dir "${OUT}" --judge rule
+
+echo "done: ${OUT}/summary.json"

eval/hle/bench.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+EVAL_DIR = Path(__file__).resolve().parent.parent  # eval/
+if str(EVAL_DIR) not in sys.path:
+    sys.path.insert(0, str(EVAL_DIR))
+
+from common.io import require_file  # noqa: E402
+from common.paths import DATA_DIR  # noqa: E402
+
+# The official three-part answer format from hle_eval
+HLE_ANSWER_FORMAT = (
+    "{question}\n"
+    "Your response should be in the following format:\n"
+    "Explanation: {{your explanation for your final answer}}\n"
+    "Answer: {{your succinct, final answer}}\n"
+    "Confidence: {{your confidence score between 0% and 100% for your answer}}"
+)
+
+
+@dataclass(frozen=True)
+class Problem:
+    task_id: str
+    question: str
+    answer: str
+    category: str = "unknown"
+    has_image: bool = False
+    image_paths: tuple[str, ...] = ()
+    image_sizes: tuple[tuple[int, int], ...] = ()
+
+
+# --------------------------------------------------------------------------- #
+# Data
+# --------------------------------------------------------------------------- #
+
+
+_IMAGE_MAGIC = (b"\xff\xd8\xff", b"\x89PNG", b"GIF8", b"RIFF", b"BM")
+
+
+def _resolve_images(row: dict[str, Any]) -> tuple[tuple[str, ...], tuple[tuple[int, int], ...]]:
+    """Resolve a row's image into (local file path, (W, H)).
+
+    Three representations are supported:
+      1. multi_modal_data.image_urls + image_sizes (the full-dataset JSONL uses this form)
+      2. an `image` field holding a local path
+      3. an `image` field holding base64 (common in the official cais/hle parquet), decoded and
+         cached under data/hle_images/
+    """
+    import base64
+
+    multi_modal = row.get("multi_modal_data") or {}
+    urls = [str(url) for url in (multi_modal.get("image_urls") or [])]
+    sizes = [(int(size[0]), int(size[1])) for size in (multi_modal.get("image_sizes") or []) if len(size) >= 2]
+
+    raw = row.get("image")
+    if not urls and raw:
+        if isinstance(raw, (bytes, bytearray)):
+            urls = [_write_cached_image(row, bytes(raw), 0)]
+        else:
+            text = str(raw).strip()
+            if os.path.isfile(text):
+                urls = [text]
+            elif text:
+                try:
+                    data = base64.b64decode(text, validate=False)
+                except Exception:
+                    data = b""
+                if any(data.startswith(magic) for magic in _IMAGE_MAGIC):
+                    urls = [_write_cached_image(row, data, 0)]
+
+    if urls and not sizes:
+        from PIL import Image  # PIL is only needed to recover the dimensions from a file
+
+        sizes = [Image.open(path).size for path in urls]
+    return tuple(urls), tuple(sizes)
+
+
+def _write_cached_image(row: dict[str, Any], data: bytes, index: int) -> str:
+    cache = DATA_DIR / "hle_images"
+    cache.mkdir(parents=True, exist_ok=True)
+    stem = str(row.get("task_id") or row.get("id") or "hle")
+    path = cache / f"{stem}_{index}.img"
+    if not path.is_file() or path.stat().st_size != len(data):
+        path.write_bytes(data)
+    return str(path)
+
+
+def _has_image(row: dict[str, Any]) -> bool:
+    """Accept several common representations: an explicit modality field, multi_modal_data.image_urls, or a single image field."""
+    if str(row.get("modality") or "").lower() == "image":
+        return True
+    multi_modal = row.get("multi_modal_data") or {}
+    if multi_modal.get("image_urls"):
+        return True
+    img = row.get("image")
+    if img is None:
+        return False
+    if isinstance(img, (bytes, bytearray)):
+        return len(img) > 0
+    return bool(str(img).strip())
+
+
+def load(data_file: str | os.PathLike) -> list[Problem]:
+    path = Path(data_file)
+    require_file(path, "HLE download: https://huggingface.co/datasets/cais/hle (gated; accept the terms on Hugging Face first)")
+
+    problems: list[Problem] = []
+    with open(path) as f:
+        for line_no, line in enumerate(f):
+            line = line.strip()
+            if not line:
+                continue
+            row = json.loads(line)
+            has_image = _has_image(row)
+            image_paths, image_sizes = _resolve_images(row) if has_image else ((), ())
+            problems.append(
+                Problem(
+                    task_id=str(row.get("task_id") or row.get("id") or f"hle_{line_no:04d}"),
+                    question=str(row["question"]),
+                    answer=str(row["answer"]),
+                    category=str(row.get("category") or row.get("raw_subject") or "unknown"),
+                    has_image=has_image,
+                    image_paths=image_paths,
+                    image_sizes=image_sizes,
+                )
+            )
+    assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
+    return problems
+
+
+def build_question(problem: Problem) -> str:
+    return HLE_ANSWER_FORMAT.format(question=problem.question.strip())
+
+
+def build_messages(problem: Problem) -> list[dict[str, Any]]:
+    return [{"role": "user", "content": build_question(problem)}]
+
+
+# --------------------------------------------------------------------------- #
+# Answer extraction and (rule-based) scoring
+# --------------------------------------------------------------------------- #
+
+_ANSWER_LINE_RE = re.compile(r"(?im)^\s*(?:\*\*)?answer(?:\*\*)?\s*[:：]\s*(.+?)\s*$")
+
+
+def extract_answer(text: str) -> str | None:
+    """Return the content of the last `Answer:` line (the official three-part format)."""
+    if not text:
+        return None
+    matches = _ANSWER_LINE_RE.findall(text)
+    if not matches:
+        return None
+    ans = matches[-1].strip()
+    ans = ans.strip("*").strip()
+    ans = re.sub(r"^\$+|\$+$", "", ans).strip()
+    return ans or None
+
+
+def normalize_answer(ans: str | None) -> str | None:
+    if ans is None:
+        return None
+    text = ans.strip()
+    text = text.strip("$").strip()
+    text = text.replace("\\left", "").replace("\\right", "")
+    text = text.replace("\\!", "").replace("\\,", "").replace("\\;", "")
+    text = text.replace("{,}", ",") if "{,}" in text else text
+    text = re.sub(r"\s+", " ", text).strip().lower()
+    text = text.rstrip(".").strip()
+    if re.fullmatch(r"-?\d+", text):
+        return str(int(text))
+    return text
+
+
+def is_correct(pred: str | None, gold: str) -> bool:
+    """Rule-based scoring: normalized string equality. **Smoke tests only**; report LLM-judge numbers."""
+    pred_norm, gold_norm = normalize_answer(pred), normalize_answer(gold)
+    return pred_norm is not None and gold_norm is not None and pred_norm == gold_norm
+
+
+# --------------------------------------------------------------------------- #
+# Entry points used by run_eval.py / score.py
+# --------------------------------------------------------------------------- #
+
+
+def build_records(
+    problems: list[Problem],
+    *,
+    limit: int | None = None,
+) -> list[dict[str, Any]]:
+    records = []
+    for problem in problems:
+        records.append(
+            {
+                "task_id": problem.task_id,
+                "question": problem.question,
+                "answer": problem.answer,
+                "category": problem.category,
+                "modality": "image" if problem.has_image else "text",
+                "image_paths": list(problem.image_paths),
+                "image_sizes": [list(size) for size in problem.image_sizes],
+                "judge_question": problem.question,
+                "messages": build_messages(problem),
+            }
+        )
+        if limit is not None and len(records) >= limit:
+            break
+    return records
+
+
+def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
+    row = dict(rollout)
+    row["pred"] = extract_answer(rollout.get("raw_text", ""))
+    row["gold"] = record["answer"]
+    row["correct"] = is_correct(row["pred"], record["answer"])
+    return row

eval/hle/run.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -euo pipefail
+
+EVAL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO="$(cd "${EVAL_DIR}/.." && pwd)"
+cd "${REPO}"
+
+MODEL="${MODEL:-deepseek-ai/DeepSeek-V4.1-Flash-1017}"
+EP="${EP:-8}"
+REASONING_EFFORT="${REASONING_EFFORT:-100}"
+TAG="${TAG:-${MODEL//\//_}-j${REASONING_EFFORT}}"
+DATA_FILE="${DATA_FILE:-${EVAL_DIR}/data/hle_test.jsonl}"
+OUT="${EVAL_DIR}/outputs/hle/${TAG}"
+
+echo "[hle] model=${MODEL} ep=${EP} effort=${REASONING_EFFORT} -> ${OUT}"
+
+# 1) inference (single sample per problem, temperature from BENCH_DEFAULTS)
+python eval/run_eval.py --bench hle --model-name "${MODEL}" --ep "${EP}" \
+    --n-sample 1 --reasoning-effort "${REASONING_EFFORT}" \
+    --tag "${TAG}" --data-file "${DATA_FILE}"
+
+# 2) evaluate (LLM judge only)
+# need EVAL_JUDGE_BASE_URL / EVAL_JUDGE_API_KEY / EVAL_JUDGE_MODEL
+python eval/score.py --out-dir "${OUT}" --judge llm --judge-concurrency 16
+
+echo "done: ${OUT}/summary.json"

eval/run_eval.py

@@ -0,0 +1,309 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+EVAL_DIR = Path(__file__).resolve().parent
+if str(EVAL_DIR) not in sys.path:
+    sys.path.insert(0, str(EVAL_DIR))
+
+from common.io import write_jsonl  # noqa: E402
+from common.loader import SUPPORTED, load_bench  # noqa: E402
+from common.paths import ensure_sys_path, resolve_out_dir  # noqa: E402
+
+REPO = ensure_sys_path()  # put repo root on sys.path so examples / deepinfer can be imported
+
+from examples._config import DEFAULT_MODEL_NAME  # noqa: E402
+
+BENCH_DEFAULTS = {
+    "aime": {"n_sample": 32, "temperature": 1.0, "top_p": 1.0},
+    "hle": {"n_sample": 1, "temperature": 1.0, "top_p": 1.0},
+    "gpqa": {"n_sample": 8, "temperature": 1.0, "top_p": 1.0},
+}
+
+
+_REASONING_EFFORT_LEVELS = ("low", "high", "max")
+
+
+def _reasoning_effort(value: str) -> int | str:
+    """Accept an integer 1-100, or a named level low/high/max."""
+    if value in _REASONING_EFFORT_LEVELS:
+        return value
+    try:
+        effort = int(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"expected an integer 1-100 or one of {_REASONING_EFFORT_LEVELS}, got {value!r}") from None
+    if not 1 <= effort <= 100:
+        raise argparse.ArgumentTypeError(f"expected 1-100, got {effort}")
+    return effort
+
+
+parser = argparse.ArgumentParser(description="eval benchmark")
+# benchmark setting
+parser.add_argument("--bench", default="aime", choices=list(SUPPORTED))
+parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME, help="HF model id or HF ckpt directory")
+parser.add_argument("--data-file", required=True, help="problem jsonl")
+parser.add_argument("--limit", type=int, default=None, help="run only the first N problems")
+parser.add_argument("--out-dir", default=None)
+parser.add_argument("--tag", default=None, help="output dir suffix, to compare settings of the same model")
+# sample param
+parser.add_argument("--n-sample", type=int, default=None, help="samples per problem (default: aime 32 / hle 1 / gpqa 8)")
+parser.add_argument("--temperature", type=float, default=None, help="default 1.0 for every benchmark")
+parser.add_argument("--top-p", type=float, default=None, help="default 1.0")
+parser.add_argument(
+    "--reasoning-effort", type=_reasoning_effort, default=100, help="integer 1-100, or low/high/max; higher means more thinking"
+)
+parser.add_argument("--max-len", type=int, default=1024 * 144, help="max prompt + output length per request")
+# engine
+parser.add_argument("--ep", type=int, default=8, help="expert parallel size = total GPUs in this job")
+parser.add_argument("--cp", type=int, default=1)
+parser.add_argument("--dp", type=int, default=None, help="defaults to ep // cp")
+parser.add_argument("--max-num-seqs", type=int, default=384)
+parser.add_argument("--max-num-batched-tokens", type=int, default=32768)
+parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
+parser.add_argument(
+    "--env", action="append", default=[], metavar="KEY=VALUE", help="extra env vars, repeatable; applied before the engine is built"
+)
+parser.add_argument("--decode-blacklist", default="0,2,129262..129264", help="tokens blocked during generation; empty string disables it")
+parser.add_argument(
+    "--batch-invariant", action="store_true", help="DEEPINFER_BATCH_INVARIANT=1: results independent of batch composition (slower)"
+)
+
+args = parser.parse_args()
+
+# fill in per-benchmark defaults (only for args not passed explicitly)
+for _key, _default in BENCH_DEFAULTS[args.bench].items():
+    if getattr(args, _key) is None:
+        setattr(args, _key, _default)
+
+REASONING_EFFORT = args.reasoning_effort
+
+OUT_DIR = resolve_out_dir(
+    args.out_dir,
+    bench=args.bench,
+    tag=args.tag or f"{args.model_name.replace('/', '_')}-n{args.n_sample}-j{args.reasoning_effort}",
+)
+META_PATH = os.path.join(OUT_DIR, "meta.json")
+
+
+# --------------------------------------------------------------------------- #
+# generation
+# --------------------------------------------------------------------------- #
+
+
+def _render_prompts(encoder: ModuleType, records: list[dict]) -> list[str]:
+    return [
+        encoder.encode_messages(
+            rec["messages"],
+            thinking_mode="thinking",
+            reasoning_effort=REASONING_EFFORT,
+        )
+        for rec in records
+    ]
+
+
+def _multi_modal_entry(record: dict) -> dict:
+    """One multi_modal_data entry: image bytes plus dimensions."""
+    return {
+        "image_bytes": [Path(path).read_bytes() for path in record.get("image_paths") or []],
+        "image_sizes": [tuple(size) for size in record.get("image_sizes") or []],
+    }
+
+
+def _setup_env() -> None:
+    # MTP / DFlash speculative decoding, set before setup_common_env (which only setdefaults it).
+    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "5"
+    os.environ["DEEPINFER_ENABLE_MTP_VARLEN"] = "1"
+
+    from examples._config import setup_common_env
+
+    setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=True, enable_eplb=True)
+    for item in args.env:
+        key, _, value = item.partition("=")
+        assert key and "=" in item, f"--env expects KEY=VALUE, got {item!r}"
+        os.environ[key] = value
+    if args.batch_invariant:
+        os.environ["DEEPINFER_BATCH_INVARIANT"] = "1"
+    if args.decode_blacklist:
+        os.environ["DEEPINFER_DECODE_TOKEN_BLACKLIST"] = args.decode_blacklist
+
+
+def _seed(global_index: int, sample_index: int) -> int:
+    """Per-sample seed: identical seeds produce identical output, so every sample needs its own."""
+    return (global_index * args.n_sample + sample_index) % (2**64)
+
+
+def _rows_from_outputs(tokenizer: Any, plan: list, records: list[dict], outputs: list) -> list[dict]:
+    rows = []
+    for (problem, sample), generated in zip(plan, outputs, strict=True):
+        assert generated, f"{records[problem]['task_id']}#{sample}: generated no tokens"
+        text = tokenizer.decode([token.token_id for token in generated], skip_special_tokens=False, clean_up_tokenization_spaces=False)
+        rows.append(
+            {
+                "task_id": records[problem]["task_id"],
+                "sample_idx": sample,
+                "seq_seed": _seed(problem, sample),
+                "prompt_len": generated[0].prompt_len,
+                "num_output_tokens": len(generated),
+                "finish_reason": generated[-1].finish_reason,
+                "raw_text": text,
+            }
+        )
+    return rows
+
+
+def run_rank(dp_size: int, dp_rank: int) -> None:
+    _setup_env()
+
+    from common.encoder import load_encoder
+
+    from examples._config import create_llm
+
+    encoder = load_encoder(args.model_name)
+
+    bench = load_bench(args.bench)
+    problems = bench.load(args.data_file)
+    records = bench.build_records(problems, limit=args.limit)
+
+    # Shard by sample so every rank gets an equal share and no rank inherits a problem's long tail.
+    plan = [(problem, sample) for problem in range(len(records)) for sample in range(args.n_sample)][dp_rank::dp_size]
+
+    llm = create_llm(
+        args.model_name,
+        dp_rank,
+        ep_size=args.ep,
+        cp_size=args.cp,
+        data_parallel_size=dp_size,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        # The engine's runtime_max_model_len is max_model_len + 1, and every request's
+        # prompt + max_tokens is already clamped to <= max_len, so a tiny margin is enough.
+        max_model_len=args.max_len + 64,
+        max_num_seqs=args.max_num_seqs,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        lookahead_tokens=min(1, int(os.environ.get("DEEPINFER_USE_NEXTN_DECODE", "0"))),
+        decode_sequence=False,
+    )
+
+    try:
+        # Import after create_llm: deepinfer imports torch, which must not happen before the engine
+        # has set up the accelerator devices.
+        from deepinfer.sampling_params import SamplingParams
+
+        tokenizer = llm.get_tokenizer()
+
+        if not plan:
+            # dp_size > number of requests leaves some ranks with an empty shard: write an empty
+            # file and stop, do not build a batch.
+            os.makedirs(OUT_DIR, exist_ok=True)
+            write_jsonl(os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl"), [])
+            print(f"[rank {dp_rank}] empty shard (dp_size={dp_size} > request count), skipping", flush=True)
+            return
+
+        # Render and tokenize each needed prompt once: a rank may use the same prompt for several
+        # samples of the same problem.
+        needed = sorted({problem for problem, _ in plan})
+        texts = dict(zip(needed, _render_prompts(encoder, [records[problem] for problem in needed]), strict=True))
+        has_images = any(records[problem].get("image_paths") for problem in needed)
+        multi_modal: dict[int, dict] = {}
+        if has_images:
+            # Replace the <image> placeholder with the image token the engine understands. Image
+            # questions then take the string path and the engine expands the token.
+            from examples._prompts import get_engine_image_token, replace_image_tokens
+
+            image_token_str, _ = get_engine_image_token(llm)
+            texts = dict(zip(needed, replace_image_tokens([texts[problem] for problem in needed], image_token_str), strict=True))
+            # The engine consumes multi_modal_data in place, so every request gets its own copy.
+            multi_modal = {problem: _multi_modal_entry(records[problem]) for problem in needed}
+        token_ids = {problem: tokenizer.encode(text, add_special_tokens=False) for problem, text in texts.items()}
+
+        # max_tokens is clamped to the remaining budget: max_len - prompt_len, at least 1.
+        sampling_params = [
+            SamplingParams(
+                temperature=args.temperature,
+                top_p=args.top_p,
+                max_tokens=max(args.max_len - len(token_ids[problem]), 1),
+                ignore_eos=False,
+                seq_seed=_seed(problem, sample),
+            )
+            for problem, sample in plan
+        ]
+
+        lens = [len(token_ids[problem]) for problem, _ in plan]
+        print(
+            f"[rank {dp_rank}] problems={len(needed)} requests={len(plan)} "
+            f"prompt_tokens[min/avg/max]={min(lens)}/{sum(lens) // len(lens)}/{max(lens)}"
+            + (f" multimodal={len(plan)}" if has_images else ""),
+            flush=True,
+        )
+
+        if has_images:
+            outputs = llm.generate(
+                prompts=[texts[problem] for problem, _ in plan],
+                multi_modal_data=[dict(multi_modal[problem]) for problem, _ in plan],
+                sampling_params=sampling_params,
+                use_tqdm=False,
+            )
+        else:
+            outputs = llm.generate(
+                prompts=None,
+                prompt_token_ids=[token_ids[problem] for problem, _ in plan],
+                sampling_params=sampling_params,
+                use_tqdm=False,
+            )
+
+        rows = _rows_from_outputs(tokenizer, plan, records, outputs)
+
+        os.makedirs(OUT_DIR, exist_ok=True)
+        path = os.path.join(OUT_DIR, f"raw_rank{dp_rank:03d}.jsonl")
+        write_jsonl(path, rows)
+        print(f"[rank {dp_rank}] wrote {path} rows={len(rows)}", flush=True)
+    finally:
+        llm.close()
+
+
+def main() -> int:
+    dp_size = max(args.ep // args.cp, 1) if args.dp is None else args.dp
+    os.makedirs(OUT_DIR, exist_ok=True)
+
+    meta = {
+        "bench": args.bench,
+        "model_name": args.model_name,
+        "ep": args.ep,
+        "cp": args.cp,
+        "dp": dp_size,
+        "data_file": args.data_file,
+        "limit": args.limit,
+        "n_sample": args.n_sample,
+        "temperature": args.temperature,
+        "top_p": args.top_p,
+        "reasoning_effort": REASONING_EFFORT,
+        "max_len": args.max_len,
+        "mtp": 5,
+        "mtp_varlen": True,
+        "batch_invariant": args.batch_invariant,
+        "decode_blacklist": args.decode_blacklist,
+        "max_num_seqs": args.max_num_seqs,
+        "max_num_batched_tokens": args.max_num_batched_tokens,
+        "env": args.env,
+    }
+
+    for old in Path(OUT_DIR).glob("raw_rank*.jsonl"):
+        old.unlink()
+    with open(META_PATH, "w") as f:
+        json.dump(meta, f, indent=2, ensure_ascii=False)
+
+    print(f"REPO={REPO}\nOUT_DIR={OUT_DIR}\nmeta={json.dumps(meta, ensure_ascii=False)}", flush=True)
+
+    from examples._launch import launch_dp
+
+    launch_dp(run_rank, dp_size)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

eval/score.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import sys
+from collections import Counter
+from pathlib import Path
+
+EVAL_DIR = Path(__file__).resolve().parent
+if str(EVAL_DIR) not in sys.path:
+    sys.path.insert(0, str(EVAL_DIR))
+
+from common.io import read_jsonl, write_jsonl  # noqa: E402
+from common.loader import load_bench  # noqa: E402
+from common.passk import default_k_list, print_report, summarize  # noqa: E402
+from common.paths import ensure_sys_path  # noqa: E402
+
+ensure_sys_path()
+
+
+# Default judge per benchmark: all three use the LLM judge, matching the reference protocol.
+BENCH_DEFAULT_JUDGE = {"aime": "llm", "hle": "llm", "gpqa": "llm"}
+
+
+def load_raw(out_dir: str) -> list[dict]:
+    """Merge every rank's raw output; for a given (task_id, sample_idx) the last write wins."""
+    merged: dict[tuple, dict] = {}
+    order: list[tuple] = []
+    paths = sorted(glob.glob(os.path.join(out_dir, "raw_rank*.jsonl")))
+    if not paths:
+        raise FileNotFoundError(f"no raw_rank*.jsonl under {out_dir} (did generation finish?)")
+    for path in paths:
+        for row in read_jsonl(path):
+            key = (row["task_id"], row["sample_idx"])
+            if key not in merged:
+                order.append(key)
+            merged[key] = row
+    return [merged[key] for key in order]
+
+
+def _judge(results: list[dict], by_id: dict, bench_name: str, judge: str, concurrency: int) -> str | None:
+    """Score `results` in place. Returns the judge model name for the LLM judge, else None."""
+    if judge == "rule" and bench_name == "hle":
+        raise SystemExit(
+            "HLE does not support rule judging: the official protocol requires an LLM judge.\n"
+            "Most HLE answers are free text, so normalized exact match marks equivalent-but-"
+            "differently-written answers wrong and systematically understates the score.\n"
+            "Set EVAL_JUDGE_BASE_URL / EVAL_JUDGE_API_KEY / EVAL_JUDGE_MODEL and use --judge llm."
+        )
+    if judge != "llm":
+        for row in results:
+            row["rule_correct"] = row["correct"]
+        return None
+
+    from common.judge import judge_answers, resolve_model, truncate_output
+
+    judge_kind = "hle" if bench_name == "hle" else "equivalence"
+
+    def judge_input(text: str) -> str:
+        # HLE official protocol: drop the thinking section, send the rest in full, no truncation.
+        if judge_kind == "hle":
+            return text.split("</think>")[-1].strip()
+        return truncate_output(text)
+
+    items = [
+        (
+            by_id[row["task_id"]].get("judge_question") or by_id[row["task_id"]]["question"],
+            by_id[row["task_id"]]["answer"],
+            judge_input(row["raw_text"]),
+        )
+        for row in results
+    ]
+    for row, verdict in zip(results, judge_answers(items, kind=judge_kind, concurrency=concurrency), strict=True):
+        row["rule_correct"] = row["correct"]
+        row["correct"] = bool(verdict["judge_correct"])
+        row["judge_raw"] = verdict.get("judge_raw", "")
+        row["judge_parsed"] = bool(verdict.get("judge_parsed", True))
+        if "judge_verdict_unparsed" in verdict:
+            row["judge_verdict_unparsed"] = bool(verdict["judge_verdict_unparsed"])
+        if verdict.get("judge_error"):
+            row["judge_error"] = verdict["judge_error"]
+    return resolve_model()
+
+
+def score_run(
+    out_dir: str,
+    *,
+    judge: str = "auto",
+    judge_concurrency: int = 8,
+    data_file: str | None = None,
+    verbose: bool = True,
+) -> int:
+    meta_path = os.path.join(out_dir, "meta.json")
+    meta: dict = {}
+    if os.path.isfile(meta_path):
+        with open(meta_path) as f:
+            meta = json.load(f)
+
+    bench_name = meta.get("bench", "aime")
+    bench = load_bench(bench_name)
+    if judge == "auto":
+        judge = BENCH_DEFAULT_JUDGE.get(bench_name, "rule")
+    data_path = data_file or meta.get("data_file")
+    if not data_path:
+        raise SystemExit(f"{meta_path} has no data_file; pass --data-file explicitly")
+    problems = bench.load(data_path)
+    records = bench.build_records(problems, limit=meta.get("limit"))
+    by_id = {r["task_id"]: r for r in records}
+
+    rollouts = load_raw(out_dir)
+    unknown = [r["task_id"] for r in rollouts if r["task_id"] not in by_id]
+    assert not unknown, f"raw output has task_ids that are not in the dataset: {sorted(set(unknown))[:5]}"
+
+    results: list[dict] = []
+    for rollout in rollouts:
+        record = by_id[rollout["task_id"]]
+        row = bench.score_rollout(rollout, record)
+        results.append(row)
+
+    judge_model = _judge(results, by_id, bench_name, judge, judge_concurrency)
+
+    if not results:
+        raise SystemExit(f"no rollout at all in {out_dir}/raw_rank*.jsonl; refusing to emit an empty score")
+
+    # Coverage: missing rank files, problems or samples must be surfaced, never silently ignored.
+    seen = {row["task_id"] for row in results}
+    n_sample = int(meta.get("n_sample", max((row["sample_idx"] for row in results), default=0) + 1))
+    missing = sorted(set(by_id) - seen)
+    expected_rollouts = len(records) * n_sample
+    coverage_complete = not missing and len(results) == expected_rollouts
+    if not coverage_complete:
+        print(
+            f"[warn] incomplete coverage: {len(seen)}/{len(records)} tasks, {len(results)}/{expected_rollouts} "
+            "rollouts; this is not a complete result",
+            flush=True,
+        )
+
+    # Truncated rollouts and fallback extraction paths both depress the score, so surface them.
+    finish_reasons = dict(Counter(str(row["finish_reason"]) for row in results if row.get("finish_reason")))
+    extract_sources = dict(Counter(str(row["extract_source"]) for row in results if row.get("extract_source")))
+    truncated = sum(finish_reasons.get(reason, 0) for reason in ("length", "prompt_length_exceeded"))
+    invalid = sum(finish_reasons.get(reason, 0) for reason in ("ignored", "aborted", "error", "preempted"))
+
+    summary = summarize(
+        results,
+        n_sample=n_sample,
+        k_list=meta.get("k_list") or default_k_list(n_sample),
+        extra={
+            "bench": meta.get("bench", "aime"),
+            "model_name": meta.get("model_name"),
+            "judge": judge,
+            "judge_model": judge_model,
+            "temperature": meta.get("temperature"),
+            "top_p": meta.get("top_p"),
+            "reasoning_effort": meta.get("reasoning_effort"),
+            "max_len": meta.get("max_len"),
+            "extract_failed": sum(1 for row in results if row.get("extract_failed", row.get("pred") is None)) / len(results),
+            "expected_tasks": len(records),
+            "expected_rollouts": expected_rollouts,
+            "coverage_complete": coverage_complete,
+            "missing_tasks": len(missing),
+            "ragged_samples": len({sum(1 for row in results if row["task_id"] == task_id) for task_id in seen}) > 1,
+            "finish_reasons": finish_reasons,
+            "extract_sources": extract_sources,
+            "truncated_rollouts": truncated,
+            "invalid_rollouts": invalid,
+        },
+    )
+    groups: dict[str, list[dict]] = {}
+    for row in results:
+        if key := str(by_id[row["task_id"]].get("modality") or ""):
+            groups.setdefault(key, []).append(row)
+    if len(groups) > 1:
+        summary["by_modality"] = {
+            modality: {
+                "num_tasks": len({row["task_id"] for row in rows}),
+                "num_rollouts": len(rows),
+                "Pass@1": sum(1 for row in rows if row["correct"]) / len(rows),
+            }
+            for modality, rows in sorted(groups.items())
+        }
+
+    if judge == "llm":
+        agree = sum(1 for row in results if bool(row["rule_correct"]) == bool(row["correct"]))
+        summary["rule_vs_llm_agreement"] = agree / len(results)
+        unparsed = sum(1 for row in results if not row.get("judge_parsed", True))
+        summary["judge_unparsed"] = unparsed / len(results)
+        verdict_unparsed = sum(1 for row in results if row.get("judge_verdict_unparsed"))
+        summary["judge_verdict_unparsed"] = verdict_unparsed / len(results)
+        if verdict_unparsed:
+            print(
+                f"[warn] {verdict_unparsed}/{len(results)} judge replies had no extractable verdict; "
+                "all three fallbacks missed, they count as wrong",
+                flush=True,
+            )
+        if unparsed:
+            print(
+                f"[warn] {unparsed}/{len(results)} judge calls returned no verdict and count as wrong; "
+                "this depresses the score -- raise the judge max_tokens and re-run",
+                flush=True,
+            )
+
+    os.makedirs(out_dir, exist_ok=True)
+    write_jsonl(os.path.join(out_dir, "scored.jsonl"), results)
+    with open(os.path.join(out_dir, "summary.json"), "w") as f:
+        json.dump(summary, f, indent=2, ensure_ascii=False)
+
+    if verbose:
+        print_report(summary, title=f"{summary['bench']}  {summary.get('model_name')}  judge={judge}")
+        print(
+            f"coverage: expected {summary['expected_rollouts']} rollouts / "
+            f"{summary['expected_tasks']} tasks, got {summary['num_rollouts']} / {summary['num_tasks']}, "
+            f"complete={summary['coverage_complete']}"
+        )
+        print(f"finish_reason={summary['finish_reasons']}")
+        for key, stat in (summary.get("by_modality") or {}).items():
+            print(f"  Pass@1[{key}] = {stat['Pass@1'] * 100:6.2f}  (tasks={stat['num_tasks']})")
+        if summary["extract_sources"]:
+            print(f"extract_source={summary['extract_sources']}")
+        if summary["truncated_rollouts"] or summary["invalid_rollouts"]:
+            print(
+                f"[warn] {summary['truncated_rollouts']} truncated / "
+                f"{summary['invalid_rollouts']} invalid rollouts, all counted as wrong; "
+                "this depresses the score"
+            )
+        print(f"extract_failed={summary['extract_failed'] * 100:.2f}%")
+        if "rule_vs_llm_agreement" in summary:
+            print(f"rule vs llm agreement = {summary['rule_vs_llm_agreement'] * 100:.2f}%")
+        print(f"wrote {os.path.join(out_dir, 'scored.jsonl')} / summary.json")
+    return 0
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="score a benchmark run and summarize Pass@k")
+    parser.add_argument("--out-dir", required=True)
+    parser.add_argument(
+        "--judge", choices=["auto", "rule", "llm"], default="auto", help="auto: per-benchmark default (all three use the LLM judge)"
+    )
+    parser.add_argument("--judge-concurrency", type=int, default=8)
+    parser.add_argument("--data-file", default=None, help="override meta.json's data_file")
+    args = parser.parse_args()
+    return score_run(
+        args.out_dir,
+        judge=args.judge,
+        judge_concurrency=args.judge_concurrency,
+        data_file=args.data_file,
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())