diff --git a/eval/aime/bench.py b/eval/aime/bench.py index b0aec2d..36727dd 100644 --- a/eval/aime/bench.py +++ b/eval/aime/bench.py @@ -124,7 +124,8 @@ def extract_answer_with_source(text: str) -> tuple[str | None, str]: A fallback for malformed output, **higher risk**; score.py reports its share separately; 3. `none`: nothing found. """ - boxed = extract_boxed(text) + # Ignore boxed intermediate results in the thinking section. + boxed = extract_boxed(text.rsplit("", 1)[-1]) if boxed is not None: norm = normalize_answer(boxed) if norm: diff --git a/tests/eval/test_aime_thinking.py b/tests/eval/test_aime_thinking.py new file mode 100644 index 0000000..56ceb3d --- /dev/null +++ b/tests/eval/test_aime_thinking.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +from eval.aime import bench + +pytestmark = [pytest.mark.ci, pytest.mark.ci_npu] + + +@pytest.mark.parametrize( + "text,expected,source", + [ + (r"Try \boxed{17}The answer is 42", "42", "tail_number"), + (r"Try \boxed{17}", None, "none"), + (r"Try \boxed{17}I cannot determine the answer", None, "none"), + (r"draft \boxed{17}42", "42", "tail_number"), + (r"\boxed{11}\boxed{17}42", "42", "tail_number"), + (r"Try \boxed{17}The answer is \boxed{42}", "42", "boxed"), + (r"Try \boxed{17}\boxed{41} then \boxed{42}", "42", "boxed"), + (r"Try \boxed{17}\boxed{} answer 42", "42", "tail_number"), + (r"Try \boxed{17}42" + " " * 450, "42", "tail_number"), + (r"\boxed{42}", "42", "boxed"), + (r"\boxed{41} then \boxed{42}", "42", "boxed"), + (r"\boxed{42}" + " " * 450, "42", "boxed"), + ("42" + " " * 450, None, "none"), + ("The answer is 42", "42", "tail_number"), + ("", None, "none"), + (r"\boxed{\frac{1}{2}}", r"\frac{1}{2}", "boxed"), + ], +) +def test_extracts_final_answer(text: str, expected: str | None, source: str) -> None: + assert bench.extract_answer_with_source(text) == (expected, source) + + +@pytest.mark.parametrize("gold,expected", [("42", True), ("17", False)]) +def test_score_rollout_preserves_raw_text(gold: str, expected: bool) -> None: + text = r"Try \boxed{17}The answer is 42" + rollout = {"task_id": "problem", "sample_idx": 0, "raw_text": text} + row = bench.score_rollout(rollout, {"answer": gold}) + assert row["correct"] is expected + assert row["pred"] == "42" + assert row["raw_text"] == text + assert rollout == {"task_id": "problem", "sample_idx": 0, "raw_text": text} + + +@pytest.mark.parametrize( + "text,gold,expected", + [ + (r"Try \boxed{17}The answer is 42", "42", 1.0), + (r"Try \boxed{17}The answer is 42", "17", 0.0), + (r"Try \boxed{17}The answer is \boxed{42}", "42", 1.0), + ], +) +def test_rule_scoring_cli(tmp_path: Path, text: str, gold: str, expected: float) -> None: + data = tmp_path / "problems.jsonl" + data.write_text(json.dumps({"task_id": "problem", "question": "Synthetic extraction fixture", "answer": gold}) + "\n") + (tmp_path / "meta.json").write_text(json.dumps({"bench": "aime", "data_file": str(data), "n_sample": 1})) + (tmp_path / "raw_rank0.jsonl").write_text( + json.dumps({"task_id": "problem", "sample_idx": 0, "raw_text": text, "finish_reason": "stop"}) + "\n" + ) + script = Path(bench.__file__).resolve().parents[1] / "score.py" + result = subprocess.run( + [sys.executable, str(script), "--out-dir", str(tmp_path), "--judge", "rule"], + capture_output=True, + text=True, + timeout=20, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + summary = json.loads((tmp_path / "summary.json").read_text()) + assert summary["coverage_complete"] is True + assert summary["Pass@1"] == expected + row = json.loads((tmp_path / "scored.jsonl").read_text()) + assert row["raw_text"] == text