run_quick_test.py

#!/usr/bin/env python3
"""
Quick test script for Adaptive CoT Framework.

This runs a quick test with a few samples to verify everything works.

Usage:
    python run_quick_test.py --model-path "/path/to/model"
    python run_quick_test.py --model-path "/path/to/model" --problem "Your custom problem"
"""

import argparse
import sys
import os
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


def create_parser():
    """Create command line argument parser."""
    parser = argparse.ArgumentParser(
        description="Quick test for Adaptive CoT Framework",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Quick test with default problem
  python run_quick_test.py --model-path "/path/to/model"
  
  # Quick test with custom problem
  python run_quick_test.py --model-path "/path/to/model" --problem "Sarah has 12 apples..."
        """
    )
    
    # 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)"
    )
    
    # Test configuration
    parser.add_argument(
        "--problem", 
        type=str, 
        default="Sarah has 12 apples. She gives 3 to her friend and buys 5 more. How many apples does she have now?",
        help="Math problem to solve"
    )
    parser.add_argument(
        "--test-adaptive", 
        action="store_true",
        help="Test adaptive branching (default: True)"
    )
    parser.add_argument(
        "--test-static", 
        action="store_true",
        help="Test static branching with 3 branches"
    )
    
    # Debug options
    parser.add_argument(
        "--verbose", 
        action="store_true",
        help="Enable verbose output"
    )
    
    return parser


def main():
    """Main test function."""
    parser = create_parser()
    args = parser.parse_args()
    
    print("🧪 Adaptive CoT Quick Test")
    print("=" * 40)
    print("Quick verification that the framework works correctly")
    print()
    
    try:
        # Create model
        model_config = {
            "model_name": args.model_path,
            "generation_params": {
                "max_new_tokens": 2048,
                "temperature": 0.6,
                "top_p": 0.95,
            }
        }
        
        print(f"📦 Loading model: {args.model_path}")
        model = ModelFactory.create_model(args.model_type, args.model_path, model_config)
        model.load_model()
        
        # Test configurations
        test_configs = []
        
        if args.test_adaptive or (not args.test_static and not args.test_adaptive):
            test_configs.append(("Adaptive", {
                "adaptive_branching": True,
                "min_branches": 1,
                "max_branches": 5,
                "research_logging": False,
            }))
        
        if args.test_static:
            test_configs.append(("Static", {
                "adaptive_branching": False,
                "min_branches": 1,
                "max_branches": 5,
                "default_branches": 3,
                "research_logging": False,
            }))
        
        # Run tests
        for test_name, config in test_configs:
            print(f"\\n🔍 Testing {test_name} Branching:")
            print("-" * 30)
            
            # Create CoT
            cot = AdaptiveCoT(model, config)
            
            # Solve problem
            result = cot.solve_problem(args.problem)
            
            # Display results
            print(f"\\n📊 {test_name} Results:")
            print(f"   Final Answer: {result['final_answer']}")
            print(f"   Branches Used: {result['num_branches']}")
            print(f"   Strategy: {result['allocation_info']['strategy']}")
            print(f"   Execution Time: {result['execution_time']:.2f}s")
            print(f"   Consensus Confidence: {result['consensus_info']['confidence']:.3f}")
            
            if 'prefill_signals' in result:
                prefill = result['prefill_signals']
                print(f"   Prefill Analysis:")
                print(f"     Entropy: {prefill['entropy']:.3f}")
                print(f"     KL Divergence: {prefill['kl_divergence']:.3f}")
                print(f"     Confidence: {prefill['confidence']:.3f}")
        
        print(f"\\n✅ Quick test completed successfully!")
        
        if args.verbose:
            print(f"\\n🔧 Technical Details:")
            print(f"   - Model: {args.model_path}")
            print(f"   - Model Type: {args.model_type}")
            print(f"   - Backend: Auto-detected")
            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"❌ Quick test failed: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()