Add MMLU-Pro 10% optimized subset for checkpoint selection

[ka00ri]

PR to support subsampling of MMLU-Pro benchmark.

• Add hash-based subset filtering for MMLU-Pro.
• Add support for 10pct_opt_v1 subset (1,203 questions) optimized for checkpoint selection while preserving topic balance and full score tracking
• Add utility funcs for question/options normalization, subset ID loading, and filtering
• Validate e2e with:
  • ns prepare_data mmlu-pro --split=10pct_opt_v1
  • ns eval --benchmarks=mmlu-pro --split=10pct_opt_v1 ...
• Raise ValueError when subset IDs do not fully match source split entries

Summary by CodeRabbit

• New Features

  • Subset-based split support for dataset preparation, allowing filtered dataset creation.
  • Added a new "10pct_opt_v1" dataset subset.
  • Utilities for loading and applying subsets to datasets based on normalized question matching.

• Improvements

  • CLI extended to accept the new split identifier.
  • Added validation with descriptive errors when subset files are missing.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a82a6a7 and e5044b1.

📒 Files selected for processing (2)

• nemo_skills/dataset/mmlu-pro/subsets/10pct_opt_v1.txt
• nemo_skills/dataset/utils.py

🚧 Files skipped from review as they are similar to previous changes (1)

• nemo_skills/dataset/mmlu-pro/subsets/10pct_opt_v1.txt

---

📝 Walkthrough

Walkthrough

Adds subset-based filtering to the MMLU-Pro preparation flow: new hashing and subset utilities, a 10pct_opt_v1 subset file, and prepare.py updated to load and apply subset IDs for non-validation/test splits before writing output.

Changes

Cohort / File(s)	Summary
Subset Utility Functions nemo_skills/dataset/utils.py	Added get_question_hash(question, options=None), load_subset_ids(ids_file), and filter_by_subset(dataset, subset_ids, question_key="question", options_key=None); introduced hashlib import.
MMLU-Pro Preparation Script nemo_skills/dataset/mmlu-pro/prepare.py	Extended main to accept an additional split (10pct_opt_v1), load subsets/{split}.txt for non-validation/test splits, validate file existence, load subset IDs, filter the test dataset via filter_by_subset, and pass the filtered dataset to the existing write flow.
Subset Definition Data nemo_skills/dataset/mmlu-pro/subsets/10pct_opt_v1.txt	Added a new subset file containing 1203 JSONL entries (each with an id field) defining the 10pct_opt_v1 subset.

Sequence Diagram(s)

sequenceDiagram
    participant prepare as prepare.py
    participant subset as Subset File
    participant utils as utils.py
    participant dataset as Dataset
    participant output as Output

    prepare->>subset: check subset file exists for split
    subset-->>prepare: file path/result
    prepare->>utils: load_subset_ids(subset_file)
    utils-->>prepare: subset_ids
    prepare->>dataset: load test dataset
    dataset-->>prepare: dataset entries
    prepare->>utils: filter_by_subset(dataset, subset_ids)
    utils->>utils: compute get_question_hash for entries
    utils-->>prepare: filtered dataset (ordered by subset_ids)
    prepare->>output: write_data_to_file(filtered_dataset)

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and concisely describes the main change: adding a 10% optimized subset for MMLU-Pro checkpoint selection, which is the primary objective of the PR.

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

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch mboubdir/subsampling-benchmarks

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

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: 2

🧹 Nitpick comments (1)

nemo_skills/dataset/mmlu-pro/prepare.py (1)

66-68: CLI choices block future subset files even though main already supports them.

Hard-coding "10pct_opt_v1" in argparse means every new subset requires code changes instead of just adding a file in subsets/.

Proposed fix

-    parser.add_argument(
-        "--split", default="test", choices=("validation", "test", "10pct_opt_v1"), help="Dataset split to process."
-    )
+    parser.add_argument("--split", default="test", help="Dataset split to process.")

🤖 Prompt for AI Agents

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

In `@nemo_skills/dataset/mmlu-pro/prepare.py` around lines 66 - 68, The CLI
currently hard-codes the --split choices in parser.add_argument (the
("validation","test","10pct_opt_v1") tuple), which prevents adding new subset
files without changing code; update the parser to either remove the static
choices argument or dynamically build the choices list from the subsets
directory (e.g., read filenames in the subsets/ folder and derive allowed split
names) and pass that list into parser.add_argument for --split; ensure the code
path that uses the split (main or functions that load subsets) still
validates/handles unknown splits gracefully if you choose to remove argparse
choices.

🤖 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/utils.py`:
- Around line 323-327: The loop building hash_to_entry silently overwrites
duplicates from dataset: in the block using
get_question_hash(entry[question_key], options) (where options_key,
question_key, dataset are used), detect if h already exists in hash_to_entry and
handle it explicitly (e.g., raise a ValueError with both entries and the
duplicate hash, or log an error and skip/keep the first entry) instead of
blindly assigning hash_to_entry[h] = entry; update the code in the hash_to_entry
construction to check "if h in hash_to_entry" and then perform the chosen
explicit duplicate-handling behavior so subset mapping is unambiguous.
- Around line 309-312: The hash preimage construction conflates question/options
boundaries by joining parts with spaces; change it to build an unambiguous
preimage (e.g., start with a marker for the question and a separator before
options) instead of plain " ".join(parts): explicitly include a question prefix
and delimiters and/or option indices when concatenating (for example build
strings like f"Q:{normalized_question}" and for each option
f"O{i}:{normalized_option}" or include a count), then pass that canonical
preimage to hashlib.sha256(...).hexdigest() so question/options boundaries
(variables question, options, parts and the hashlib.sha256 call) are
unambiguous.

---

Nitpick comments:
In `@nemo_skills/dataset/mmlu-pro/prepare.py`:
- Around line 66-68: The CLI currently hard-codes the --split choices in
parser.add_argument (the ("validation","test","10pct_opt_v1") tuple), which
prevents adding new subset files without changing code; update the parser to
either remove the static choices argument or dynamically build the choices list
from the subsets directory (e.g., read filenames in the subsets/ folder and
derive allowed split names) and pass that list into parser.add_argument for
--split; ensure the code path that uses the split (main or functions that load
subsets) still validates/handles unknown splits gracefully if you choose to
remove argparse choices.

---

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d4ef8c0 and a82a6a7.

📒 Files selected for processing (3)

• nemo_skills/dataset/mmlu-pro/prepare.py
• nemo_skills/dataset/mmlu-pro/subsets/10pct_opt_v1.txt
• nemo_skills/dataset/utils.py

[gwarmstrong]

@ka00ri just apply suggestions from code rabbit then I think everything looks good