Load AIME problems from HuggingFace. Supports: - dataset="aime24" → AIME 2024 (30 problems) - dataset="aime25" → AIME 2025 I & II (30 problems) - dataset="all" → Both datasets (60 problems)
(dataset="all", problem_ids=None, limit=None)
| 36 | |
| 37 | |
| 38 | def _load_problems(dataset="all", problem_ids=None, limit=None) -> List[dict]: |
| 39 | """Load AIME problems from HuggingFace. |
| 40 | |
| 41 | Supports: |
| 42 | - dataset="aime24" → AIME 2024 (30 problems) |
| 43 | - dataset="aime25" → AIME 2025 I & II (30 problems) |
| 44 | - dataset="all" → Both datasets (60 problems) |
| 45 | """ |
| 46 | from datasets import load_dataset |
| 47 | |
| 48 | datasets_to_load = [dataset] if dataset != "all" else list(DATASETS.keys()) |
| 49 | |
| 50 | problems = [] |
| 51 | for ds_key in datasets_to_load: |
| 52 | ds_info = DATASETS[ds_key] |
| 53 | |
| 54 | for config_name in ds_info["configs"]: |
| 55 | # Load from HuggingFace |
| 56 | if config_name is not None: |
| 57 | ds = load_dataset( |
| 58 | ds_info["hf_path"], config_name, split=ds_info["split"] |
| 59 | ) |
| 60 | else: |
| 61 | ds = load_dataset(ds_info["hf_path"], split=ds_info["split"]) |
| 62 | |
| 63 | for row_idx, row in enumerate(ds): |
| 64 | # Build problem ID |
| 65 | if ds_info["id_field"] and ds_info["id_field"] in row: |
| 66 | pid = str(row[ds_info["id_field"]]) |
| 67 | else: |
| 68 | # Generate ID from config name and index |
| 69 | # e.g. "AIME2025-I" → "2025-I-1" |
| 70 | if config_name: |
| 71 | prefix = config_name.replace("AIME", "") |
| 72 | pid = f"{prefix}-{row_idx + 1}" |
| 73 | else: |
| 74 | pid = f"{ds_key}-{row_idx + 1}" |
| 75 | |
| 76 | question = row[ds_info["question_field"]] |
| 77 | answer = str(row[ds_info["answer_field"]]) |
| 78 | |
| 79 | # Determine dataset label for grouping |
| 80 | if config_name: |
| 81 | dataset_label = f"{ds_key}-{config_name.split('-')[-1]}" |
| 82 | else: |
| 83 | dataset_label = ds_key |
| 84 | |
| 85 | problems.append( |
| 86 | { |
| 87 | "id": pid, |
| 88 | "problem": question, |
| 89 | "answer": answer, |
| 90 | "dataset": dataset_label, |
| 91 | } |
| 92 | ) |
| 93 | |
| 94 | # Apply filters |
| 95 | if problem_ids: |