refactor

Author: JamesDeng42

.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

@@ -130,31 +130,14 @@ 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 one-shot chat-path diagnostic with the
-same model/server settings. For an 8-chip v6e GPT-OSS deployment in HF mode:
+Before a long TPU benchmark, run the simple-evals unit suite, which checks
+Harmony output separation, request-error behavior, concurrency, and result
+reporting:
 
 ```bash
-python -m maxtext.eval.runner.debug_simple_evals \
-  --model_name gpt-oss-20b \
-  --hf_path openai/gpt-oss-20b \
-  --hf_mode \
-  --base_output_directory /tmp \
-  --run_name simple_evals_debug \
-  --max_model_len 32768 \
-  --tensor_parallel_size 8 \
-  --debug_max_tokens 1024 \
-  --reasoning_effort high \
-  --debug_output /tmp/simple_evals_tpu_debug.json
+pytest tests/unit/eval/test_simple_evals_runner.py
 ```
 
-The script launches the normal in-process server and checks Harmony prompt
-rendering, final/reasoning/raw separation, forced analysis-only truncation,
-diagnostic opt-in, token accounting, and two-request batch response ordering.
-It prints `PASS` and exits zero only when every applicable check succeeds;
-otherwise it writes the failure to the JSON report and exits nonzero. When
-testing a MaxText checkpoint, pass the same `--checkpoint_path` and
-`--model_name` used by the benchmark instead of `--hf_mode`.
-
 ## Common Flags
 
 | Flag | Description |

src/maxtext/eval/native_evals/aime_eval.py

@@ -12,6 +12,8 @@
 # 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
@@ -23,6 +25,9 @@
 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"
 )
@@ -37,7 +42,7 @@
 
 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 '
+    "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}'
 )
 
@@ -71,22 +76,20 @@ def __init__(self, year: int, num_examples: int | None = None, n_repeats: int =
 
   def __call__(self, sampler: SamplerBase) -> EvalResult:
     def fn(row: dict):
-      prompt_messages = [
-          sampler._pack_message(content=QUERY_TEMPLATE.format(problem=row["problem"]), role="user")
-      ]
+      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)
+      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=dict(content=response_text, role="assistant"),
+          next_message={"content": response_text, "role": "assistant"},
           score=score,
           correct_answer=row["answer"],
           extracted_answer=extracted_answer,
       )
-      convo = actual_queried_prompt_messages + [dict(content=response_text, role="assistant")]
+      convo = actual_queried_prompt_messages + [{"content": response_text, "role": "assistant"}]
       return SingleEvalResult(
           html=html,
           score=score,

src/maxtext/eval/native_evals/gsm8k_eval.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""GSM8K simple-evals task."""
+
 from __future__ import annotations
 
 import random
@@ -24,6 +26,9 @@
 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"
 
@@ -83,22 +88,20 @@ def __init__(self, num_examples: int | None = None):
 
   def __call__(self, sampler: SamplerBase) -> EvalResult:
     def fn(row: dict):
-      prompt_messages = [
-          sampler._pack_message(content=QUERY_TEMPLATE.format(question=row["question"]), role="user")
-      ]
+      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)
+      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=dict(content=response_text, role="assistant"),
+          next_message={"content": response_text, "role": "assistant"},
           score=score,
           correct_answer=row["target"],
           extracted_answer=extracted_answer or None,
       )
-      convo = actual_queried_prompt_messages + [dict(content=response_text, role="assistant")]
+      convo = actual_queried_prompt_messages + [{"content": response_text, "role": "assistant"}]
       return SingleEvalResult(
           html=html,
           score=score,