run_experiment.py

#!/usr/bin/env python3
"""
Main experiment runner for Adaptive CoT Framework.

This is the main entry point for running research experiments.

Usage:
    python run_experiment.py --model-path "/path/to/model" --datasets gsm8k math --strategies adaptive static
    python run_experiment.py --model-path "/path/to/model" --datasets gsm8k --max-samples 50 --verbose
"""

import argparse
import sys
import os
import json
from pathlib import Path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.experiments.experiment_runner import ExperimentRunner


def create_parser():
    """Create command line argument parser."""
    parser = argparse.ArgumentParser(
        description="Run Adaptive CoT research experiments",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run full evaluation on GSM8K and MATH datasets
  python run_experiment.py --model-path "/path/to/model" --datasets gsm8k math --strategies adaptive static
  
  # Run with limited samples for quick testing
  python run_experiment.py --model-path "/path/to/model" --datasets gsm8k --max-samples 50 --verbose
  
  # Run overnight experiment
  python run_experiment.py --model-path "/path/to/model" --datasets gsm8k math aime --overnight
        """
    )
    
    # Model configuration
    parser.add_argument(
        "--model-path", 
        type=str, 
        required=True,
        help="Path to the model"
    )
    parser.add_argument(
        "--model-type", 
        type=str, 
        default="deepseek-r1-distill-qwen",
        help="Model type (default: deepseek-r1-distill-qwen)"
    )
    
    # Dataset configuration
    parser.add_argument(
        "--datasets", 
        nargs="+", 
        default=["gsm8k"],
        help="Datasets to evaluate on (default: gsm8k)"
    )
    parser.add_argument(
        "--max-samples", 
        type=int, 
        default=100,
        help="Maximum samples per dataset (default: 100)"
    )
    
    # Strategy configuration
    parser.add_argument(
        "--strategies", 
        nargs="+", 
        default=["adaptive", "static"],
        help="Strategies to compare (default: adaptive static)"
    )
    parser.add_argument(
        "--static-branches", 
        type=int, 
        default=3,
        help="Number of branches for static strategy (default: 3)"
    )
    
    # Experiment configuration
    parser.add_argument(
        "--experiment-dir", 
        type=str, 
        default="experiments",
        help="Directory to save experiment results (default: experiments)"
    )
    parser.add_argument(
        "--overnight", 
        action="store_true",
        help="Run overnight experiment with more samples"
    )
    
    # Debug options
    parser.add_argument(
        "--verbose", 
        action="store_true",
        help="Enable verbose output"
    )
    
    return parser


def create_experiment_config(args):
    """Create experiment configuration from arguments."""
    config = {
        "model": {
            "model_name": args.model_path,
            "model_type": args.model_type,
            "generation_params": {
                "max_new_tokens": 2048,
                "temperature": 0.6,
                "top_p": 0.95,
            }
        },
        "datasets": args.datasets,
        "max_samples_per_dataset": args.max_samples,
        "strategies": args.strategies,
        "static_branches": args.static_branches,
        "experiment_dir": args.experiment_dir,
        "verbose": args.verbose,
        "research_logging": True,
    }
    
    # Overnight experiment configuration
    if args.overnight:
        config["max_samples_per_dataset"] = 1000
        config["datasets"] = ["gsm8k", "math", "aime"]
        config["strategies"] = ["adaptive", "static"]
        print("🌙 Running overnight experiment with extended configuration")
    
    return config


def main():
    """Main experiment function."""
    parser = create_parser()
    args = parser.parse_args()
    
    print("🔬 Adaptive CoT Research Experiment")
    print("=" * 50)
    print("Running systematic evaluation of adaptive branching strategies")
    print()
    
    try:
        # Create experiment configuration
        config = create_experiment_config(args)
        
        print(f"📦 Model: {args.model_path}")
        print(f"📊 Datasets: {', '.join(args.datasets)}")
        print(f"🎯 Strategies: {', '.join(args.strategies)}")
        print(f"📈 Max samples per dataset: {args.max_samples}")
        print(f"📁 Experiment directory: {args.experiment_dir}")
        print()
        
        # Create experiment runner
        runner = ExperimentRunner(config)
        
        # Run full evaluation
        print("🚀 Starting experiment...")
        results = runner.run_full_evaluation(
            model_name=args.model_path,
            datasets=args.datasets,
            strategies=args.strategies,
            max_samples_per_dataset=args.max_samples
        )
        
        # Display results summary
        print("\\n📊 Experiment Results Summary:")
        print("-" * 40)
        
        for dataset, dataset_results in results.items():
            if isinstance(dataset_results, dict) and "overall" in dataset_results:
                overall = dataset_results["overall"]
                print(f"\\n{dataset.upper()}:")
                print(f"  Accuracy: {overall.get('accuracy', 0.0):.3f}")
                print(f"  Total samples: {overall.get('total_samples', 0)}")
                print(f"  Execution time: {overall.get('execution_time', 0.0):.2f}s")
                
                # Show strategy comparison
                if "strategy_comparison" in dataset_results:
                    print(f"  Strategy comparison:")
                    for strategy, strategy_results in dataset_results["strategy_comparison"].items():
                        accuracy = strategy_results.get("accuracy", 0.0)
                        branches = strategy_results.get("avg_branches", 0)
                        print(f"    {strategy}: {accuracy:.3f} accuracy, {branches:.1f} avg branches")
        
        print(f"\\n✅ Experiment completed successfully!")
        print(f"📁 Results saved to: {runner.run_dir}")
        
        if args.verbose:
            print(f"\\n🔧 Technical Details:")
            print(f"   - Model type: {args.model_type}")
            print(f"   - Backend: Auto-detected")
            print(f"   - Adaptive branching: Based on prefill analysis")
            print(f"   - Self-consistency: Majority voting with do_sample=True")
            print(f"   - Generation: num_return_sequences for parallel processing")
        
    except Exception as e:
        print(f"❌ Experiment failed: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()