Function main

examples/run_aime_reflection_experiment.py:648–729

()

def main() -> None:
    load_dotenv(override=True)
    args = parse_args()
    ensure_output_path(args.output)

    experiment_meta = {
        "timestamp": datetime.utcnow().isoformat() + "Z",
        "system_prompt": INFERENCE_SYSTEM_PROMPT,
        "evaluation_instruction": EVALUATION_INSTRUCTION,
        "evaluation_system_prompt": EVALUATION_SYSTEM_PROMPT,
        "max_step": args.max_step,
        "max_samples": args.max_samples,
    }

    if args.output.exists():
        with args.output.open("r", encoding="utf-8") as f:
            results: Dict[str, Any] = json.load(f)
    else:
        results = {}

    results.setdefault("meta", {})
    results["meta"].update(experiment_meta)
    results.setdefault("models", {})

    save_results(results, args.output)

    agent_runner: Optional[ToolCallingAgentRunner] = None
    if args.use_tool_agent:
        try:
            agent_runner = build_tool_calling_agent_runner(
                agent_name=args.agent_name,
                config_path=args.agent_config,
                cfg_options=args.agent_cfg_options,
            )
        except Exception as exc:
            print(f"[工具代理初始化失败] {exc}")
            sys.exit(1)

    for model_name in args.models:
        try:
            engine = tg.get_engine(model_name)
        except ValueError as exc:
            print(f"[跳过模型] {model_name}: {exc}")
            results["models"][model_name] = {"error": str(exc)}
            save_results(results, args.output)
            continue

        system_prompt_var = Variable(
            INFERENCE_SYSTEM_PROMPT,
            requires_grad=False,
            role_description="system prompt specifying reasoning format",
        )
        if agent_runner:
            model = ToolCallingAgentLLM(agent_runner, system_prompt_var)
        else:
            model = BlackboxLLM(engine=engine, system_prompt=system_prompt_var)

        model_results = results["models"].setdefault(model_name, {})