diff --git a/docs/guides/eval.md b/docs/guides/eval.md index 74eee7a004a..9e679fabd0d 100644 --- a/docs/guides/eval.md +++ b/docs/guides/eval.md @@ -103,7 +103,7 @@ score=0.1000 (3.0/30) ## List of currently supported benchmarks -- [AIME-2024 and AIME-2025](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"` and `"aime2025"`. +- [AIME-2024, AIME-2025, and AIME-2026](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"`, `"aime2025"`, and `"aime2026"`. - [GPQA and GPQA-diamond](../../nemo_rl/data/datasets/eval_datasets/gpqa.py): the corresponding `data.dataset_name` are `"gpqa"` and `"gpqa_diamond"`. - [MATH and MATH-500](../../nemo_rl/data/datasets/eval_datasets/math.py): the corresponding `data.dataset_name` are `"math"` and `"math500"`. - [MMLU](../../nemo_rl/data/datasets/eval_datasets/mmlu.py): this also includes MMMLU (Multilingual MMLU), a total of 14 languages. When `data.dataset_name` is set to `mmlu`, the English version is used. If one wants to run evaluation on another language, `data.dataset_name` should be set to `mmlu_{language}` where `language` is one of following 14 values, `["AR-XY", "BN-BD", "DE-DE", "ES-LA", "FR-FR", "HI-IN", "ID-ID", "IT-IT", "JA-JP", "KO-KR", "PT-BR", "ZH-CN", "SW-KE", "YO-NG"]`. diff --git a/nemo_rl/data/datasets/eval_datasets/__init__.py b/nemo_rl/data/datasets/eval_datasets/__init__.py index d813ed040cc..296323efda8 100644 --- a/nemo_rl/data/datasets/eval_datasets/__init__.py +++ b/nemo_rl/data/datasets/eval_datasets/__init__.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset +from typing import cast + +from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset, AIMEVariant from nemo_rl.data.datasets.eval_datasets.gpqa import GPQADataset from nemo_rl.data.datasets.eval_datasets.local_math_dataset import LocalMathDataset from nemo_rl.data.datasets.eval_datasets.math import MathDataset @@ -53,15 +55,9 @@ def load_eval_dataset(data_config): system_prompt_file=data_config["system_prompt_file"], ) # aime - elif dataset_name == "aime2024": - base_dataset = AIMEDataset( - variant="2024", - prompt_file=data_config["prompt_file"], - system_prompt_file=data_config["system_prompt_file"], - ) - elif dataset_name == "aime2025": + elif dataset_name in ["aime2024", "aime2025", "aime2026"]: base_dataset = AIMEDataset( - variant="2025", + variant=cast(AIMEVariant, dataset_name[4:]), prompt_file=data_config["prompt_file"], system_prompt_file=data_config["system_prompt_file"], ) diff --git a/nemo_rl/data/datasets/eval_datasets/aime.py b/nemo_rl/data/datasets/eval_datasets/aime.py index 671b410a69e..d63916ff56d 100644 --- a/nemo_rl/data/datasets/eval_datasets/aime.py +++ b/nemo_rl/data/datasets/eval_datasets/aime.py @@ -21,11 +21,13 @@ from nemo_rl.data import processors from nemo_rl.data.interfaces import TaskDataSpec +AIMEVariant = Literal["2024", "2025", "2026"] + class AIMEDataset: def __init__( self, - variant: Literal["2024", "2025"] = "2025", + variant: AIMEVariant = "2025", prompt_file: Optional[str] = None, system_prompt_file: Optional[str] = None, ): @@ -37,6 +39,9 @@ def __init__( ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test") ds = concatenate_datasets([ds0, ds1]) self.input_key = "question" + elif variant == "2026": + ds = load_dataset("MathArena/aime_2026", split="train") + self.input_key = "problem" else: raise ValueError(f"Invalid variant for aime dataset: aime{variant}") diff --git a/tests/unit/data/datasets/test_eval_dataset.py b/tests/unit/data/datasets/test_eval_dataset.py index dc567d09d75..5e33b9181c0 100644 --- a/tests/unit/data/datasets/test_eval_dataset.py +++ b/tests/unit/data/datasets/test_eval_dataset.py @@ -82,6 +82,39 @@ def test_math_dataset(): ) +@pytest.mark.parametrize("variant", ["aime2024", "aime2025", "aime2026"]) +@pytest.mark.skip(reason="dataset download is flaky") +def test_aime_dataset(variant): + # load the dataset + data_config = { + "dataset_name": variant, + "prompt_file": None, + "system_prompt_file": None, + } + aime_dataset = load_eval_dataset(data_config) + + # load the tokenizer + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct") + + # check that the dataset is formatted correctly + for example in aime_dataset.rekeyed_ds.take(5): + assert "problem" in example + assert "expected_answer" in example + + ## check that applying chat template works as expected + default_templated = tokenizer.apply_chat_template( + [{"role": "user", "content": example["problem"]}], + tokenize=False, + add_generation_prompt=False, + add_special_tokens=False, + ) + + assert ( + default_templated + == f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{example['problem']}<|im_end|>\n" + ) + + @pytest.mark.skip(reason="dataset download is flaky") def test_mmlu_dataset(): # load the dataset