test_aime_2024_samples.py

#!/usr/bin/env python3
"""
Test script for AIME 2024 samples using the adaptive CoT framework.
"""

import argparse
import json
import time
from pathlib import Path

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

from src.benchmarks.math_benchmarks import MathBenchmarkLoader
from src.adaptive.adaptive_cot import AdaptiveCoT
from src.models.generic_model import GenericModel as HuggingFaceModel


def test_aime_2024_samples(
    model_path: str,
    branches: int = 8,
    samples: int = 3,
    output_file: str = "results/aime_2024_test.json",
    mode: str = "static"
):
    """Test AIME 2024 samples with the adaptive CoT framework."""
    
    print("🔬 AIME 2024 Sample Testing")
    print("=" * 60)
    print(f"Mode: {mode}")
    print(f"Branch Count: {branches}")
    print(f"Number of Samples: {samples}")
    print(f"Output File: {output_file}")
    print("=" * 60)
    
    # Load model
    print("🔧 Loading model...")
    model_config = {
        "model_config": {},
        "generation_params": {}
    }
    model = HuggingFaceModel(model_path, model_config)
    model.load_model()
    print(f"Model loaded successfully on {model.device}")
    
    # Load AIME 2024 dataset
    print("📚 Loading AIME 2024 dataset...")
    benchmark_loader = MathBenchmarkLoader()
    dataset = benchmark_loader.load_dataset("aime_2024", max_samples=samples)
    
    if len(dataset) < samples:
        print(f"📊 Limited to {len(dataset)} samples")
    else:
        print(f"📊 Using all {len(dataset)} samples")
    
    # Configure adaptive CoT
    config = {
        "adaptive_branching": mode == "adaptive",
        "min_branches": branches,
        "max_branches": branches,
        "default_branches": branches,
        "num_fewshot": 0,  # Zero-shot evaluation
        "temperature": 0.7,  # Non-deterministic for self-consistency
        "top_p": 0.95,
        "max_tokens": 32768,
        "benchmark": "aime_2024"  # Use AIME 2024-specific answer extraction
    }
    
    print(f"⚙️  Configuration:")
    for key, value in config.items():
        print(f"   {key}: {value}")
    print()
    
    # Initialize adaptive CoT
    adaptive_cot = AdaptiveCoT(model, config)
    
    # Test problems
    results = []
    total_start_time = time.time()
    
    print("🚀 Starting evaluation...")
    
    for i, sample in enumerate(dataset, 1):
        print(f"📝 Problem {i}/{len(dataset)}: {sample['question'][:100]}...")
        
        # Solve problem
        start_time = time.time()
        result = adaptive_cot.solve_problem(sample['question'])
        solve_time = time.time() - start_time
        
        # Extract ground truth answer
        ground_truth_answer = str(sample['answer']).strip()
        
        # Check correctness
        is_correct = result['final_answer'] == ground_truth_answer
        
        # Store result
        problem_result = {
            "problem_id": i,
            "question": sample['question'],
            "ground_truth": ground_truth_answer,
            "predicted_answer": result['final_answer'],
            "correct": is_correct,
            "solve_time": solve_time,
            "num_branches": result['num_branches'],
            "consensus_confidence": result.get('consensus_confidence', 0.0),
            "reasoning_paths": result['reasoning_paths'],
            "extracted_answers": result['extracted_answers']
        }
        results.append(problem_result)
        
        # Print result
        print(f"🎯 Final answer: {result['final_answer']}")
        print(f"📊 Consensus confidence: {result.get('consensus_confidence', 0.0):.3f}")
        print(f"   Answer: {result['final_answer']}")
        print(f"   Ground Truth: {ground_truth_answer}")
        print(f"   Correct: {'✅' if is_correct else '❌'}")
        print(f"   Confidence: {result.get('consensus_confidence', 0.0):.3f}")
        print()
    
    total_time = time.time() - total_start_time
    
    # Calculate statistics
    correct_count = sum(1 for r in results if r['correct'])
    accuracy = correct_count / len(results) if results else 0
    
    # Print summary
    print("📊 Evaluation Results")
    print("=" * 60)
    print(f"Total Problems: {len(results)}")
    print(f"Correct: {correct_count}")
    print(f"Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
    print(f"Total Duration: {total_time:.2f}s")
    print(f"Average per Problem: {total_time/len(results):.2f}s")
    print(f"Branch Count: {branches}")
    print("=" * 60)
    
    # Save results
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    summary = {
        "dataset": "aime_2024",
        "model_path": model_path,
        "config": config,
        "total_problems": len(results),
        "correct": correct_count,
        "accuracy": accuracy,
        "total_time": total_time,
        "average_time_per_problem": total_time / len(results),
        "branch_count": branches,
        "mode": mode,
        "results": results
    }
    
    with open(output_path, 'w') as f:
        json.dump(summary, f, indent=2)
    
    print(f"💾 Results saved to: {output_file}")
    
    return summary


def main():
    parser = argparse.ArgumentParser(description="Test AIME 2024 samples with adaptive CoT")
    parser.add_argument("--model_path", type=str, 
                       default="/raid/LLM/deepseek-r1-distill-qwen-14b",
                       help="Path to the model")
    parser.add_argument("--branches", type=int, default=8,
                       help="Number of reasoning branches")
    parser.add_argument("--samples", type=int, default=3,
                       help="Number of samples to test")
    parser.add_argument("--output", type=str, default="results/aime_2024_test.json",
                       help="Output file path")
    parser.add_argument("--mode", type=str, choices=["static", "adaptive"], default="static",
                       help="Branching mode")
    
    args = parser.parse_args()
    
    test_aime_2024_samples(
        model_path=args.model_path,
        branches=args.branches,
        samples=args.samples,
        output_file=args.output,
        mode=args.mode
    )


if __name__ == "__main__":
    main()