DeepInfer/tests/eval/test_aime.py

48 lines
1.3 KiB
Python

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.",
"<think>I first tried 17.</think>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")