"""

Evaluation script for IMO-Bench ProofBench dataset (60 problems)

Tests model performance on rigorous mathematical proof construction

Uses IMO25-style verification system for grading

"""

import argparse

import json

import os

import logging

import time

import re

import pandas as pd

import requests

from typing import List, Dict, Optional

from datetime import datetime

from openai import OpenAI

from tqdm import tqdm

# Add sys path to import optillm modules

import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Configure logging

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

# Dataset URL

PROOFBENCH_URL = "https://raw.githubusercontent.com/google-deepmind/superhuman/main/imobench/proofbench.csv"

SYSTEM_PROMPT = '''You are solving IMO-level mathematical proof problems. These require rigorous, complete proofs.

Key requirements:

1. **Complete proofs**: Provide rigorous, step-by-step mathematical proofs

2. **Mathematical rigor**: Every step must be logically justified

3. **Clear structure**: Organize your solution with clear logical flow

4. **Proper notation**: Use correct mathematical notation and formatting

5. **Verification**: Double-check your reasoning and conclusions

For different problem types:

- Functional equations: Consider injectivity, surjectivity, and special values

- Geometry: Use coordinate systems, trigonometry, or synthetic methods as appropriate

- Number theory: Apply divisibility, modular arithmetic, and prime factorization

- Combinatorics: Use counting techniques, pigeonhole principle, and extremal arguments

- Algebra: Handle polynomials, inequalities, and algebraic structures

Provide a complete, rigorous proof that addresses all aspects of the problem.'''

VERIFICATION_PROMPT = """You are an expert mathematician and IMO grader. Your task is to rigorously verify this mathematical solution.

**Grading Scale (0-7 points):**

- **7 points**: Complete, rigorous, correct proof

- **6 points**: Correct approach, minor gaps or notation issues

- **5 points**: Mostly correct, some gaps in rigor

- **3-4 points**: Significant progress, partial solution

- **1-2 points**: Some correct ideas, incomplete

- **0 points**: No progress or completely wrong

**Verification Instructions:**

1. Check logical correctness of each step

2. Verify mathematical rigor and completeness

3. Identify any critical errors or gaps

4. Assess proof structure and clarity

**Problem:**

{problem}

**Solution to verify:**

{solution}

Provide your assessment in the following format:

**SCORE:** [0-7]

**VERDICT:** [Correct/Partially Correct/Incorrect]

**REASONING:** [Detailed explanation of your assessment]

**CRITICAL ERRORS:** [List any critical errors found, or "None"]

**GAPS:** [List any gaps in rigor, or "None"]"""

def download_proofbench() -> pd.DataFrame:

"""

Download and parse the ProofBench CSV dataset

"""

logger.info("Downloading ProofBench dataset...")

try:

response = requests.get(PROOFBENCH_URL, timeout=30)

response.raise_for_status()

# Save to temp file and load with pandas

temp_file = "/tmp/proofbench.csv"

with open(temp_file, 'wb') as f:

f.write(response.content)

df = pd.read_csv(temp_file)

logger.info(f"Loaded {len(df)} problems from ProofBench")

return df

except Exception as e:

logger.error(f"Error downloading ProofBench: {e}")

raise

def verify_proof(problem: str, solution: str, grading_guidelines: str, model: str, client: OpenAI) -> Dict:

"""

Verify a proof using IMO25-style two-stage verification

Returns score on 0-7 scale and detailed assessment

"""

try:

# Format verification prompt

verification_text = VERIFICATION_PROMPT.format(

problem=problem,

solution=solution

)

# Add grading guidelines if available

if grading_guidelines and pd.notna(grading_guidelines):

verification_text += f"\n\n**Grading Guidelines:**\n{grading_guidelines}"

# Get verification response

response = client.chat.completions.create(

model=model,

messages=[

{"role": "system", "content": "You are an expert IMO grader. Provide rigorous assessment."},

{"role": "user", "content": verification_text}

],

max_tokens=4000,

temperature=0.1

)

verification_response = response.choices[0].message.content.strip()

# Extract score

score_match = re.search(r'\*\*SCORE:\*\*\s*(\d+)', verification_response)

score = int(score_match.group(1)) if score_match else 0

# Extract verdict

verdict_match = re.search(r'\*\*VERDICT:\*\*\s*([^\n]+)', verification_response)

verdict = verdict_match.group(1).strip() if verdict_match else "Unknown"

# Determine if correct (7 points = full marks)

is_correct = (score == 7)

# Check for critical errors

errors_match = re.search(r'\*\*CRITICAL ERRORS:\*\*\s*([^\n]+)', verification_response)

has_critical_errors = errors_match and "None" not in errors_match.group(1) if errors_match else False

return {

"score": score,

"verdict": verdict,

"is_correct": is_correct,

"has_critical_errors": has_critical_errors,

"verification_response": verification_response,

"success": True

}

except Exception as e:

logger.error(f"Error in proof verification: {e}")

return {

"score": 0,

"verdict": "Error",

"is_correct": False,

"has_critical_errors": True,

"verification_response": f"Verification error: {str(e)}",

"success": False

}

def extract_solution_quality(solution: str) -> Dict:

"""

Analyze the quality of a mathematical proof

"""

analysis = {

"has_proof_structure": False,

"uses_mathematical_notation": False,

"has_logical_steps": False,

"addresses_cases": False,

"has_conclusion": False,

"length_score": 0

}

if not solution:

return analysis

solution_lower = solution.lower()

# Check for proof structure

proof_keywords = ["proof:", "solution:", "we prove", "to show", "suppose", "assume", "let", "consider"]

if any(keyword in solution_lower for keyword in proof_keywords):

analysis["has_proof_structure"] = True

# Check for mathematical notation

math_patterns = [r'\$.*\$', r'\\[a-zA-Z]+', r'\\geq', r'\\leq', r'\\in', r'\\sum', r'\\prod']

if any(re.search(pattern, solution) for pattern in math_patterns):

analysis["uses_mathematical_notation"] = True

# Check for logical flow

logical_words = ["therefore", "thus", "hence", "consequently", "since", "because", "implies"]

logical_count = sum(1 for word in logical_words if word in solution_lower)

if logical_count >= 3:

analysis["has_logical_steps"] = True

# Check for case analysis

case_words = ["case", "if", "when", "suppose"]

case_count = sum(1 for word in case_words if word in solution_lower)

if case_count >= 2:

analysis["addresses_cases"] = True

# Check for conclusion

conclusion_words = ["therefore", "thus", "q.e.d", "qed", "proven", "concluded"]

if any(word in solution_lower for word in conclusion_words):

analysis["has_conclusion"] = True

# Length score (normalized)

analysis["length_score"] = min(len(solution) / 2000, 1.0)

return analysis

def get_llm_response(problem: str, model: str, client: OpenAI, extra_body: dict = None, timeout: int = 600) -> Dict:

"""

Get response from the LLM for a proof problem

"""

try:

kwargs = {}

if extra_body:

kwargs["extra_body"] = extra_body

response = client.with_options(timeout=timeout).chat.completions.create(

model=model,

messages=[

{"role": "system", "content": SYSTEM_PROMPT},

{"role": "user", "content": problem}

],

max_tokens=64000, # Extended for complex proofs

temperature=0.1,

**kwargs

)

solution_text = response.choices[0].message.content.strip()

reasoning_tokens = getattr(response.usage, 'reasoning_tokens', 0)

total_tokens = response.usage.total_tokens if hasattr(response.usage, 'total_tokens') else 0

return {

"solution": solution_text,

"reasoning_tokens": reasoning_tokens,

"total_tokens": total_tokens,

"success": True

}

except Exception as e:

logger.error(f"Error getting LLM response: {e}")

return {

"solution": f"Error: {str(e)}",

"reasoning_tokens": 0,

"total_tokens": 0,

"success": False

}

def save_result(filename: str, result: Dict):

"""Save a single result with incremental updates"""

results = []

if os.path.exists(filename):

try:

with open(filename, 'r') as f:

results = json.load(f)

except (FileNotFoundError, json.JSONDecodeError):

results = []

results.append(result)

with open(filename, 'w') as f:

json.dump(results, f, indent=2)

def load_existing_results(filename: str) -> List[Dict]:

"""Load existing results from file"""

try:

with open(filename, 'r') as f:

return json.load(f)

except FileNotFoundError:

return []

def calculate_subset_scores(results: List[Dict]) -> Dict:

"""

Calculate full credit scores for various subsets (Novel, IMO 2024, USAMO 2025)

Returns dictionary with subset names and their (solved, total, percentage) tuples

"""

subsets = {

'Novel': {'full': 0, 'total': 0},

'IMO 2024': {'full': 0, 'total': 0},

'USAMO 2025': {'full': 0, 'total': 0}

}

for r in results:

source = r.get('source', '')

score = r['verification'].get('score', 0)

is_full = (score == 7)

# Categorize by source

if 'Novel Problem' in source:

subsets['Novel']['total'] += 1

if is_full:

subsets['Novel']['full'] += 1

elif 'IMO 2024' in source:

subsets['IMO 2024']['total'] += 1

if is_full:

subsets['IMO 2024']['full'] += 1

elif 'USAMO 2025' in source:

subsets['USAMO 2025']['total'] += 1

if is_full:

subsets['USAMO 2025']['full'] += 1

# Calculate percentages

subset_stats = {}

for name, counts in subsets.items():

total = counts['total']

full = counts['full']

pct = (full / total * 100) if total > 0 else 0

subset_stats[name] = (full, total, pct)

return subset_stats

def analyze_results(results: List[Dict]):

"""Analyze and print comprehensive statistics with full credit prioritized"""

if not results:

print("No results to analyze")

return

total_problems = len(results)

full_marks = sum(1 for r in results if r['verification'].get('score', 0) == 7)

partial_credit = sum(1 for r in results if 1 <= r['verification'].get('score', 0) <= 6)

no_credit = total_problems - full_marks - partial_credit

total_score = sum(r['verification'].get('score', 0) for r in results)

avg_score = total_score / total_problems

print("\n" + "="*80)

print("IMO-Bench ProofBench Evaluation Results")

print("="*80)

# ========================================================================

# SECTION 1: FULL CREDIT SCORES (PRIMARY METRIC)

# ========================================================================

print("\n" + "="*80)

print("FULL CREDIT SCORES (7/7 = Solved) - PRIMARY METRIC")

print("="*80)

print(f"\nOverall: {full_marks}/{total_problems} = {full_marks/total_problems*100:.1f}%")

# Basic vs Advanced breakdown (full credit only)

basic_full = sum(1 for r in results if 'Basic' in r.get('problem_id', '') and r['verification'].get('score', 0) == 7)

basic_total = sum(1 for r in results if 'Basic' in r.get('problem_id', ''))

adv_full = sum(1 for r in results if 'Advanced' in r.get('problem_id', '') and r['verification'].get('score', 0) == 7)

adv_total = sum(1 for r in results if 'Advanced' in r.get('problem_id', ''))

print(f"\nBasic problems: {basic_full}/{basic_total} = {basic_full/basic_total*100 if basic_total > 0 else 0:.1f}%")

print(f"Advanced problems: {adv_full}/{adv_total} = {adv_full/adv_total*100 if adv_total > 0 else 0:.1f}%")

# ========================================================================

# SECTION 2: SUBSET BREAKDOWN (Novel, IMO 2024, USAMO 2025)

# ========================================================================

subset_stats = calculate_subset_scores(results)

if any(total > 0 for _, total, _ in subset_stats.values()):

print("\n" + "-"*80)

print("Subset Breakdown (Full Credit Only):")

print("-"*80)

for name in ['Novel', 'IMO 2024', 'USAMO 2025']:

full, total, pct = subset_stats[name]

if total > 0:

print(f"{name:15s}: {full}/{total} = {pct:.1f}%")

# ========================================================================

# SECTION 3: DETAILED ANALYSIS (Average Scores and Distributions)

# ========================================================================

print("\n" + "="*80)

print("DETAILED ANALYSIS (Average Scores)")

print("="*80)

print(f"\nAverage score: {avg_score:.2f}/7 ({avg_score/7*100:.1f}%)")

print(f"Full credit (7/7): {full_marks} ({full_marks/total_problems*100:.1f}%)")

print(f"Partial credit (1-6): {partial_credit} ({partial_credit/total_problems*100:.1f}%)")

print(f"No credit (0): {no_credit} ({no_credit/total_problems*100:.1f}%)")

# Basic vs Advanced (average scores)

basic_scores = [r['verification'].get('score', 0) for r in results if 'Basic' in r.get('problem_id', '')]

adv_scores = [r['verification'].get('score', 0) for r in results if 'Advanced' in r.get('problem_id', '')]

if basic_scores or adv_scores:

print("\n" + "-"*80)

print("Basic vs Advanced (Average Scores):")

print("-"*80)

if basic_scores:

basic_avg = sum(basic_scores) / len(basic_scores)

print(f"Basic ({len(basic_scores)}): {basic_avg:.2f}/7 ({basic_avg/7*100:.1f}%)")

if adv_scores:

adv_avg = sum(adv_scores) / len(adv_scores)

print(f"Advanced ({len(adv_scores)}): {adv_avg:.2f}/7 ({adv_avg/7*100:.1f}%)")

# Category breakdown (average scores)

categories = {}

for r in results:

cat = r.get('category', 'Unknown')

if cat not in categories:

categories[cat] = {'total': 0, 'scores': [], 'full': 0}

categories[cat]['total'] += 1

score = r['verification'].get('score', 0)

categories[cat]['scores'].append(score)

if score == 7:

categories[cat]['full'] += 1

if categories:

print("\n" + "-"*80)

print("Performance by Category:")

print("-"*80)

for cat, stats in sorted(categories.items()):

avg = sum(stats['scores']) / len(stats['scores']) if stats['scores'] else 0

full = stats['full']

total = stats['total']

print(f"{cat:20s}: Avg {avg:.2f}/7 ({avg/7*100:.1f}%) | Solved {full}/{total} ({full/total*100:.1f}%)")

# Level breakdown (average scores)

levels = {}

for r in results:

level = r.get('level', 'Unknown')

if level not in levels:

levels[level] = {'total': 0, 'scores': [], 'full': 0}

levels[level]['total'] += 1

score = r['verification'].get('score', 0)

levels[level]['scores'].append(score)

if score == 7:

levels[level]['full'] += 1

if levels:

print("\n" + "-"*80)

print("Performance by Level:")

print("-"*80)

for level, stats in sorted(levels.items()):

avg = sum(stats['scores']) / len(stats['scores']) if stats['scores'] else 0

full = stats['full']

total = stats['total']

print(f"{level:20s}: Avg {avg:.2f}/7 ({avg/7*100:.1f}%) | Solved {full}/{total} ({full/total*100:.1f}%)")

# Token statistics

try:

total_tokens = sum(r['response'].get('total_tokens', 0) for r in results)

reasoning_tokens = sum(r['response'].get('reasoning_tokens', 0) for r in results)

print("\n" + "-"*80)

print("Token Statistics:")

print("-"*80)

print(f"Total tokens: {total_tokens:,}")

print(f"Reasoning tokens: {reasoning_tokens:,}")

print(f"Avg tokens per problem: {total_tokens/total_problems:.0f}")

except (KeyError, TypeError):

pass # Skip token stats if data not available

# Time statistics

total_time = sum(r.get('solve_time_seconds', 0) for r in results)

print("\n" + "-"*80)

print(f"Total solve time: {total_time:.1f}s ({total_time/60:.1f} minutes, {total_time/3600:.1f} hours)")

print(f"Avg time per problem: {total_time/total_problems:.1f}s")

print("="*80 + "\n")

def main():

parser = argparse.ArgumentParser(description="Evaluate on IMO-Bench ProofBench")

parser.add_argument("--model", type=str, required=True,

help="Model to use (e.g., google/gemini-2.5-flash-preview-09-2025 or mars-...)")

parser.add_argument("--base-url", type=str, default="http://localhost:8001/v1",

help="Base URL for OptiLLM server")

parser.add_argument("--verifier-model", type=str, default=None,

help="Model to use for verification (defaults to same as solver)")

parser.add_argument("--timeout", type=int, default=600,

help="Timeout in seconds for each problem")

parser.add_argument("--limit", type=int, default=None,

help="Limit number of problems (for testing)")

parser.add_argument("--subset", type=str, default=None,

help="Evaluate only 'basic' or 'advanced' subset")

args = parser.parse_args()

# Initialize OpenAI client

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

# Verifier model defaults to solver model

verifier_model = args.verifier_model or args.model.replace("mars-", "")

# Setup results directory

os.makedirs("results", exist_ok=True)

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

# Determine if using MARS

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_proof_{approach_name}_{model_name.replace('/', '_')}_{timestamp}.json"

# Download dataset

df = download_proofbench()

# Filter by subset if specified

if args.subset:

if args.subset.lower() == 'basic':

df = df[df['Level'].str.contains('Basic', case=False, na=False)]

elif args.subset.lower() == 'advanced':

df = df[df['Level'].str.contains('Advanced', case=False, na=False)]

print(f"Filtered to {args.subset} subset")

# Limit problems if specified

if args.limit:

df = df.head(args.limit)

print(f"\nEvaluating {len(df)} ProofBench problems")

print(f"Model: {args.model}")

print(f"Approach: {approach_name}")

print(f"Verifier: {verifier_model}")

if is_mars:

print("MARS Config: use_thinking_tags=False, answer_extraction_mode='none'")

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 proof problems"):

problem_id = row.get('Problem ID', f'problem_{idx}')

problem_text = row['Problem']

reference_solution = row.get('Solution', '')

grading_guidelines = row.get('Grading guidelines', '')

category = row.get('Category', 'Unknown')

level = row.get('Level', 'Unknown')

source = row.get('Source', 'Unknown')

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

start_time = time.time()

# Get LLM response (model prefix handles MARS configuration automatically)

response = get_llm_response(

problem_text,

args.model,

client,

extra_body=None,

timeout=args.timeout

)

solve_time = time.time() - start_time

# Verify the proof

verification = verify_proof(

problem_text,

response['solution'],

grading_guidelines,

verifier_model,

client

)

# Analyze solution quality

quality = extract_solution_quality(response['solution'])

# Compile result

result = {

"timestamp": datetime.now().isoformat(),

"model": args.model,

"approach": approach_name,

"verifier_model": verifier_model,

"problem_id": problem_id,

"category": category,

"level": level,

"source": source,

"problem": problem_text,

"reference_solution": reference_solution,

"grading_guidelines": grading_guidelines,

"response": response,

"verification": verification,

"quality": quality,

"solve_time_seconds": solve_time

}

# Save result immediately

save_result(results_file, result)

score = verification.get('score', 0)

status = "✓ SOLVED" if score == 7 else f"✗ {score}/7"

logger.info(f"{status} {problem_id}")

# Load and analyze all results

print("\n" + "="*80)

print("Evaluation Complete!")

print("="*80)

results = load_existing_results(results_file)

analyze_results(results)

print(f"Results saved to: {results_file}")

if __name__ == "__main__":

main()