()
| 191 | |
| 192 | |
| 193 | def main(): |
| 194 | args = parse_args() |
| 195 | |
| 196 | # Step 1: Load problems from HuggingFace |
| 197 | ds_label = args.dataset if args.dataset != "all" else "AIME 2024 + 2025" |
| 198 | print(f"Loading {ds_label} from HuggingFace...") |
| 199 | problems = _load_problems( |
| 200 | dataset=args.dataset, |
| 201 | problem_ids=args.problem_ids, |
| 202 | limit=args.limit, |
| 203 | ) |
| 204 | print(f"\n{len(problems)} problems to process") |
| 205 | |
| 206 | # Step 2: Run agent on each problem |
| 207 | output_dir = Path(args.output_dir) / "runs" |
| 208 | output_dir.mkdir(parents=True, exist_ok=True) |
| 209 | |
| 210 | # Resume: load existing completed results and skip them |
| 211 | existing_results = {} |
| 212 | if args.resume: |
| 213 | existing_results = _load_existing_results(output_dir) |
| 214 | if existing_results: |
| 215 | print( |
| 216 | f"Resuming: {len(existing_results)} completed problems found, skipping them" |
| 217 | ) |
| 218 | |
| 219 | results: List[AIMEResult] = [] |
| 220 | to_run = [] |
| 221 | for i, p in enumerate(problems): |
| 222 | if p["id"] in existing_results: |
| 223 | results.append(existing_results[p["id"]]) |
| 224 | else: |
| 225 | to_run.append((i, p)) |
| 226 | |
| 227 | if not to_run: |
| 228 | print("All problems already completed!") |
| 229 | elif args.max_parallel > 1: |
| 230 | with ThreadPoolExecutor(max_workers=args.max_parallel) as executor: |
| 231 | futures = { |
| 232 | executor.submit( |
| 233 | _run_single_problem, i + 1, len(problems), p, args, output_dir |
| 234 | ): p |
| 235 | for i, p in to_run |
| 236 | } |
| 237 | for future in as_completed(futures): |
| 238 | try: |
| 239 | results.append(future.result()) |
| 240 | except Exception as e: |
| 241 | p = futures[future] |
| 242 | results.append( |
| 243 | AIMEResult( |
| 244 | problem_id=p["id"], |
| 245 | problem=p["problem"], |
| 246 | correct_answer=p["answer"], |
| 247 | dataset=p.get("dataset", ""), |
| 248 | status="failed", |
| 249 | error=str(e), |
| 250 | ) |
no test coverage detected