Add Simple-evals supports to maxtext runner

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
         args:
           - '-w'
           - '--skip="*.txt,pylintrc,.*,src/maxtext/assets/*,src/maxtext/input_pipeline/protos/*"'
-          - '-L ND,nd,sems,TE,ROUGE,rouge,astroid,ags,dout'
+          - '-L ND,nd,sems,TE,te,als,ROUGE,rouge,astroid,ags,dout'
           - '.'
         additional_dependencies:
           - tomli
@@ -30,7 +30,7 @@ repos:
         args:
           - '--disable=R0401,R0917,W0201,W0613'
           - "--ignore-patterns='.pytype,.*pyi$'"
-          - '--ignore-paths=src/maxtext/input_pipeline/protos'
+          - '--ignore-paths=src/maxtext/input_pipeline/protos|src/maxtext/eval/third_party'
 
   #  - repo: https://github.com/google/pytype
   #    rev: 2024.10.11

src/maxtext/eval/README.md

@@ -7,7 +7,7 @@ A vLLM-native evaluation framework for MaxText models supporting harness-based e
 All runners share a single entry point:
 
 ```bash
-python -m maxtext.eval.runner.run --runner <eval|lm_eval|evalchemy> [flags]
+python -m maxtext.eval.runner.run --runner <eval|lm_eval|evalchemy|simple_evals> [flags]
 ```
 
 ### Custom dataset (MLPerf OpenOrca, ROUGE scoring, Other)
@@ -78,6 +78,66 @@ python -m maxtext.eval.runner.run \
   --hf_token $HF_TOKEN
 ```
 
+### Simple Evals (OpenAI simple-evals)
+
+Runs OpenAI's [simple-evals](https://github.com/openai/simple-evals) benchmarks
+against the vLLM server's OpenAI-compatible `/v1/chat/completions` endpoint.
+
+Requires: `pip install aiohttp openai jinja2 numpy pandas scipy requests tqdm`
+
+Phase 1 supports grader-free evals only: `mmlu`, `gpqa`, `gsm8k`, `drop`, `mgsm`,
+`mgsm_en`, `aime2024`, `aime2025`. Grader-dependent evals (math, simpleqa, browsecomp,
+healthbench) need an LLM grader endpoint and are not yet supported.
+
+`mmlu`, `gpqa`, `drop`, `mgsm` are based on the vendored implementations from
+[openai/simple-evals](https://github.com/openai/simple-evals); `gsm8k` and
+`aime2024`/`aime2025` are not part of upstream simple-evals and are
+MaxText-authored (`maxtext.eval.native_evals`) following the same
+grader-free Eval conventions. `mgsm` matches upstream's 11-language suite;
+`mgsm_en` is the explicitly named English-only variant.
+
+```bash
+python -m maxtext.eval.runner.run \
+  --runner simple_evals \
+  --checkpoint_path gs://<bucket>/checkpoints/0/items \
+  --model_name llama3.1-8b \
+  --hf_path meta-llama/Llama-3.1-8B-Instruct \
+  --tasks mmlu gpqa gsm8k drop mgsm aime2024 aime2025 \
+  --base_output_directory gs://<bucket>/ \
+  --run_name simple_evals_run \
+  --max_model_len 8192 \
+  --tensor_parallel_size 4 \
+  --num_samples 50 \
+  --hf_token $HF_TOKEN
+```
+
+Simple-evals specific flags:
+
+| Flag | Description |
+|---|---|
+| `--tasks` | Space-separated task names (`mmlu`, `gpqa`, `gsm8k`, `drop`, `mgsm`, `mgsm_en`, `aime2024`, `aime2025`). |
+| `--num_samples` | Limit examples per task; for MGSM variants this is per language (None = full dataset). |
+| `--n_repeats` | Repeats per example (defaults: upstream GPQA=4, AIME=1; forced to 1 with `--num_samples`). |
+| `--max_tokens` | Max tokens per generation (default: 2048). |
+| `--temperature` | Sampling temperature (upstream default: 0.5). |
+| `--concurrency` | Task worker count and maximum in-flight requests. Omit for automatic selection from CPU count, accelerator count, `max_num_seqs`, and chat batch capacity. This is the only request-pressure setting users normally need to tune. |
+| `--log-debug-info` | Write a timestamp-matched `.debug.txt` report beside the result JSON. The report includes request/token throughput, latency percentiles, retry and terminal-error rates, output truncation, answer-extraction failures, confidence intervals, and bounded samples of final, reasoning, and raw model output. Reasoning and raw output remain diagnostic-only and are never graded. |
+| `--continue-on-request-error` | Score terminal request failures as zero instead of aborting. This is independent of debug logging. |
+
+Debug logging is observational and does not change request-failure behavior.
+When `--continue-on-request-error` is enabled, failed requests count as zero in
+official accuracy. The debug report also shows diagnostic accuracy excluding
+infrastructure failures; that diagnostic must not be used as the published
+benchmark score.
+
+Before a long TPU benchmark, run the simple-evals unit suite, which checks
+Harmony output separation, request-error behavior, concurrency, and result
+reporting:
+
+```bash
+pytest tests/unit/eval/test_simple_evals_runner.py
+```
+
 ## Common Flags
 
 | Flag | Description |

src/maxtext/eval/native_evals/aime_eval.py

@@ -0,0 +1,108 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AIME 2024 and 2025 simple-evals tasks."""
+
+from __future__ import annotations
+
+import random
+import re
+
+import pandas
+
+from maxtext.eval.third_party.simple_evals import common
+from maxtext.eval.third_party.simple_evals.common import HTML_JINJA
+from maxtext.eval.third_party.simple_evals.types import Eval, EvalResult, SamplerBase, SingleEvalResult
+
+# SamplerBase intentionally preserves the vendored simple-evals API.
+# pylint: disable=protected-access
+
+_AIME_2024_URL = (
+    "https://huggingface.co/datasets/Maxwell-Jia/AIME_2024/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
+)
+_AIME_2025_URL = (
+    "https://huggingface.co/datasets/math-ai/aime25/resolve/refs%2Fconvert%2Fparquet/default/test/0000.parquet"
+)
+
+_DATASET_BY_YEAR = {
+    2024: (_AIME_2024_URL, "Problem", "Answer"),
+    2025: (_AIME_2025_URL, "problem", "answer"),
+}
+
+QUERY_TEMPLATE = (
+    "Solve this competition math problem. The final answer is always an integer between 0 and 999. "
+    "Think step by step, then write the final answer by itself on the last line in the format "
+    '"Answer: $N" (without quotes), where $N is the integer answer.\n\n{problem}'
+)
+
+_ANSWER_LINE_RE = re.compile(r"(?i)(?:^|\n)\s*Answer\s*:\s*\$?([0-9]{1,3})\$?\s*\Z")
+
+
+def extract_answer(response_text: str) -> str | None:
+  """Pull the integer answer out of a response.
+
+  Requires the explicit final-line format requested by QUERY_TEMPLATE.
+  """
+  match = _ANSWER_LINE_RE.search(response_text)
+  return match.group(1) if match else None
+
+
+class AIMEEval(Eval):
+  """AIME eval for a single contest year (2024 or 2025)."""
+
+  def __init__(self, year: int, num_examples: int | None = None, n_repeats: int = 1):
+    if year not in _DATASET_BY_YEAR:
+      raise ValueError(f"Unsupported AIME year: {year}. Supported: {sorted(_DATASET_BY_YEAR)}.")
+    url, problem_col, answer_col = _DATASET_BY_YEAR[year]
+    df = pandas.read_parquet(url)
+    examples = [{"problem": row[problem_col], "answer": str(int(row[answer_col]))} for _, row in df.iterrows()]
+    rng = random.Random(0)
+    if num_examples:
+      assert n_repeats == 1, "n_repeats only supported for num_examples = None"
+      examples = rng.sample(examples, num_examples)
+    self.examples = examples * n_repeats
+    self.year = year
+
+  def __call__(self, sampler: SamplerBase) -> EvalResult:
+    def fn(row: dict):
+      prompt_messages = [sampler._pack_message(content=QUERY_TEMPLATE.format(problem=row["problem"]), role="user")]
+      sampler_response = sampler(prompt_messages)
+      response_text = sampler_response.response_text
+      actual_queried_prompt_messages = sampler_response.actual_queried_message_list
+      extracted_answer = extract_answer(response_text) if common.request_succeeded(sampler_response) else None
+      score = 1.0 if extracted_answer is not None and int(extracted_answer) == int(row["answer"]) else 0.0
+      html = common.jinja_env.from_string(HTML_JINJA).render(
+          prompt_messages=actual_queried_prompt_messages,
+          next_message={"content": response_text, "role": "assistant"},
+          score=score,
+          correct_answer=row["answer"],
+          extracted_answer=extracted_answer,
+      )
+      convo = actual_queried_prompt_messages + [{"content": response_text, "role": "assistant"}]
+      return SingleEvalResult(
+          html=html,
+          score=score,
+          convo=convo,
+          metrics={"chars": len(response_text)},
+          example_level_metadata={
+              "request_id": sampler_response.response_metadata.get("request_id"),
+              "request_status": sampler_response.response_metadata.get("status", "success"),
+              "score": score,
+              "correct_answer": row["answer"],
+              "extracted_answer": extracted_answer,
+          },
+      )
+
+    results = common.map_with_progress(fn, self.examples)
+    return common.aggregate_results(results)

src/maxtext/eval/native_evals/gsm8k_eval.py

@@ -0,0 +1,119 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""GSM8K simple-evals task."""
+
+from __future__ import annotations
+
+import random
+import re
+
+import pandas
+
+from maxtext.eval.third_party.simple_evals import common
+from maxtext.eval.third_party.simple_evals.common import HTML_JINJA
+from maxtext.eval.third_party.simple_evals.mgsm_eval import score_mgsm
+from maxtext.eval.third_party.simple_evals.types import Eval, EvalResult, SamplerBase, SingleEvalResult
+
+# SamplerBase intentionally preserves the vendored simple-evals API.
+# pylint: disable=protected-access
+
+_DATASET_REPO_ID = "openai/gsm8k"
+_TEST_PARQUET_FILENAME = "main/test-00000-of-00001.parquet"
+
+QUERY_TEMPLATE = (
+    "Solve this math problem. Give the reasoning steps before giving the final answer on the last line "
+    'by itself in the format of "Answer:". Do not add anything other than the integer answer after '
+    '"Answer:".\n\n{question}'
+)
+
+_ANSWER_LINE_RE = re.compile(r"(?i)(?:^|\n)\s*Answer\s*:\s*(-?[\d,]+(?:\.\d+)?)\s*\Z")
+
+
+def extract_answer(response_text: str) -> str | None:
+  """Extract the last answer that obeys GSM8K's requested final-line format."""
+  match = _ANSWER_LINE_RE.search(response_text)
+  return match.group(1).replace(",", "") if match else None
+
+
+def _load_gsm8k_examples() -> list[dict]:
+  """Download/cache the canonical GSM8K test parquet through HF Hub."""
+  try:
+    from huggingface_hub import hf_hub_download  # pylint: disable=import-outside-toplevel
+  except ImportError as exc:
+    raise ImportError(
+        "GSM8K evaluation requires huggingface_hub. Install it with `pip install huggingface_hub`."
+    ) from exc
+
+  try:
+    parquet_path = hf_hub_download(
+        repo_id=_DATASET_REPO_ID,
+        repo_type="dataset",
+        filename=_TEST_PARQUET_FILENAME,
+    )
+  except Exception as exc:  # pylint: disable=broad-except
+    raise RuntimeError(
+        "Could not download the canonical GSM8K main/test split from Hugging Face. "
+        "Check outbound Hub access and HF_TOKEN for any proxy policy."
+    ) from exc
+
+  df = pandas.read_parquet(parquet_path)
+  examples = [row.to_dict() for _, row in df.iterrows()]
+  if len(examples) != 1319:
+    raise ValueError(f"Expected 1319 GSM8K main/test examples, found {len(examples)}.")
+  return examples
+
+
+class GSM8KEval(Eval):
+  """GSM8K eval: zero-shot CoT prompt, gold answer parsed from the '#### N' suffix."""
+
+  def __init__(self, num_examples: int | None = None):
+    examples = _load_gsm8k_examples()
+    for example in examples:
+      example["target"] = example["answer"].split("####")[-1].strip().replace(",", "")
+    if num_examples:
+      examples = random.Random(0).sample(examples, num_examples)
+    self.examples = examples
+
+  def __call__(self, sampler: SamplerBase) -> EvalResult:
+    def fn(row: dict):
+      prompt_messages = [sampler._pack_message(content=QUERY_TEMPLATE.format(question=row["question"]), role="user")]
+      sampler_response = sampler(prompt_messages)
+      response_text = sampler_response.response_text
+      actual_queried_prompt_messages = sampler_response.actual_queried_message_list
+      extracted_answer = extract_answer(response_text) if common.request_succeeded(sampler_response) else None
+      score = score_mgsm(row["target"], extracted_answer or "")
+      html = common.jinja_env.from_string(HTML_JINJA).render(
+          prompt_messages=actual_queried_prompt_messages,
+          next_message={"content": response_text, "role": "assistant"},
+          score=score,
+          correct_answer=row["target"],
+          extracted_answer=extracted_answer or None,
+      )
+      convo = actual_queried_prompt_messages + [{"content": response_text, "role": "assistant"}]
+      return SingleEvalResult(
+          html=html,
+          score=score,
+          convo=convo,
+          example_level_metadata={
+              "request_id": sampler_response.response_metadata.get("request_id"),
+              "request_status": sampler_response.response_metadata.get("status", "success"),
+              "score": score,
+              "correct_answer": row["target"],
+              "extracted_answer": extracted_answer or None,
+          },
+      )
+
+    results = common.map_with_progress(fn, self.examples)
+    return common.aggregate_results(results)