Add aime26

[bzantium]

resolved: #1255

Summary by CodeRabbit

• New Features
  • Added AIME 2026 dataset support with preconfigured math evaluation defaults (overridable at runtime).
  • Added dataset preparation tooling to export AIME 2026 into JSONL for evaluation, including a command-line entry point and progress feedback.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds an AIME 2026 dataset package with module-level evaluation/generation constants and a CLI-able prepare script that loads the MathArena/aime_2026 dataset and exports split-specific JSONL manifests.

Changes

Cohort / File(s)	Summary
Package init nemo_skills/dataset/aime26/__init__.py	Adds module-level constants: DATASET_GROUP="math", METRICS_TYPE="math", GENERATION_ARGS="++prompt_config=generic/math ++eval_type=math".
Data preparation script nemo_skills/dataset/aime26/prepare.py	Adds format_entry(), write_data_to_file(), and main(); loads MathArena/aime_2026 (specified split), formats entries (id prefixed with aime26-, problem, expected_answer) and writes UTF-8 JSONL with progress reporting.

Sequence Diagram(s)

sequenceDiagram
    actor User
    participant CLI
    participant Prepare as prepare.py
    participant HF as "HuggingFace Dataset\n(MathArena/aime_2026)"
    participant FS as "Local Filesystem"

    User->>CLI: run prepare.py --split <split>
    CLI->>Prepare: invoke main(args)
    Prepare->>HF: load_dataset(..., split=split)
    HF-->>Prepare: dataset iterator
    loop for each entry
        Prepare->>Prepare: format_entry(entry)
        Prepare->>FS: append JSONL line
        FS-->>Prepare: write ACK
    end
    Prepare->>User: exit (file written)

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• Add apex-shortlist dataset #1080: Adds the same dataset-level constants and a prepare.py that exports MathArena datasets to JSONL.
• add musan dataset #1139: Adds a dataset package with identical module-level constants and a prepare.py entrypoint following the same export pattern.
• first version of 3 imo-bench datasets #1047: Introduces dataset package initializers with the same configuration constants and analogous preparation scripts.

Suggested labels

run GPU tests

Suggested reviewers

• gwarmstrong
• Kipok

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Add aime26' clearly and concisely describes the main change: adding support for the AIME 2026 benchmark to the codebase.
Linked Issues check	✅ Passed	The PR successfully implements the objective from issue #1255 by adding the AIME 2026 benchmark with dataset loading, formatting, and preparation utilities.
Out of Scope Changes check	✅ Passed	All changes are directly aligned with adding AIME 2026 benchmark support; no unrelated modifications are present.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

nemo_skills/dataset/aime26/prepare.py (1)

23-28: Missing type hints on functions.

As per coding guidelines, "Use type hints for simple types (dict, list, int, float, existing classes) in Python code."

Suggested type hints

-def format_entry(entry):
-    return {
+def format_entry(entry: dict) -> dict:
+    return {
         "id": f"aime26-{entry['problem_idx']}",
         "problem": entry["problem"],
         "expected_answer": str(entry["answer"]),
     }

-def write_data_to_file(output_file, data):
+def write_data_to_file(output_file: Path, data: list) -> None:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/dataset/aime26/prepare.py` around lines 23 - 28, The function
format_entry lacks type hints; update its signature to accept a typed parameter
(e.g., entry: Dict[str, Any] or Mapping[str, Any]) and declare the return type
(e.g., Dict[str, str]) and add the necessary typing imports (from typing import
Any, Dict or Mapping) so format_entry(entry: Dict[str, Any]) -> Dict[str, str]
is explicit; ensure the keys "id", "problem", and "expected_answer" are
reflected in the return type annotation.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@nemo_skills/dataset/aime26/__init__.py`:
- Line 1: Update the copyright year in the module header of
nemo_skills/dataset/aime26/__init__.py from 2024 to 2026 to match prepare.py;
open the file, locate the top-of-file copyright comment (the string "Copyright
(c) 2024, NVIDIA CORPORATION.  All rights reserved."), and change the year to
2026 so both files use the same copyright year.

In `@nemo_skills/dataset/aime26/prepare.py`:
- Around line 38-43: In main(), the call load_dataset("MathArena/aime_2026",
split="train") is hardcoded to "train" while the output file is named using
args.split (output_file = data_dir / f"{args.split}.jsonl"); add a brief
clarifying comment above the load_dataset line (inside the main function)
explaining that the HuggingFace dataset only exposes a single "train" split but
the benchmark uses that data as the target split (e.g., test/validation), so we
intentionally load "train" regardless of args.split to produce the appropriately
named output file.
- Line 48: The parser.add_argument call for "--split" uses choices=("test")
which is a string, causing argparse to check membership against characters;
change it to a proper iterable like a single-element tuple or list (e.g.,
choices=("test",) or choices=["test"]) in the parser.add_argument call to ensure
only "test" is accepted; update the choices parameter in the
parser.add_argument("--split", ...) invocation and mirror the style used by
other dataset prepare.py files.

---

Nitpick comments:
In `@nemo_skills/dataset/aime26/prepare.py`:
- Around line 23-28: The function format_entry lacks type hints; update its
signature to accept a typed parameter (e.g., entry: Dict[str, Any] or
Mapping[str, Any]) and declare the return type (e.g., Dict[str, str]) and add
the necessary typing imports (from typing import Any, Dict or Mapping) so
format_entry(entry: Dict[str, Any]) -> Dict[str, str] is explicit; ensure the
keys "id", "problem", and "expected_answer" are reflected in the return type
annotation.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Copyright year mismatch with prepare.py.

This file says Copyright (c) 2024 while prepare.py says Copyright (c) 2026. Since this is a new file for a 2026 benchmark, the year should likely be 2026 to be consistent.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/dataset/aime26/__init__.py` at line 1, Update the copyright year
in the module header of nemo_skills/dataset/aime26/__init__.py from 2024 to 2026
to match prepare.py; open the file, locate the top-of-file copyright comment
(the string "Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved."),
and change the year to 2026 so both files use the same copyright year.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check if other prepare.py scripts in this repo also load a different split than the output name
rg -n 'load_dataset.*split=' --glob 'nemo_skills/dataset/*/prepare.py' -A2 -B2

Repository: NVIDIA-NeMo/Skills

Length of output: 10510

---

Add a clarifying comment explaining the hardcoded split.

The dataset is loaded with split="train" (line 39) regardless of args.split, but the output file is named {args.split}.jsonl. This pattern—where the HuggingFace split differs from the output filename—appears elsewhere in the codebase (e.g., simpleqa), suggesting it's intentional because the HuggingFace dataset only exposes a train split but the benchmark treats it as a test set. A brief comment explaining this would prevent future confusion.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/dataset/aime26/prepare.py` around lines 38 - 43, In main(), the
call load_dataset("MathArena/aime_2026", split="train") is hardcoded to "train"
while the output file is named using args.split (output_file = data_dir /
f"{args.split}.jsonl"); add a brief clarifying comment above the load_dataset
line (inside the main function) explaining that the HuggingFace dataset only
exposes a single "train" split but the benchmark uses that data as the target
split (e.g., test/validation), so we intentionally load "train" regardless of
args.split to produce the appropriately named output file.

[coderabbitai]

🧹 Nitpick comments (1)

nemo_skills/dataset/aime26/prepare.py (1)

31-35: Missing type hints; consider pre-computing entries before opening the output file.

Two points:

1. Type hints — same guideline violation as format_entry:

♻️ Proposed fix for type hints

-def write_data_to_file(output_file, data):
+def write_data_to_file(output_file: Path, data) -> None:

2. Computation / I/O ordering — format_entry is called inside the open() block. A KeyError on a malformed entry mid-stream would leave a partially-written JSONL file on disk with no indication. Pre-materializing the formatted list before opening the file costs negligible memory for an AIME dataset (30 problems) and eliminates partial-output risk:

♻️ Proposed fix for computation ordering

 def write_data_to_file(output_file: Path, data) -> None:
+    formatted = [format_entry(entry) for entry in tqdm(data, desc="Processing")]
     with open(output_file, "wt", encoding="utf-8") as fout:
-        for entry in tqdm(data, desc=f"Writing {output_file.name}"):
-            json.dump(format_entry(entry), fout, ensure_ascii=False)
+        for entry in formatted:
+            json.dump(entry, fout, ensure_ascii=False)
             fout.write("\n")

As per coding guidelines, "Use type hints for simple types" and "perform all computations before re-opening files for writing to avoid accidental data loss if code fails during execution."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@nemo_skills/dataset/aime26/prepare.py` around lines 31 - 35, Add type hints
to write_data_to_file (e.g., def write_data_to_file(output_file: Path, data:
Iterable[dict]) -> None) and precompute formatted entries before opening the
file to avoid partial writes: build a list (e.g., formatted =
[format_entry(entry) for entry in data]) then open the output_file and iterate
over formatted with tqdm, json.dump each item and write a newline; reference the
existing write_data_to_file and format_entry symbols when making the change.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@nemo_skills/dataset/aime26/prepare.py`:
- Around line 38-43: The call in main that uses
load_dataset("MathArena/aime_2026", split="train") is hardcoded and should be
annotated so future readers understand this is intentional; update the main
function (around load_dataset, dataset variable) to add a concise clarifying
comment noting that the HF dataset only exposes a "train" split and therefore
args.split is not used for loading (but still used for output_file), so the
hardcoded split should remain; ensure the comment references args.split and
write_data_to_file to make the intent obvious.

---

Nitpick comments:
In `@nemo_skills/dataset/aime26/prepare.py`:
- Around line 31-35: Add type hints to write_data_to_file (e.g., def
write_data_to_file(output_file: Path, data: Iterable[dict]) -> None) and
precompute formatted entries before opening the file to avoid partial writes:
build a list (e.g., formatted = [format_entry(entry) for entry in data]) then
open the output_file and iterate over formatted with tqdm, json.dump each item
and write a newline; reference the existing write_data_to_file and format_entry
symbols when making the change.

[bzantium]

@Kipok I fixed the bug above and attached an evaluation.

Model Evaluation: Qwen3-30B-A3B-Thinking-2507

The following table summarizes the performance scores across different AIME (American Invitational Mathematics Examination) datasets.

Task	Score
aime24	91.00
aime25	82.33
aime26	87.67

---

Click the arrow below to expand the full technical breakdown, including pass@k and majority@n statistics.

View JSON Execution Data

{
  "aime26": {
    "pass@1": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.66666666666667,
      "no_answer": 0.0
    },
    "pass@2": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 91.25925925925927,
      "no_answer": 0.0
    },
    "pass@1[avg-of-2]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 91.66666666666667,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 17334.65,
        "std_dev_across_runs": 215.1725935150659,
        "avg_sample_std_dev": 2548.012145553645,
        "std_err_across_runs": 152.14999999999964
      },
      "symbolic_correct_statistics": {
        "avg": 0.9166666666666666,
        "std_dev_across_runs": 0.02357022603955158,
        "avg_sample_std_dev": 0.023570226039551587,
        "std_err_across_runs": 0.016666666666666663
      }
    },
    "pass@3": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 92.25,
      "no_answer": 0.0
    },
    "pass@1[avg-of-3]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.77777777777777,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 17662.333333333332,
        "std_dev_across_runs": 587.6042233794223,
        "avg_sample_std_dev": 3435.5203326507794,
        "std_err_across_runs": 339.25345654507043
      },
      "symbolic_correct_statistics": {
        "avg": 0.8777777777777778,
        "std_dev_across_runs": 0.06938886664887108,
        "avg_sample_std_dev": 0.07698003589195011,
        "std_err_across_runs": 0.04006168083848876
      }
    },
    "pass@4": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 92.76190476190476,
      "no_answer": 0.0
    },
    "pass@1[avg-of-4]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 86.66666666666667,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18331.816666666666,
        "std_dev_across_runs": 1422.3282145833991,
        "avg_sample_std_dev": 4074.4175305478043,
        "std_err_across_runs": 711.1641072916996
      },
      "symbolic_correct_statistics": {
        "avg": 0.8666666666666667,
        "std_dev_across_runs": 0.06085806194501844,
        "avg_sample_std_dev": 0.0859116756396542,
        "std_err_across_runs": 0.03042903097250922
      }
    },
    "pass@5": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.05555555555557,
      "no_answer": 0.0
    },
    "pass@1[avg-of-5]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 86.66666666666667,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18328.713333333333,
        "std_dev_across_runs": 1231.791912576507,
        "avg_sample_std_dev": 4199.473728063009,
        "std_err_across_runs": 550.8740901311095
      },
      "symbolic_correct_statistics": {
        "avg": 0.8666666666666667,
        "std_dev_across_runs": 0.05270462766947298,
        "avg_sample_std_dev": 0.08123619671700687,
        "std_err_across_runs": 0.02357022603955158
      }
    },
    "pass@6": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.22222222222221,
      "no_answer": 0.0
    },
    "pass@1[avg-of-6]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 86.66666666666669,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18596.200000000004,
        "std_dev_across_runs": 1281.8516894624663,
        "avg_sample_std_dev": 4426.99811124109,
        "std_err_across_runs": 523.3137608512666
      },
      "symbolic_correct_statistics": {
        "avg": 0.8666666666666668,
        "std_dev_across_runs": 0.04714045207910316,
        "avg_sample_std_dev": 0.08885962402813656,
        "std_err_across_runs": 0.019245008972987525
      }
    },
    "pass@7": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.30555555555556,
      "no_answer": 0.0
    },
    "pass@1[avg-of-7]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.14285714285714,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18356.109523809522,
        "std_dev_across_runs": 1331.4617977182336,
        "avg_sample_std_dev": 4374.646122160459,
        "std_err_across_runs": 503.2452567064904
      },
      "symbolic_correct_statistics": {
        "avg": 0.8714285714285714,
        "std_dev_across_runs": 0.044839513942303275,
        "avg_sample_std_dev": 0.08814386586003546,
        "std_err_across_runs": 0.016947743257192553
      }
    },
    "pass@8": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "pass@1[avg-of-8]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.5,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18257.5375,
        "std_dev_across_runs": 1263.8300049487398,
        "avg_sample_std_dev": 4177.041148241753,
        "std_err_across_runs": 446.83138338314086
      },
      "symbolic_correct_statistics": {
        "avg": 0.875,
        "std_dev_across_runs": 0.04272466296411038,
        "avg_sample_std_dev": 0.08346764803530463,
        "std_err_across_runs": 0.015105449452916093
      }
    },
    "pass@9": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "pass@1[avg-of-9]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.03703703703702,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18251.462962962964,
        "std_dev_across_runs": 1182.345167975312,
        "avg_sample_std_dev": 4085.528060506555,
        "std_err_across_runs": 394.1150559917707
      },
      "symbolic_correct_statistics": {
        "avg": 0.8703703703703703,
        "std_dev_across_runs": 0.042309850588132816,
        "avg_sample_std_dev": 0.0914098904572087,
        "std_err_across_runs": 0.014103283529377605
      }
    },
    "pass@10": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "pass@1[avg-of-10]": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 87.66666666666667,
      "no_answer": 0.0,
      "reasoning_tokens_statistics": {
        "avg": 0.0,
        "std_dev_across_runs": 0.0,
        "avg_sample_std_dev": 0.0,
        "std_err_across_runs": 0.0
      },
      "answer_tokens_statistics": {
        "avg": 18102.21,
        "std_dev_across_runs": 1210.5279378880143,
        "avg_sample_std_dev": 4033.7664918246287,
        "std_err_across_runs": 382.8025454992963
      },
      "symbolic_correct_statistics": {
        "avg": 0.8766666666666667,
        "std_dev_across_runs": 0.044583116994975275,
        "avg_sample_std_dev": 0.08953305854857022,
        "std_err_across_runs": 0.01409841948938835
      }
    },
    "majority@2": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 91.66666666666667,
      "no_answer": 0.0
    },
    "majority@3": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 91.1111111111111,
      "no_answer": 0.0
    },
    "majority@4": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 90.83333333333333,
      "no_answer": 0.0
    },
    "majority@5": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 90.66666666666667,
      "no_answer": 0.0
    },
    "majority@6": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 90.55555555555557,
      "no_answer": 0.0
    },
    "majority@7": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "majority@8": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "majority@9": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    },
    "majority@10": {
      "num_entries": 30,
      "avg_tokens": 18102,
      "gen_seconds": 6223,
      "symbolic_correct": 93.33333333333333,
      "no_answer": 0.0
    }
  }
}

[Kipok]

thanks!