research_framework.py

#!/usr/bin/env python3
"""
Research Framework for Adaptive CoT with Parallel Test-Time Scaling

This framework implements your specific requirements:
- Parallel test-time scaling with self-consistency + CoT
- Adaptive branch allocation based on prefill signals
- Default 8 branches for static baseline
- Reliable accuracy measurement
- Comprehensive research logging
- Support for reasoning models (DeepSeek-R1-Distill-Qwen, etc.)
- Math benchmarks (GSM8K, AIME, Olympiad, MATH)
"""

import argparse
import sys
import os
import yaml
import time
import json
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import pandas as pd

# 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 ResearchFramework:
    """Research framework for adaptive parallel test-time scaling."""
    
    def __init__(self, config_path: str = "research_config.yaml"):
        """Initialize the research framework."""
        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 Framework Initialized")
        print(f"📁 Experiment directory: {self.experiment_dir}")
        print(f"📊 Default static branches: {self.config['adaptive_branching']['default_branches']}")
    
    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_adaptive_vs_static(self, model: Any, problem: str) -> Dict[str, Any]:
        """Compare adaptive vs static branching on a single problem."""
        print(f"\n🔄 Comparing Adaptive vs Static Branching")
        print(f"Problem: {problem}")
        print("=" * 60)
        
        results = {}
        
        # Test Adaptive Branching
        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 Result:")
            print(f"   Answer: {adaptive_result['final_answer']}")
            print(f"   Branches: {adaptive_result['num_branches']}")
            print(f"   Strategy: {adaptive_result['allocation_info']['strategy']}")
            print(f"   Time: {adaptive_result['execution_time']:.2f}s")
            print(f"   Consensus: {adaptive_result['consensus_info']['confidence']:.3f}")
            
            if 'prefill_signals' in adaptive_result:
                prefill = adaptive_result['prefill_signals']
                print(f"   Prefill: entropy={prefill['entropy']:.3f}, confidence={prefill['confidence']:.3f}")
            
        except Exception as e:
            print(f"❌ Adaptive failed: {e}")
            results["adaptive"] = {"error": str(e)}
        
        # Test Static Branching (8 branches default)
        print(f"\n🔍 Testing Static Branching (8 branches):")
        print("-" * 40)
        
        static_config = {
            "adaptive_branching": False,
            "min_branches": 1,
            "max_branches": 10,
            "default_branches": 8,
            "research_logging": False
        }
        
        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 Result:")
            print(f"   Answer: {static_result['final_answer']}")
            print(f"   Branches: {static_result['num_branches']}")
            print(f"   Strategy: {static_result['allocation_info']['strategy']}")
            print(f"   Time: {static_result['execution_time']:.2f}s")
            print(f"   Consensus: {static_result['consensus_info']['confidence']:.3f}")
            
        except Exception as e:
            print(f"❌ Static failed: {e}")
            results["static"] = {"error": str(e)}
        
        # Compare results
        print(f"\n📊 Comparison Summary:")
        print("-" * 40)
        
        if "error" not in results.get("adaptive", {}) and "error" not in results.get("static", {}):
            adaptive = results["adaptive"]
            static = results["static"]
            
            print(f"   Answer Match: {adaptive['final_answer'] == static['final_answer']}")
            print(f"   Branch Efficiency: {static['num_branches'] - adaptive['num_branches']} fewer branches with adaptive")
            print(f"   Time Difference: {abs(adaptive['execution_time'] - static['execution_time']):.2f}s")
            print(f"   Consensus Quality: Adaptive={adaptive['consensus_info']['confidence']:.3f}, Static={static['consensus_info']['confidence']:.3f}")
            
            # Calculate efficiency gain
            if adaptive['num_branches'] < static['num_branches']:
                efficiency_gain = (static['num_branches'] - adaptive['num_branches']) / static['num_branches'] * 100
                print(f"   🎯 Efficiency Gain: {efficiency_gain:.1f}% fewer branches with adaptive")
        
        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("=" * 60)
        
        # 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 run_strategy_comparison(self, model: Any, dataset: str, max_samples: int = 100) -> Dict[str, Any]:
        """Compare adaptive vs static strategies on a dataset."""
        print(f"\n🔄 Running Strategy Comparison")
        print(f"Dataset: {dataset}")
        print(f"Max samples: {max_samples}")
        print("=" * 60)
        
        # 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:")
        print("-" * 50)
        
        comparison_data = []
        for strategy, result in results.items():
            if "error" not in result:
                metrics = result.get("metrics", {})
                accuracy = metrics.get("accuracy", {}).get("exact_match", 0.0)
                avg_branches = metrics.get("efficiency", {}).get("average_branch_count", 0.0)
                avg_time = metrics.get("efficiency", {}).get("average_execution_time", 0.0)
                
                print(f"   {strategy.upper()}:")
                print(f"     Accuracy: {accuracy:.3f}")
                print(f"     Avg Branches: {avg_branches:.1f}")
                print(f"     Avg Time: {avg_time:.2f}s")
                
                comparison_data.append({
                    "strategy": strategy,
                    "accuracy": accuracy,
                    "avg_branches": avg_branches,
                    "avg_time": avg_time
                })
            else:
                print(f"   {strategy.upper()}: ERROR - {result['error']}")
        
        # Calculate efficiency gains
        if len(comparison_data) == 2:
            adaptive = next((d for d in comparison_data if d["strategy"] == "adaptive"), None)
            static = next((d for d in comparison_data if d["strategy"] == "static"), None)
            
            if adaptive and static:
                branch_efficiency = (static["avg_branches"] - adaptive["avg_branches"]) / static["avg_branches"] * 100
                time_efficiency = (static["avg_time"] - adaptive["avg_time"]) / static["avg_time"] * 100
                
                print(f"\n🎯 Efficiency Analysis:")
                print(f"   Branch Efficiency: {branch_efficiency:.1f}% fewer branches with adaptive")
                print(f"   Time Efficiency: {time_efficiency:.1f}% time difference")
                print(f"   Accuracy Difference: {abs(adaptive['accuracy'] - static['accuracy']):.3f}")
                
                # Determine if adaptive is more efficient
                if branch_efficiency > 0 and abs(adaptive['accuracy'] - static['accuracy']) < 0.05:
                    print(f"   ✅ Adaptive is more efficient while maintaining accuracy!")
                elif branch_efficiency > 0:
                    print(f"   ⚠️  Adaptive uses fewer branches but accuracy differs")
                else:
                    print(f"   ❌ Adaptive uses more branches than static")
        
        return results
    
    def save_results(self, results: Dict[str, Any], filename: str) -> None:
        """Save results to JSON file."""
        filepath = self.experiment_dir / filename
        with open(filepath, 'w') as f:
            json.dump(results, f, indent=2, default=str)
        print(f"💾 Results saved to: {filepath}")


def main():
    """Main research framework function."""
    parser = argparse.ArgumentParser(
        description="Research Framework for Adaptive CoT with Parallel Test-Time Scaling",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Test single problem comparison
  python research_framework.py --model-path "/path/to/model" --problem "What is 2+2?" --compare-single
  
  # Run benchmark evaluation
  python research_framework.py --model-path "/path/to/model" --benchmark gsm8k --max-samples 100
  
  # Compare strategies on dataset
  python research_framework.py --model-path "/path/to/model" --compare-dataset 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("--compare-single", action="store_true", help="Compare adaptive vs static on single problem")
    
    parser.add_argument("--benchmark", type=str, choices=["gsm8k", "aime", "olympiad", "math"], 
                       help="Benchmark dataset to evaluate")
    parser.add_argument("--compare-dataset", 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 framework
        framework = ResearchFramework(args.config)
        
        # Load model
        model = framework.load_model(args.model_path, args.model_type)
        
        # Run experiments
        if args.problem and args.compare_single:
            # Single problem comparison
            results = framework.test_adaptive_vs_static(model, args.problem)
            framework.save_results(results, "single_problem_comparison.json")
        
        elif args.benchmark:
            # Benchmark evaluation
            results = framework.run_benchmark_evaluation(model, args.benchmark, args.max_samples)
            framework.save_results(results, f"{args.benchmark}_evaluation.json")
        
        elif args.compare_dataset:
            # Strategy comparison
            results = framework.run_strategy_comparison(model, args.compare_dataset, args.max_samples)
            framework.save_results(results, f"{args.compare_dataset}_strategy_comparison.json")
        
        else:
            print("❌ Please specify --compare-single, --benchmark, or --compare-dataset")
            return 1
        
        print(f"\n✅ Research experiment completed successfully!")
        print(f"📁 All results saved to: {framework.experiment_dir}")
        
        return 0
        
    except Exception as e:
        print(f"❌ Research experiment failed: {e}")
        import traceback
        traceback.print_exc()
        return 1


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