test_5_samples.py

#!/usr/bin/env python3
"""
Test 5 samples with zero-shot to verify the fix.
"""

import sys
import os
from pathlib import Path
import json
import time
import torch
import numpy as np

# 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):
    """Test our framework on the samples."""
    print(f"🔧 Testing Our Framework (zero-shot)")
    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": 0,
        "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 '❌'}")
            
        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):
    """Test direct model generation on the samples."""
    print(f"\n🔧 Testing Direct Generation (zero-shot)")
    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
            prompt = f"Q: {sample['question']}\nA:"
            
            # Set seed for deterministic generation
            torch.manual_seed(42)
            np.random.seed(42)
            # Also set CUDA seed if available
            if torch.cuda.is_available():
                torch.cuda.manual_seed(42)
                torch.cuda.manual_seed_all(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 '❌'}")
            
        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 main():
    """Main test function."""
    print("🔬 5-Sample Test: Zero-Shot (Fixed)")
    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=5)
        
        samples = []
        for item in gsm8k_data:
            samples.append({
                'question': item['question'],
                'answer': item['answer']
            })
        
        print(f"Loaded {len(samples)} samples")
        
        # Test our framework
        our_results = test_our_framework(model, samples)
        
        # Test direct generation
        direct_results = test_direct_generation(model, samples)
        
        # Print comparison
        print(f"\n📊 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'})")
        
        # Check if reasoning is identical
        print(f"\n🔍 Reasoning Comparison:")
        identical_count = 0
        for i in range(len(samples)):
            our_reasoning = our_results['results'][i]['our_reasoning']
            direct_reasoning = direct_results['results'][i]['direct_reasoning']
            identical = our_reasoning == direct_reasoning
            if identical:
                identical_count += 1
            print(f"  Problem {i+1}: {'✅ Identical' if identical else '❌ Different'}")
        
        print(f"\n📊 Reasoning Identical: {identical_count}/{len(samples)} ({identical_count/len(samples):.1%})")
        
        # Save detailed results
        output_file = "test_5_samples_results.json"
        all_results = {
            "our_framework": our_results,
            "direct_generation": direct_results
        }
        
        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}")
        
    except Exception as e:
        print(f"❌ Error during test: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()