simple_lighteval_comparison.py

#!/usr/bin/env python3
"""
Simple comparison between our framework and direct model generation (LightEval-style).
"""

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

# 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_our_framework(model, samples, num_fewshot=0):
    """Test our framework on the samples."""
    print(f"🔧 Testing Our Framework (few-shot={num_fewshot})")
    print("-" * 50)
    
    # Create Adaptive CoT configuration for single branch
    config = {
        "adaptive_branching": False,  # Disable adaptive branching for single branch
        "min_branches": 1,
        "max_branches": 1,
        "default_branches": 1,
        "num_fewshot": num_fewshot,
        "temperature": 0.7,
        "top_p": 0.95,
        "max_tokens": 512,
    }
    
    adaptive_cot = AdaptiveCoT(model, config)
    
    results = []
    correct = 0
    start_time = time.time()
    
    for i, sample in enumerate(samples):
        print(f"📝 Problem {i+1}/{len(samples)}: {sample['question'][:80]}...")
        
        try:
            result = adaptive_cot.solve_problem(sample['question'])
            
            answer = result['final_answer']
            reasoning_path = result.get('reasoning_paths', [''])[0] if result.get('reasoning_paths') else ''
            
            # Check accuracy
            is_correct = check_accuracy(answer, sample['answer'])
            if is_correct:
                correct += 1
            
            results.append({
                "problem_id": i + 1,
                "question": sample['question'],
                "ground_truth": sample['answer'],
                "our_answer": answer,
                "our_reasoning": reasoning_path,
                "correct": is_correct
            })
            
            print(f"  Our Answer: {answer}")
            print(f"  Ground Truth: {sample['answer']}")
            print(f"  Correct: {'✅' if is_correct else '❌'}")
            
        except Exception as e:
            print(f"  ❌ Error: {e}")
            results.append({
                "problem_id": i + 1,
                "question": sample['question'],
                "ground_truth": sample['answer'],
                "our_answer": "",
                "our_reasoning": "",
                "correct": False,
                "error": str(e)
            })
    
    end_time = time.time()
    duration = end_time - start_time
    accuracy = correct / len(samples)
    
    print(f"\n📊 Our Framework Results:")
    print(f"  Accuracy: {accuracy:.3f} ({correct}/{len(samples)})")
    print(f"  Duration: {duration:.2f}s")
    
    return {
        "results": results,
        "accuracy": accuracy,
        "correct": correct,
        "total": len(samples),
        "duration": duration
    }


def test_direct_generation(model, samples, num_fewshot=0):
    """Test direct model generation (LightEval-style) on the samples."""
    print(f"\n🔧 Testing Direct Generation (few-shot={num_fewshot})")
    print("-" * 50)
    
    results = []
    correct = 0
    start_time = time.time()
    
    for i, sample in enumerate(samples):
        print(f"📝 Problem {i+1}/{len(samples)}: {sample['question'][:80]}...")
        
        try:
            # Create prompt
            if num_fewshot > 0:
                from src.adaptive.fewshot_examples import FewShotExampleLoader
                fewshot_loader = FewShotExampleLoader()
                examples = fewshot_loader.get_fewshot_examples("gsm8k", num_fewshot)
                prompt = fewshot_loader.format_fewshot_prompt(examples, sample['question'])
            else:
                prompt = f"Q: {sample['question']}\nA:"
            
            # Generate using the model directly
            generated = model.generate(
                prompt,
                max_tokens=512,
                temperature=0.7,
                top_p=0.95,
                do_sample=True,
                num_return_sequences=1
            )
            
            if isinstance(generated, list):
                answer_text = generated[0]
            else:
                answer_text = generated
            
            # Apply stop sequences
            for stop_seq in ["Q:", "</s>", "<|im_end|>", "\n\nQ:"]:
                if stop_seq in answer_text:
                    answer_text = answer_text.split(stop_seq)[0]
            
            # Extract answer using improved method
            answer = extract_answer_improved(answer_text)
            
            # Check accuracy
            is_correct = check_accuracy(answer, sample['answer'])
            if is_correct:
                correct += 1
            
            results.append({
                "problem_id": i + 1,
                "question": sample['question'],
                "ground_truth": sample['answer'],
                "direct_answer": answer,
                "direct_reasoning": answer_text,
                "correct": is_correct
            })
            
            print(f"  Direct Answer: {answer}")
            print(f"  Ground Truth: {sample['answer']}")
            print(f"  Correct: {'✅' if is_correct else '❌'}")
            
        except Exception as e:
            print(f"  ❌ Error: {e}")
            results.append({
                "problem_id": i + 1,
                "question": sample['question'],
                "ground_truth": sample['answer'],
                "direct_answer": "",
                "direct_reasoning": "",
                "correct": False,
                "error": str(e)
            })
    
    end_time = time.time()
    duration = end_time - start_time
    accuracy = correct / len(samples)
    
    print(f"\n📊 Direct Generation Results:")
    print(f"  Accuracy: {accuracy:.3f} ({correct}/{len(samples)})")
    print(f"  Duration: {duration:.2f}s")
    
    return {
        "results": results,
        "accuracy": accuracy,
        "correct": correct,
        "total": len(samples),
        "duration": duration
    }


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)
    
    # Remove trailing punctuation (periods, commas, etc.)
    answer = re.sub(r'[.,;:!?]+$', '', answer)
    
    # Extract numbers
    numbers = re.findall(r'-?\d+\.?\d*', answer)
    if numbers:
        return numbers[-1]
    
    return answer.strip()


def extract_answer(text):
    """Extract answer from text using our framework's logic."""
    import re
    
    # Answer extraction patterns from our framework
    answer_patterns = [
        r"\\boxed\{([^}]+)\}",  # \boxed{answer}
        r"\$([^$]+)\$",         # $answer$
        r"Answer: ([^\n]+)",
        r"Final answer: ([^\n]+)",
        r"The final answer is:? ([^\n]+)",
        r"The answer is:? ([^\n]+)",
        r"= ([0-9]+(?:\.[0-9]+)?)",
    ]
    
    # First, try to find the answer using patterns
    for pattern in answer_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            answer = match.group(1).strip()
            cleaned = clean_answer(answer)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    # If no pattern match, try to find the first valid number
    lines = text.strip().split('\n')
    for line in lines:
        number_matches = re.findall(r'(\d+(?:\.\d+)?)', line)
        for number in number_matches:
            if is_valid_answer(number):
                return number
    
    # Last resort: return the last line if it contains a number
    if lines:
        last_line = lines[-1].strip()
        number_match = re.search(r'(\d+(?:\.\d+)?)', last_line)
        if number_match:
            return number_match.group(1)
    
    return ""


def extract_answer_improved(text):
    """Improved answer extraction that looks for the final answer more carefully."""
    import re
    
    # First, try to find explicit answer patterns
    answer_patterns = [
        r"\\boxed\{([^}]+)\}",  # \boxed{answer}
        r"\$([^$]+)\$",         # $answer$
        r"Answer: ([^\n]+)",
        r"Final answer: ([^\n]+)",
        r"The final answer is:? ([^\n]+)",
        r"The answer is:? ([^\n]+)",
        r"= ([0-9]+(?:\.[0-9]+)?)",
    ]
    
    for pattern in answer_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            answer = match.group(1).strip()
            cleaned = clean_answer(answer)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    # Look for patterns like "The answer is X" or "So the answer is X"
    final_answer_patterns = [
        r"(?:So )?the answer is:? ([0-9,]+)",
        r"(?:Therefore,? )?the answer is:? ([0-9,]+)",
        r"(?:Thus,? )?the answer is:? ([0-9,]+)",
        r"(?:Hence,? )?the answer is:? ([0-9,]+)",
        r"answer:? ([0-9,]+)",
    ]
    
    for pattern in final_answer_patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            answer = match.group(1).strip()
            cleaned = clean_answer(answer)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    # Look for the last number in the text that could be an answer
    # Split by sentences and look for numbers
    sentences = re.split(r'[.!?]\s*', text)
    for sentence in reversed(sentences):
        # Look for numbers in this sentence
        numbers = re.findall(r'([0-9,]+)', sentence)
        for number in reversed(numbers):
            cleaned = clean_answer(number)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    # Fallback: look for any valid number
    lines = text.strip().split('\n')
    for line in reversed(lines):
        number_matches = re.findall(r'([0-9,]+)', line)
        for number in reversed(number_matches):
            cleaned = clean_answer(number)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    return ""


def is_valid_answer(answer):
    """Check if an answer is valid."""
    if not answer or answer.strip() == "":
        return False
    
    try:
        num = float(answer)
        return -1000 <= num <= 10000
    except ValueError:
        return False


def main():
    """Main comparison function."""
    print("🔬 Our Framework vs Direct Generation Comparison")
    print("=" * 60)
    print("Testing on 30 GSM8K samples with single-branch generation")
    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=30)
        
        samples = []
        for item in gsm8k_data:
            samples.append({
                'question': item['question'],
                'answer': item['answer']
            })
        
        print(f"Loaded {len(samples)} samples")
        
        # Test configurations
        configs = [
            {"name": "No Few-Shot", "num_fewshot": 0},
            {"name": "2 Few-Shot", "num_fewshot": 2},
            {"name": "8 Few-Shot", "num_fewshot": 8}
        ]
        
        all_results = {}
        
        for config in configs:
            print(f"\n{'='*60}")
            print(f"🧪 Testing {config['name']}")
            print(f"{'='*60}")
            
            # Test our framework
            our_results = test_our_framework(model, samples, config['num_fewshot'])
            
            # Test direct generation
            direct_results = test_direct_generation(model, samples, config['num_fewshot'])
            
            # Store results
            all_results[config['name']] = {
                "our_framework": our_results,
                "direct_generation": direct_results
            }
            
            # Print comparison
            print(f"\n📊 {config['name']} Comparison:")
            print(f"  Our Framework:     {our_results['accuracy']:.3f} ({our_results['correct']}/{our_results['total']}) - {our_results['duration']:.2f}s")
            print(f"  Direct Generation: {direct_results['accuracy']:.3f} ({direct_results['correct']}/{direct_results['total']}) - {direct_results['duration']:.2f}s")
            
            # Calculate difference
            acc_diff = our_results['accuracy'] - direct_results['accuracy']
            print(f"  Difference:        {acc_diff:+.3f} ({'Our framework' if acc_diff > 0 else 'Direct generation'} {'wins' if abs(acc_diff) > 0.01 else 'tie'})")
        
        # Save detailed results
        output_file = "simple_comparison_results.json"
        with open(output_file, 'w') as f:
            json.dump(all_results, f, indent=2, default=str)
        
        print(f"\n💾 Detailed results saved to: {output_file}")
        
        # Print final summary
        print(f"\n{'='*60}")
        print("📊 FINAL SUMMARY")
        print(f"{'='*60}")
        
        for config_name, results in all_results.items():
            print(f"\n{config_name}:")
            our_acc = results['our_framework']['accuracy']
            our_time = results['our_framework']['duration']
            direct_acc = results['direct_generation']['accuracy']
            direct_time = results['direct_generation']['duration']
            print(f"  Our Framework:     {our_acc:.3f} accuracy, {our_time:.2f}s")
            print(f"  Direct Generation: {direct_acc:.3f} accuracy, {direct_time:.2f}s")
            print(f"  Difference:        {our_acc - direct_acc:+.3f} accuracy")
        
    except Exception as e:
        print(f"❌ Error during comparison: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()