merge aime

Signed-off-by: Yuki Huang <[email protected]>

Author: yuki-97

docs/guides/eval.md

@@ -89,8 +89,7 @@ score=0.1000 (3.0/30)
 
 ## List of currently supported benchmarks
 
-- [AIME-2024](../../nemo_rl/data/datasets/eval_datasets/aime2024.py): the corresponding `data.dataset_name` is `"aime2024"`.
-- [AIME-2025](../../nemo_rl/data/datasets/eval_datasets/aime2025.py): the corresponding `data.dataset_name` is `"aime2025"`.
+- [AIME-2024 and AIME-2025](../../nemo_rl/data/datasets/eval_datasets/aime.py): the corresponding `data.dataset_name` are `"aime2024"` and `"aime2025"`.
 - [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"]`.

docs/guides/sft.md

@@ -71,7 +71,6 @@ NeMo RL SFT uses Hugging Face chat templates to format the individual examples.
     custom_template: "{% for message in messages %}{%- if message['role'] == 'system'  %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user'  %}{{' Question: ' + message['content'].strip() + ' Answer: '}}{%- elif message['role'] == 'assistant'  %}{{message['content'].strip()}}{%- endif %}{% endfor %}"
     ```
 
-
 By default, NeMo RL has support for [OpenAssistant](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data/datasets/response_datasets/oasst.py), [Squad](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data/datasets/response_datasets/squad.py) and [OpenMathInstruct-2](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/data/datasets/response_datasets/openmathinstruct2.py) datasets. All of these datasets are downloaded from Hugging Face and preprocessed on-the-fly, so there's no need to provide a path to any datasets on disk.
 
 We provide a [ResponseDataset](../../nemo_rl/data/datasets/response_datasets/response_dataset.py) class that is compatible with jsonl-formatted response datasets. You can use `input_key`, `output_key` to specify which fields in your data correspond to the question and answer respectively. Here's an example configuration:

nemo_rl/data/datasets/eval_datasets/__init__.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemo_rl.data.datasets.eval_datasets.aime2024 import AIME2024Dataset
-from nemo_rl.data.datasets.eval_datasets.aime2025 import AIME2025Dataset
+from nemo_rl.data.datasets.eval_datasets.aime import AIMEDataset
 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
@@ -46,12 +45,14 @@ def load_eval_dataset(data_config):
         )
     # aime
     elif dataset_name == "aime2024":
-        base_dataset = AIME2024Dataset(
+        base_dataset = AIMEDataset(
+            variant="2024",
             prompt_file=data_config["prompt_file"],
             system_prompt_file=data_config["system_prompt_file"],
         )
     elif dataset_name == "aime2025":
-        base_dataset = AIME2025Dataset(
+        base_dataset = AIMEDataset(
+            variant="2025",
             prompt_file=data_config["prompt_file"],
             system_prompt_file=data_config["system_prompt_file"],
         )
@@ -98,8 +99,7 @@ def load_eval_dataset(data_config):
 
 
 __all__ = [
-    "AIME2024Dataset",
-    "AIME2025Dataset",
+    "AIMEDataset",
     "GPQADataset",
     "LocalMathDataset",
     "MathDataset",

nemo_rl/data/datasets/eval_datasets/aime2025.py renamed to nemo_rl/data/datasets/eval_datasets/aime.py

@@ -12,35 +12,44 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""AIME 2025 dataset."""
+"""AIME dataset."""
 
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 from datasets import concatenate_datasets, load_dataset
 
 from nemo_rl.data import processors
 from nemo_rl.data.interfaces import TaskDataSpec
 
 
-class AIME2025Dataset:
+class AIMEDataset:
     def __init__(
         self,
+        variant: Literal["2024", "2025"] = "2025",
         prompt_file: Optional[str] = None,
         system_prompt_file: Optional[str] = None,
     ):
-        ds0 = load_dataset("opencompass/AIME2025", "AIME2025-I", split="test")
-        ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test")
-        ds = concatenate_datasets([ds0, ds1])
+        if variant == "2024":
+            ds = load_dataset("HuggingFaceH4/aime_2024", split="train")
+            self.input_key = "problem"
+        elif variant == "2025":
+            ds0 = load_dataset("opencompass/AIME2025", "AIME2025-I", split="test")
+            ds1 = load_dataset("opencompass/AIME2025", "AIME2025-II", split="test")
+            ds = concatenate_datasets([ds0, ds1])
+            self.input_key = "question"
+        else:
+            raise ValueError(f"Invalid variant for aime dataset: aime{variant}")
+
         self.rekeyed_ds = ds.map(self._rekey, remove_columns=ds.column_names)
         self.task_spec = TaskDataSpec(
-            task_name="aime2025",
+            task_name=f"aime{variant}",
             prompt_file=prompt_file,
             system_prompt_file=system_prompt_file,
         )
         self.processor = processors.math_data_processor
 
     def _rekey(self, data: dict[str, Any]):
         return {
-            "problem": data["question"],
+            "problem": data[self.input_key],
             "expected_answer": data["answer"],
         }