| 586 | |
| 587 | |
| 588 | def parse_args() -> argparse.Namespace: |
| 589 | parser = argparse.ArgumentParser( |
| 590 | description="Compare direct reasoning vs reflection-optimized solutions on AIME datasets." |
| 591 | ) |
| 592 | parser.add_argument( |
| 593 | "--models", |
| 594 | nargs="+", |
| 595 | default=["gpt-4.1"], |
| 596 | help="List of base model names to evaluate.", |
| 597 | ) |
| 598 | parser.add_argument( |
| 599 | "--datasets", |
| 600 | nargs="+", |
| 601 | default=["AIME24", "AIME25"], |
| 602 | choices=list(DATASET_REGISTRY.keys()), |
| 603 | help="Datasets to evaluate.", |
| 604 | ) |
| 605 | parser.add_argument( |
| 606 | "--max-step", |
| 607 | type=int, |
| 608 | default=3, |
| 609 | help="Maximum number of reflection update steps per sample.", |
| 610 | ) |
| 611 | parser.add_argument( |
| 612 | "--max-samples", |
| 613 | type=int, |
| 614 | default=None, |
| 615 | help="Optional limit on number of samples per dataset.", |
| 616 | ) |
| 617 | parser.add_argument( |
| 618 | "--output", |
| 619 | type=Path, |
| 620 | default=Path("logs/aime_reflection_experiment_4.1.json"), |
| 621 | help="Path to store the experiment log (JSON).", |
| 622 | ) |
| 623 | parser.add_argument( |
| 624 | "--use-tool-agent", |
| 625 | action="store_true", |
| 626 | help="Use the tool-calling agent for direct inference instead of the base model.", |
| 627 | ) |
| 628 | parser.add_argument( |
| 629 | "--agent-config", |
| 630 | type=Path, |
| 631 | default=Path("configs/tool_calling_agent.py"), |
| 632 | help="Config file for the tool-calling agent (used when --use-tool-agent is set).", |
| 633 | ) |
| 634 | parser.add_argument( |
| 635 | "--agent-name", |
| 636 | default="tool_calling", |
| 637 | help="Agent registered in ACP to use for inference (requires --use-tool-agent).", |
| 638 | ) |
| 639 | parser.add_argument( |
| 640 | "--agent-cfg-options", |
| 641 | nargs="+", |
| 642 | action=DictAction, |
| 643 | help="Override options for the agent config, same format as mmengine DictAction.", |
| 644 | ) |
| 645 | return parser.parse_args() |