test_10_samples.py

#!/usr/bin/env python3
"""
Test 10 samples with no few-shot and 8-shot to analyze answer extraction issues.
"""

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_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.0,  # Use deterministic generation for single branch
        "top_p": 1.0,        # Use deterministic generation for single branch
        "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 '❌'}")
            
            # Show reasoning snippet for debugging
            if not is_correct:
                print(f"  Reasoning snippet: {reasoning_path[-200:]}...")
            
        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 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:"
            
            # Set seed for deterministic generation
            import torch
            import numpy as np
            torch.manual_seed(42)
            np.random.seed(42)
            
            # Generate using the model directly
            generated = model.generate(
                prompt,
                max_tokens=512,
                temperature=0.0,  # Use deterministic generation for single branch
                top_p=1.0,        # Use deterministic generation for single branch
                do_sample=False,  # Use deterministic generation for single branch
                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]
            
            # Strip whitespace to match our framework's behavior
            answer_text = answer_text.strip()
            
            # Extract answer using the same method as our framework
            answer = extract_answer_like_framework(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 '❌'}")
            
            # Show reasoning snippet for debugging
            if not is_correct:
                print(f"  Reasoning snippet: {answer_text[-200:]}...")
            
        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 predicted answer
    pred_clean = clean_answer(predicted)
    
    # Extract final answer from ground truth (look for #### pattern)
    gt_clean = extract_final_answer_from_ground_truth(ground_truth)
    
    return pred_clean == gt_clean

def extract_final_answer_from_ground_truth(ground_truth):
    """Extract the final answer from ground truth format."""
    import re
    
    # Look for #### pattern at the end
    match = re.search(r'####\s*([^\n]+)', ground_truth)
    if match:
        return clean_answer(match.group(1))
    
    # Fallback: look for the last number in the text
    numbers = re.findall(r'([0-9,]+(?:\.[0-9]+)?)', ground_truth)
    if numbers:
        return clean_answer(numbers[-1])
    
    return clean_answer(ground_truth)


def clean_answer(answer):
    """Clean answer for comparison using the same method as our framework."""
    import re
    
    if not answer:
        return ""
    
    answer = str(answer).strip()
    
    # Remove common prefixes
    answer = re.sub(r'^(The answer is|Answer:|Final answer:?)\s*', '', answer, flags=re.IGNORECASE)
    
    # Remove dollar signs and other currency symbols
    answer = re.sub(r'[\$\s]+', '', answer)
    
    # Remove boxed formatting (handle nested cases like $\boxed{$70,000}$), but keep the number
    answer = re.sub(r'\\boxed\{([^}]+)\}', r'\1', answer)
    
    # Remove brackets, parentheses
    answer = re.sub(r'^[\\[\\](){}]+|[\\[\\](){}]+$', '', answer)
    
    # Remove trailing punctuation (periods, commas, etc.)
    answer = re.sub(r'[.,;:!?]+$', '', answer)
    
    # Remove commas from numbers (e.g., "70,000" -> "70000")
    answer = answer.replace(",", "")
    
    # Convert to float and back to remove unnecessary decimals (e.g., "18.0" -> "18")
    try:
        num = float(answer)
        if num == int(num):
            return str(int(num))
        else:
            return str(num)
    except ValueError:
        return answer


def extract_answer_like_framework(text):
    """Extract answer using the same method as our framework."""
    import re
    
    if not text or not text.strip():
        return ""
    
    # Use the same patterns as our framework
    answer_patterns = [
        re.compile(r"####\s*([-+]?\d+(?:\.\d+)?)", re.I),  # #### answer
        re.compile(r"final answer.*?([-+]?\d+(?:\.\d+)?)", re.I),  # final answer
        re.compile(r"answer is\s*[:\s]?([-+]?\d[\d,]*(?:\.\d+)?)(?=[\.\n]|$)", re.I),  # answer is
    ]
    number_pattern = re.compile(r"[-+]?\d[\d,]*(?:\.\d+)?")
    
    # Try to find explicit answer patterns first
    for pattern in answer_patterns:
        match = pattern.search(text)
        if match:
            extracted_answer = match.group(1).strip()
            cleaned_answer = clean_answer(extracted_answer)
            if is_valid_answer(cleaned_answer):
                return cleaned_answer
    
    # Fallback: find the last number in the text
    all_numbers = number_pattern.findall(text)
    if all_numbers:
        last_number = clean_answer(all_numbers[-1])
        if is_valid_answer(last_number):
            return last_number
    
    return ""

def is_valid_answer(answer: str) -> bool:
    """Check if an answer is valid (not empty, reasonable number)."""
    if not answer or answer.strip() == "":
        return False
    
    try:
        num = float(answer)
        # Check if it's a reasonable number (not too large or too small)
        return -100000 <= num <= 1000000
    except ValueError:
        return False


def extract_answer_improved(text):
    """Improved answer extraction that looks for the final answer more carefully."""
    import re
    
    if not text or not text.strip():
        return ""
    
    # Split into lines for analysis
    lines = text.strip().split('\n')
    
    # Strategy 1: Look for explicit answer patterns in the last few lines
    # (Most reliable - look for "final answer", "answer is", etc.)
    explicit_patterns = [
        r"(?:The )?final answer is:?\s*([^\n]+)",
        r"(?:The )?answer is:?\s*([^\n]+)",
        r"Answer:?\s*([^\n]+)",
        r"Therefore,?\s*(?:the )?answer is:?\s*([^\n]+)",
        r"Thus,?\s*(?:the )?answer is:?\s*([^\n]+)",
        r"Hence,?\s*(?:the )?answer is:?\s*([^\n]+)",
        r"So,?\s*(?:the )?answer is:?\s*([^\n]+)",
    ]
    
    # Check the last 5 lines for explicit answer patterns
    for line in lines[-5:]:
        for pattern in explicit_patterns:
            match = re.search(pattern, line, re.IGNORECASE)
            if match:
                answer = match.group(1).strip()
                cleaned = clean_answer(answer)
                if cleaned and is_valid_answer(cleaned):
                    return cleaned
    
    # Strategy 2: Look for boxed answers
    boxed_patterns = [
        r"\\boxed\{([^}]+)\}",
        r"\$\$([^$]+)\$\$",
        r"\$([^$]+)\$",
    ]
    
    for line in lines:
        for pattern in boxed_patterns:
            match = re.search(pattern, line)
            if match:
                answer = match.group(1).strip()
                cleaned = clean_answer(answer)
                if cleaned and is_valid_answer(cleaned):
                    return cleaned
    
    # Strategy 3: Look for "= number" patterns in the last few lines
    # (Often used in final calculations) - IMPROVED to handle "100 + 20 + 40 = 160"
    for line in lines[-3:]:
        # Look for patterns like "= 70" or "= $70" or "= 70,000" or "100 + 20 + 40 = 160"
        # First try to find the result of an equation
        equation_match = re.search(r'([0-9,]+(?:\.[0-9]+)?)\s*\+\s*[0-9,]+(?:\.[0-9]+)?\s*\+\s*[0-9,]+(?:\.[0-9]+)?\s*=\s*([0-9,]+(?:\.[0-9]+)?)', line)
        if equation_match:
            answer = equation_match.group(2).strip()  # Take the result, not the first number
            cleaned = clean_answer(answer)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
        
        # Then try simple equals patterns
        equals_match = re.search(r'=\s*\$?([0-9,]+(?:\.[0-9]+)?)', line)
        if equals_match:
            answer = equals_match.group(1).strip()
            cleaned = clean_answer(answer)
            if cleaned and is_valid_answer(cleaned):
                return cleaned
    
    # Strategy 4: Look for the last reasonable number in the text
    # (This is more careful than the original approach)
    all_numbers = []
    for line in lines:
        # Find all numbers in this line
        numbers = re.findall(r'([0-9,]+(?:\.[0-9]+)?)', line)
        for num in numbers:
            cleaned = clean_answer(num)
            if cleaned and is_valid_answer(cleaned):
                all_numbers.append((cleaned, line))
    
    # Return the last valid number found
    if all_numbers:
        return all_numbers[-1][0]
    
    return ""


def is_valid_answer(answer):
    """Check if an answer is valid."""
    if not answer or answer.strip() == "":
        return False
    
    try:
        num = float(answer)
        # Updated range to handle larger numbers like 70,000
        return -100000 <= num <= 1000000
    except ValueError:
        return False


def main():
    """Main test function."""
    print("🔬 10-Sample Test: No Few-Shot vs 8-Shot")
    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)
        
        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": "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 = "test_10_samples_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 test: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()