test_framework_comparison.py

#!/usr/bin/env python3
"""
Direct comparison test: Our framework vs simple generation.
This will help us understand if the issue is with our framework or the model itself.
"""

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_simple_generation(model, problems, num_samples=10):
    """Test simple generation without adaptive CoT."""
    print("🧪 Testing Simple Generation (No Adaptive CoT)")
    print("-" * 50)
    
    results = []
    
    for i, problem in enumerate(problems[:num_samples]):
        print(f"\n📝 Problem {i+1}/{num_samples}:")
        print(f"Q: {problem['question']}")
        
        # Simple generation
        prompt = f"Q: {problem['question']}\nA:"
        
        try:
            # Generate single response
            response = model.generate(
                prompt,
                max_tokens=512,
                temperature=0.7,
                top_p=0.95,
                do_sample=True
            )
            
            # Extract answer
            answer = extract_simple_answer(response)
            
            print(f"Generated: {response[:200]}...")
            print(f"Extracted Answer: {answer}")
            print(f"Ground Truth: {problem['answer']}")
            
            # Check accuracy
            is_correct = check_accuracy(answer, problem['answer'])
            print(f"Correct: {'✅' if is_correct else '❌'}")
            
            results.append({
                'problem': problem['question'],
                'ground_truth': problem['answer'],
                'generated': response,
                'extracted_answer': answer,
                'correct': is_correct
            })
            
        except Exception as e:
            print(f"❌ Error: {e}")
            results.append({
                'problem': problem['question'],
                'ground_truth': problem['answer'],
                'generated': '',
                'extracted_answer': '',
                'correct': False,
                'error': str(e)
            })
    
    return results


def test_adaptive_cot(model, problems, num_samples=10):
    """Test with our Adaptive CoT framework."""
    print("\n🧪 Testing Adaptive CoT Framework")
    print("-" * 50)
    
    # Create Adaptive CoT configuration
    adaptive_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,
    }
    
    # Create Adaptive CoT instance
    adaptive_cot = AdaptiveCoT(model, adaptive_config)
    
    results = []
    
    for i, problem in enumerate(problems[:num_samples]):
        print(f"\n📝 Problem {i+1}/{num_samples}:")
        print(f"Q: {problem['question']}")
        
        try:
            # Use Adaptive CoT
            result = adaptive_cot.solve_problem(problem['question'])
            
            answer = result['final_answer']
            branches_used = result.get('num_branches', 0)
            reasoning_paths = result.get('reasoning_paths', [])
            
            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 '❌'}")
            
            # Show first reasoning path
            if reasoning_paths:
                print(f"First Reasoning Path: {reasoning_paths[0][:200]}...")
            
            results.append({
                'problem': problem['question'],
                'ground_truth': problem['answer'],
                'answer': answer,
                'branches_used': branches_used,
                'reasoning_paths': reasoning_paths,
                'correct': is_correct
            })
            
        except Exception as e:
            print(f"❌ Error: {e}")
            results.append({
                'problem': problem['question'],
                'ground_truth': problem['answer'],
                'answer': '',
                'branches_used': 0,
                'reasoning_paths': [],
                'correct': False,
                'error': str(e)
            })
    
    return results


def extract_simple_answer(text):
    """Extract answer from simple generation."""
    import re
    
    # Look for patterns
    patterns = [
        r"\\boxed\{([^}]+)\}",
        r"\$([^$]+)\$",
        r"The final answer is:? ([^\n]+)",
        r"Final answer:? ([^\n]+)",
        r"Answer:? ([^\n]+)",
    ]
    
    for pattern in patterns:
        match = re.search(pattern, text)
        if match:
            answer = match.group(1).strip()
            if is_valid_answer(answer):
                return answer
    
    # Fallback: look for numbers in the last few lines
    lines = text.split('\n')[-5:]
    for line in lines:
        numbers = re.findall(r'-?\d+\.?\d*', line)
        if numbers:
            return numbers[-1]
    
    return ""


def is_valid_answer(answer):
    """Check if answer is valid."""
    if not answer or len(answer) > 50:
        return False
    
    import re
    return bool(re.search(r'\d', answer))


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
    
    # 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()


def main():
    print("🔬 Framework Comparison Test - GSM8K (10 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 1: Simple generation
        print("\n" + "="*60)
        start_time = time.time()
        simple_results = test_simple_generation(model, problems, 10)
        simple_duration = time.time() - start_time
        
        # Test 2: Adaptive CoT
        print("\n" + "="*60)
        start_time = time.time()
        adaptive_results = test_adaptive_cot(model, problems, 10)
        adaptive_duration = time.time() - start_time
        
        # Calculate metrics
        simple_correct = sum(1 for r in simple_results if r['correct'])
        adaptive_correct = sum(1 for r in adaptive_results if r['correct'])
        
        simple_accuracy = simple_correct / len(simple_results)
        adaptive_accuracy = adaptive_correct / len(adaptive_results)
        
        # Calculate average branches
        avg_branches = sum(r.get('branches_used', 0) for r in adaptive_results) / len(adaptive_results)
        
        # Print comparison
        print("\n" + "="*60)
        print("📊 COMPARISON RESULTS")
        print("="*60)
        print(f"Simple Generation:")
        print(f"  Accuracy: {simple_accuracy:.3f} ({simple_correct}/{len(simple_results)})")
        print(f"  Duration: {simple_duration:.2f}s")
        print(f"  Average Branches: 1")
        print()
        print(f"Adaptive CoT:")
        print(f"  Accuracy: {adaptive_accuracy:.3f} ({adaptive_correct}/{len(adaptive_results)})")
        print(f"  Duration: {adaptive_duration:.2f}s")
        print(f"  Average Branches: {avg_branches:.1f}")
        print()
        print(f"Improvement: {((adaptive_accuracy - simple_accuracy) / simple_accuracy * 100):+.1f}%")
        
        # Save detailed results
        comparison_results = {
            'simple_generation': {
                'results': simple_results,
                'accuracy': simple_accuracy,
                'duration': simple_duration,
                'avg_branches': 1
            },
            'adaptive_cot': {
                'results': adaptive_results,
                'accuracy': adaptive_accuracy,
                'duration': adaptive_duration,
                'avg_branches': avg_branches
            },
            'comparison': {
                'accuracy_improvement': ((adaptive_accuracy - simple_accuracy) / simple_accuracy * 100) if simple_accuracy > 0 else 0,
                'efficiency_ratio': adaptive_duration / simple_duration if simple_duration > 0 else 0
            }
        }
        
        with open('framework_comparison_results.json', 'w') as f:
            json.dump(comparison_results, f, indent=2, default=str)
        
        print(f"\n💾 Detailed results saved to: framework_comparison_results.json")
        
    except Exception as e:
        print(f"❌ Error during comparison: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()