research_experiment.py

#!/usr/bin/env python3
"""
Research Experiment Runner for Adaptive CoT Framework

Focused on parallel test-time scaling with self-consistency:
- Adaptive branch allocation based on prefill signals
- Default 8 branches for static baseline
- Reliable accuracy measurement
- Comprehensive research logging
"""

import argparse
import sys
import os
import yaml
import time
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.models.model_factory import ModelFactory
from src.adaptive.adaptive_cot import AdaptiveCoT
from src.evaluation.evaluator import AdaptiveCoTEvaluator
from src.utils.research_logger import ResearchLogger
from src.utils.memory_monitor import MemoryMonitor


class ResearchExperiment:
    """Research experiment runner focused on adaptive branching."""
    
    def __init__(self, config_path: str):
        """Initialize with research configuration."""
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)
        
        # Setup logging
        self.research_logger = ResearchLogger(
            output_dir=self.config["research_logging"]["output_dir"]
        )
        
        # Setup memory monitoring
        self.memory_monitor = MemoryMonitor(verbose=True)
        
        # Create experiment directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.experiment_dir = Path(f"research_experiments/adaptive_cot_{timestamp}")
        self.experiment_dir.mkdir(parents=True, exist_ok=True)
        
        # Save config
        with open(self.experiment_dir / "config.yaml", 'w') as f:
            yaml.dump(self.config, f, default_flow_style=False)
        
        print(f"🔬 Research Experiment Initialized")
        print(f"📁 Experiment directory: {self.experiment_dir}")
        print(f"📊 Logging level: {self.config['research_logging']['log_level']}")
    
    def load_model(self, model_path: str, model_type: str = "deepseek") -> Any:
        """Load the specified model."""
        print(f"📦 Loading model: {model_path}")
        
        model_config = self.config["models"]["deepseek_r1_distill_qwen"].copy()
        model_config["model_name"] = model_path
        
        model = ModelFactory.create_model(model_type, model_path, model_config)
        model.load_model()
        
        print(f"✅ Model loaded successfully")
        return model
    
    def test_single_problem(self, model: Any, problem: str, test_adaptive: bool = True, test_static: bool = True) -> Dict[str, Any]:
        """Test a single problem with both adaptive and static approaches."""
        results = {}
        
        # Test adaptive branching
        if test_adaptive:
            print(f"\n🔍 Testing Adaptive Branching:")
            print("-" * 40)
            
            adaptive_config = self.config["adaptive_branching"].copy()
            adaptive_config["enabled"] = True
            
            cot_adaptive = AdaptiveCoT(model, adaptive_config)
            start_time = time.time()
            
            try:
                adaptive_result = cot_adaptive.solve_problem(problem)
                adaptive_result["execution_time"] = time.time() - start_time
                results["adaptive"] = adaptive_result
                
                print(f"✅ Adaptive: {adaptive_result['final_answer']} ({adaptive_result['num_branches']} branches)")
                print(f"   Prefill signals: entropy={adaptive_result['prefill_signals']['entropy']:.3f}")
                print(f"   Execution time: {adaptive_result['execution_time']:.2f}s")
                
            except Exception as e:
                print(f"❌ Adaptive failed: {e}")
                results["adaptive"] = {"error": str(e)}
        
        # Test static branching (8 branches default)
        if test_static:
            print(f"\n🔍 Testing Static Branching (8 branches):")
            print("-" * 40)
            
            static_config = self.config["adaptive_branching"].copy()
            static_config["enabled"] = False
            static_config["default_branches"] = 8
            
            cot_static = AdaptiveCoT(model, static_config)
            start_time = time.time()
            
            try:
                static_result = cot_static.solve_problem(problem)
                static_result["execution_time"] = time.time() - start_time
                results["static"] = static_result
                
                print(f"✅ Static: {static_result['final_answer']} ({static_result['num_branches']} branches)")
                print(f"   Execution time: {static_result['execution_time']:.2f}s")
                
            except Exception as e:
                print(f"❌ Static failed: {e}")
                results["static"] = {"error": str(e)}
        
        return results
    
    def run_benchmark_evaluation(self, model: Any, dataset: str, max_samples: int = 100) -> Dict[str, Any]:
        """Run evaluation on a benchmark dataset."""
        print(f"\n📊 Running Benchmark Evaluation:")
        print(f"   Dataset: {dataset}")
        print(f"   Max samples: {max_samples}")
        print("-" * 50)
        
        # Create evaluator
        evaluator = AdaptiveCoTEvaluator(model, self.config)
        
        # Run evaluation
        results = evaluator.evaluate_dataset(
            dataset_name=dataset,
            max_samples=max_samples,
            save_results=True
        )
        
        # Display results
        print(f"\n📈 Evaluation Results:")
        print(f"   Problems: {results['num_problems']}")
        print(f"   Successful: {results['num_successful']}")
        print(f"   Accuracy: {results['metrics']['accuracy']['exact_match']:.3f}")
        print(f"   Avg branches: {results['metrics']['efficiency']['average_branch_count']:.1f}")
        print(f"   Avg time: {results['metrics']['efficiency']['average_execution_time']:.2f}s")
        
        return results
    
    def compare_strategies(self, model: Any, dataset: str, max_samples: int = 100) -> Dict[str, Any]:
        """Compare adaptive vs static strategies."""
        print(f"\n🔄 Comparing Strategies:")
        print(f"   Dataset: {dataset}")
        print(f"   Max samples: {max_samples}")
        print("-" * 50)
        
        # Create evaluator
        evaluator = AdaptiveCoTEvaluator(model, self.config)
        
        # Compare strategies
        strategies = ["adaptive", "static"]
        results = evaluator.compare_strategies(
            dataset_name=dataset,
            strategies=strategies,
            max_samples=max_samples
        )
        
        # Display comparison
        print(f"\n📊 Strategy Comparison Results:")
        for strategy, result in results.items():
            if "error" not in result:
                accuracy = result.get("metrics", {}).get("accuracy", {}).get("exact_match", 0.0)
                avg_branches = result.get("metrics", {}).get("efficiency", {}).get("average_branch_count", 0.0)
                print(f"   {strategy}: {accuracy:.3f} accuracy, {avg_branches:.1f} avg branches")
            else:
                print(f"   {strategy}: ERROR - {result['error']}")
        
        return results


def main():
    """Main research experiment function."""
    parser = argparse.ArgumentParser(
        description="Research Experiment Runner for Adaptive CoT Framework",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Test single problem
  python research_experiment.py --model-path "/path/to/model" --problem "What is 2+2?" --test-both
  
  # Run benchmark evaluation
  python research_experiment.py --model-path "/path/to/model" --benchmark gsm8k --max-samples 100
  
  # Compare strategies
  python research_experiment.py --model-path "/path/to/model" --compare gsm8k --max-samples 100
        """
    )
    
    # Model configuration
    parser.add_argument("--model-path", type=str, required=True, help="Path to model")
    parser.add_argument("--model-type", type=str, default="deepseek", help="Model type")
    parser.add_argument("--config", type=str, default="research_config.yaml", help="Config file")
    
    # Experiment types
    parser.add_argument("--problem", type=str, help="Single problem to test")
    parser.add_argument("--test-adaptive", action="store_true", help="Test adaptive branching")
    parser.add_argument("--test-static", action="store_true", help="Test static branching")
    parser.add_argument("--test-both", action="store_true", help="Test both adaptive and static")
    
    parser.add_argument("--benchmark", type=str, choices=["gsm8k", "aime", "olympiad", "math"], 
                       help="Benchmark dataset to evaluate")
    parser.add_argument("--compare", type=str, choices=["gsm8k", "aime", "olympiad", "math"], 
                       help="Dataset to compare strategies on")
    
    # Experiment parameters
    parser.add_argument("--max-samples", type=int, default=100, help="Maximum samples")
    parser.add_argument("--gpu-id", type=int, default=1, help="GPU ID to use")
    
    args = parser.parse_args()
    
    # Set GPU
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
    
    try:
        # Initialize experiment
        experiment = ResearchExperiment(args.config)
        
        # Load model
        model = experiment.load_model(args.model_path, args.model_type)
        
        # Run experiments
        if args.problem:
            # Single problem test
            test_adaptive = args.test_adaptive or args.test_both
            test_static = args.test_static or args.test_both
            
            if not test_adaptive and not test_static:
                test_adaptive = test_static = True  # Default to both
            
            results = experiment.test_single_problem(
                model, args.problem, test_adaptive, test_static
            )
            
            # Save results
            with open(experiment.experiment_dir / "single_problem_results.json", 'w') as f:
                import json
                json.dump(results, f, indent=2, default=str)
        
        elif args.benchmark:
            # Benchmark evaluation
            results = experiment.run_benchmark_evaluation(
                model, args.benchmark, args.max_samples
            )
            
            # Save results
            with open(experiment.experiment_dir / f"{args.benchmark}_results.json", 'w') as f:
                import json
                json.dump(results, f, indent=2, default=str)
        
        elif args.compare:
            # Strategy comparison
            results = experiment.compare_strategies(
                model, args.compare, args.max_samples
            )
            
            # Save results
            with open(experiment.experiment_dir / f"{args.compare}_comparison.json", 'w') as f:
                import json
                json.dump(results, f, indent=2, default=str)
        
        else:
            print("❌ Please specify --problem, --benchmark, or --compare")
            return 1
        
        print(f"\n✅ Experiment completed successfully!")
        print(f"📁 Results saved to: {experiment.experiment_dir}")
        
        return 0
        
    except Exception as e:
        print(f"❌ Experiment failed: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())