test_fixed_framework.py

#!/usr/bin/env python3
"""
Test the fixed framework with 10-sample GSM8K.
"""

import sys
import os
from pathlib import Path
import json
import time

# Add src to path
sys.path.append(str(Path(__file__).parent / "src"))

from models.model_factory import ModelFactory
from adaptive.adaptive_cot import AdaptiveCoT
from benchmarks.math_benchmarks import MathBenchmarkLoader


def test_fixed_framework():
    """Test the fixed framework with 10 GSM8K samples."""
    
    print("🔧 Testing Fixed Framework - 10 GSM8K Samples")
    print("=" * 60)
    
    try:
        # Load model
        print("🔧 Loading model...")
        model = ModelFactory.create_model(
            model_type="deepseek",
            model_name="/raid/LLM/llama3.1-8b-instruct",
            config={"gpu_id": 0}
        )
        model.load_model()
        
        # Load GSM8K dataset
        print("📚 Loading GSM8K dataset...")
        benchmark_loader = MathBenchmarkLoader(cache_dir="data_cache")
        gsm8k_data = benchmark_loader.load_dataset("gsm8k", max_samples=10)
        
        problems = []
        for item in gsm8k_data:
            problems.append({
                'question': item['question'],
                'answer': item['answer']
            })
        
        print(f"Loaded {len(problems)} problems")
        
        # Test configurations
        configs = [
            {
                "name": "No Few-Shot",
                "config": {
                    "adaptive_branching": True,
                    "min_branches": 3,
                    "max_branches": 8,
                    "default_branches": 5,
                    "num_fewshot": 0,
                    "temperature": 0.7,
                    "top_p": 0.95,
                    "max_tokens": 512,
                }
            },
            {
                "name": "2 Few-Shot",
                "config": {
                    "adaptive_branching": True,
                    "min_branches": 3,
                    "max_branches": 8,
                    "default_branches": 5,
                    "num_fewshot": 2,
                    "temperature": 0.7,
                    "top_p": 0.95,
                    "max_tokens": 512,
                }
            },
            {
                "name": "Static 5 Branches",
                "config": {
                    "adaptive_branching": False,
                    "min_branches": 5,
                    "max_branches": 5,
                    "default_branches": 5,
                    "num_fewshot": 0,
                    "temperature": 0.7,
                    "top_p": 0.95,
                    "max_tokens": 512,
                }
            }
        ]
        
        results = {}
        
        for config_info in configs:
            print(f"\n🧪 Testing {config_info['name']}")
            print("-" * 40)
            
            # Create Adaptive CoT instance
            adaptive_cot = AdaptiveCoT(model, config_info['config'])
            
            correct = 0
            total_branches = 0
            start_time = time.time()
            
            for i, problem in enumerate(problems):
                print(f"\n📝 Problem {i+1}/{len(problems)}: {problem['question'][:100]}...")
                
                try:
                    # Use Adaptive CoT
                    result = adaptive_cot.solve_problem(problem['question'])
                    
                    answer = result['final_answer']
                    branches_used = result.get('num_branches', 0)
                    total_branches += branches_used
                    
                    print(f"Answer: {answer}")
                    print(f"Branches Used: {branches_used}")
                    print(f"Ground Truth: {problem['answer']}")
                    
                    # Check accuracy
                    is_correct = check_accuracy(answer, problem['answer'])
                    print(f"Correct: {'✅' if is_correct else '❌'}")
                    
                    if is_correct:
                        correct += 1
                    
                except Exception as e:
                    print(f"❌ Error: {e}")
            
            end_time = time.time()
            duration = end_time - start_time
            accuracy = correct / len(problems)
            avg_branches = total_branches / len(problems)
            
            results[config_info['name']] = {
                'accuracy': accuracy,
                'correct': correct,
                'total': len(problems),
                'duration': duration,
                'avg_branches': avg_branches
            }
            
            print(f"\n📊 {config_info['name']} Results:")
            print(f"  Accuracy: {accuracy:.3f} ({correct}/{len(problems)})")
            print(f"  Duration: {duration:.2f}s")
            print(f"  Avg Branches: {avg_branches:.1f}")
        
        # Print comparison
        print("\n" + "=" * 60)
        print("📊 COMPARISON RESULTS")
        print("=" * 60)
        
        for name, result in results.items():
            print(f"{name}:")
            print(f"  Accuracy: {result['accuracy']:.3f} ({result['correct']}/{result['total']})")
            print(f"  Duration: {result['duration']:.2f}s")
            print(f"  Avg Branches: {result['avg_branches']:.1f}")
            print()
        
        # Save results
        with open('fixed_framework_results.json', 'w') as f:
            json.dump(results, f, indent=2, default=str)
        
        print(f"💾 Results saved to: fixed_framework_results.json")
        
    except Exception as e:
        print(f"❌ Error during testing: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


def check_accuracy(predicted, ground_truth):
    """Check if predicted answer matches ground truth."""
    if not predicted or not ground_truth:
        return False
    
    # Clean both answers
    pred_clean = clean_answer(predicted)
    gt_clean = clean_answer(ground_truth)
    
    return pred_clean == gt_clean


def clean_answer(answer):
    """Clean answer for comparison."""
    import re
    
    if not answer:
        return ""
    
    # Remove common prefixes
    answer = re.sub(r'^(The answer is|Answer:|Final answer:?)\s*', '', answer, flags=re.IGNORECASE)
    
    # Extract numbers
    numbers = re.findall(r'-?\d+\.?\d*', answer)
    if numbers:
        return numbers[-1]
    
    return answer.strip()


if __name__ == "__main__":
    test_fixed_framework()