Function run_sample

examples/run_aime_reflection_experiment.py:381–480

(
    *,
    question_text: str,
    ground_truth_text: str,
    model: Module,  # BlackboxLLM or ToolCallingAgentLLM
    evaluation_module: MultiFieldEvaluation,
    engine: tg.EngineLM,
    model_name: str,
    max_step: int,
)

def run_sample(
    *,
    question_text: str,
    ground_truth_text: str,
    model: Module,  # BlackboxLLM or ToolCallingAgentLLM
    evaluation_module: MultiFieldEvaluation,
    engine: tg.EngineLM,
    model_name: str,
    max_step: int,
) -> Dict[str, Any]:
    question_var = Variable(
        question_text,
        requires_grad=False,
        role_description="contest problem statement",
    )

    # Direct inference
    direct_response = model(question_var)
    direct_response_text = direct_response.value
    direct_eval = evaluate_prediction(direct_response_text, ground_truth_text)
    direct_think_stats = compute_think_token_stats(model_name, direct_response_text)

    # Reflection loop
    current_solution = direct_response_text
    step_logs: List[Dict[str, Any]] = []

    for step_idx in range(1, max_step + 1):
        solution_var = Variable(
            current_solution,
            requires_grad=False,
            role_description="candidate solution with reasoning and boxed answer",
        )
        evaluation_output = evaluation_module([question_var, solution_var])
        evaluation_text = evaluation_output.value

        reflection_prompt = build_reflection_prompt(
            question_text=question_text,
            current_solution=current_solution,
            evaluation_text=evaluation_text,
        )

        try:
            improved_text = engine(
                reflection_prompt,
                system_prompt=REFLECTION_SYSTEM_PROMPT,
            )
        except Exception as exc:  # pragma: no cover - runtime safety
            logger.error(f"Reflection step {step_idx} failed: {exc}")
            step_logs.append(
                {
                    "step": step_idx,
                    "error": f"reflection_failed: {exc}",
                    "evaluation_output": evaluation_text,
                    "reflection_prompt": reflection_prompt,
                }
            )
            break