LRM/test_eval_utils.py at main · 1e4ves/LRM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
#!/usr/bin/env python3
"""Test official evaluation utilities."""

from eval_utils_official import (
    parse_gsm8k_answer,
    score_gsm8k,
    extract_math_answer,
    check_math_equality,
    extract_aime_answer,
    score_aime,
)


def test_gsm8k():
    print("Testing GSM8K evaluation...")

    # Test answer extraction
    text1 = "The answer is 42. #### 42"
    assert parse_gsm8k_answer(text1) == "42"

    text2 = "Let me calculate: 10 + 20 = 30\n#### 30"
    assert parse_gsm8k_answer(text2) == "30"

    text3 = "The total is 1,234.5\n#### 1,234.5"
    assert parse_gsm8k_answer(text3) == "1234.5"

    # Test scoring
    assert score_gsm8k("42", "42") == True
    assert score_gsm8k("42", "43") == False
    assert score_gsm8k("1234", "1,234") == True
    assert score_gsm8k("10.0", "10") == True

    print("✓ GSM8K tests passed")


def test_math():
    print("\nTesting MATH evaluation...")

    # Test answer extraction
    text1 = "The solution is x = 5.\nAnswer: 5"
    assert extract_math_answer(text1) == "5"

    text2 = "After simplification:\nanswer: \\frac{1}{2}"
    assert extract_math_answer(text2) == "\\frac{1}{2}"

    # Test equality checking
    assert check_math_equality("5", "5") == True
    assert check_math_equality("1/2", "0.5") == True
    assert check_math_equality("2", "3") == False

    print("✓ MATH tests passed")


def test_aime():
    print("\nTesting AIME evaluation...")

    # Test answer extraction
    text1 = "The answer is \\boxed{123}"
    assert extract_aime_answer(text1) == 123

    text2 = "Therefore, the answer is 456.\n#### 456"
    assert extract_aime_answer(text2) == 456

    text3 = "The final answer is 789"
    assert extract_aime_answer(text3) == 789

    # Test out of range
    text4 = "The answer is 1234"  # > 999
    assert extract_aime_answer(text4) is None

    # Test scoring
    assert score_aime(123, 123) == True
    assert score_aime(123, 456) == False
    assert score_aime(123, None) == False

    print("✓ AIME tests passed")


if __name__ == "__main__":
    test_gsm8k()
    test_math()
    test_aime()
    print("\n✅ All tests passed!")