Function main

scripts/eval_imobench_answer.py:295–403

()

def main():
    parser = argparse.ArgumentParser(description="Evaluate on IMO-Bench AnswerBench")
    parser.add_argument("--model", type=str, required=True,
                       help="Model to use (e.g., google/gemini-2.5-flash-preview-09-2025)")
    parser.add_argument("--base-url", type=str, default="http://localhost:8001/v1",
                       help="Base URL for OptiLLM server")
    parser.add_argument("--timeout", type=int, default=300,
                       help="Timeout in seconds for each problem")
    parser.add_argument("--limit", type=int, default=None,
                       help="Limit number of problems to evaluate (for testing)")
    parser.add_argument("--categories", type=str, default=None,
                       help="Comma-separated list of categories to evaluate (e.g., 'Algebra,Geometry')")

    args = parser.parse_args()

    # Initialize OpenAI client
    client = OpenAI(api_key="optillm", base_url=args.base_url)

    # Setup results directory
    os.makedirs("results", exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # Determine if using MARS approach
    is_mars = args.model.startswith("mars-")
    approach_name = "mars" if is_mars else "baseline"
    model_name = args.model.replace("mars-", "") if is_mars else args.model

    results_file = f"results/imobench_answer_{approach_name}_{model_name.replace('/', '_')}_{timestamp}.json"

    # Download dataset
    df = download_answerbench()

    # Filter by categories if specified
    if args.categories:
        selected_cats = [c.strip() for c in args.categories.split(',')]
        df = df[df['Category'].isin(selected_cats)]
        print(f"Filtered to categories: {selected_cats}")

    # Limit problems if specified
    if args.limit:
        df = df.head(args.limit)

    print(f"\nEvaluating {len(df)} AnswerBench problems")
    print(f"Model: {args.model}")
    print(f"Approach: {approach_name}")
    print(f"Results will be saved to: {results_file}\n")

    # Evaluate each problem
    for idx, row in tqdm(df.iterrows(), total=len(df), desc="Solving problems"):
        problem_id = row.get('Problem ID', f'problem_{idx}')
        problem_text = row['Problem']
        ground_truth = row['Short Answer']
        category = row.get('Category', 'Unknown')
        subcategory = row.get('Subcategory', '')
        difficulty = row.get('Difficulty', '')

        logger.info(f"Evaluating {problem_id}: {category}")