diff --git a/eval/aime/bench.py b/eval/aime/bench.py index b0aec2d..12ba8b3 100644 --- a/eval/aime/bench.py +++ b/eval/aime/bench.py @@ -89,7 +89,8 @@ def extract_boxed(text: str) -> str | None: return None -_NUMBER_RE = re.compile(r"(? str | None: diff --git a/tests/eval/test_aime.py b/tests/eval/test_aime.py new file mode 100644 index 0000000..60e02ea --- /dev/null +++ b/tests/eval/test_aime.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import pytest + +from eval.aime.bench import extract_answer_with_source, score_rollout + +pytestmark = [pytest.mark.ci, pytest.mark.ci_npu] + + +@pytest.mark.parametrize( + "text", + [ + "The answer is 42", + "The answer is 42.", + "The answer is 42. This is the final answer.", + "The answer is 42.\n", + "The answer is 42!", + "The answer is 42...", + "The answer is 042.", + "I first tried 17.The answer is 42.", + ], +) +def test_tail_number_with_sentence_punctuation(text: str) -> None: + assert extract_answer_with_source(text) == ("42", "tail_number") + row = score_rollout({"raw_text": text}, {"answer": "42"}) + assert row["correct"] is True + + +@pytest.mark.parametrize( + "text", + [ + "The answer is 42.5", + "The answer is 42.5.", + "The answer is .42", + "The answer is 0.042.", + "The answer is 1234.", + "The answer is x42.", + "The answer is 42x.", + "There is no answer.", + ], +) +def test_tail_number_rejects_non_integer_tokens(text: str) -> None: + assert extract_answer_with_source(text) == (None, "none") + + +def test_boxed_answer_takes_priority_over_tail_number() -> None: + assert extract_answer_with_source(r"The answer is \boxed{17}. I rejected 42.") == ("17", "boxed")