test_adaptive_cot.py

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

Usage:
    python test_adaptive_cot.py --problem "Sarah has 12 apples..." --adaptive
    python test_adaptive_cot.py --problem "What is 2+2?" --static --branches 3
"""

import argparse
import sys
import os
import time
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="Test Adaptive CoT Framework",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Test adaptive branching
  python test_adaptive_cot.py --problem "Sarah has 12 apples..." --adaptive
  
  # Test static branching
  python test_adaptive_cot.py --problem "What is 2+2?" --static --branches 3
        """
    )
    
    # Problem input
    parser.add_argument(
        "--problem", 
        type=str, 
        required=True,
        help="Math problem to solve"
    )
    
    # Test mode
    parser.add_argument(
        "--adaptive", 
        action="store_true",
        help="Use adaptive branching based on prefill analysis"
    )
    parser.add_argument(
        "--static", 
        action="store_true",
        help="Use static branching with fixed number of branches"
    )
    
    # Branching configuration
    parser.add_argument(
        "--branches", 
        type=int, 
        default=3,
        help="Number of branches for static mode (default: 3)"
    )
    parser.add_argument(
        "--min-branches", 
        type=int, 
        default=1,
        help="Minimum number of branches for adaptive mode (default: 1)"
    )
    parser.add_argument(
        "--max-branches", 
        type=int, 
        default=10,
        help="Maximum number of branches for adaptive mode (default: 10)"
    )
    
    # Model configuration
    parser.add_argument(
        "--model-path", 
        type=str, 
        default="/raid/LLM/deepseek-r1-distill-qwen-1.5b",
        help="Path to model (default: /raid/LLM/deepseek-r1-distill-qwen-1.5b)"
    )
    parser.add_argument(
        "--model-type", 
        type=str, 
        default="deepseek-r1-distill-qwen",
        help="Model type (default: deepseek-r1-distill-qwen)"
    )
    
    # 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 Framework Test")
    print("=" * 50)
    print("Adaptive Chain-of-Thought with self-consistency:")
    print("1. First prefill: Analyze problem difficulty → get signals")
    print("2. Determine branch count based on signals")
    print("3. Second prefill: Generate multiple reasoning paths")
    print("4. Apply self-consistency (majority voting)")
    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()
        
        # Determine test mode
        adaptive = args.adaptive or not args.static  # Default to adaptive unless --static
        
        # Create CoT configuration
        cot_config = {
            "adaptive_branching": adaptive,
            "min_branches": args.min_branches,
            "max_branches": args.max_branches,
            "default_branches": args.branches,
            "research_logging": False,
        }
        
        # Create AdaptiveCoT
        cot = AdaptiveCoT(model, cot_config)
        
        print(f"✅ AdaptiveCoT created successfully!")
        print(f"   - Backend: {cot.backend_type}")
        print(f"   - Adaptive branching: {cot.adaptive_branching}")
        print(f"   - Branch range: {cot.min_branches}-{cot.max_branches}")
        print()
        
        # Solve problem
        print(f"🔍 Solving problem: {args.problem}")
        print("-" * 50)
        
        start_time = time.time()
        result = cot.solve_problem(args.problem)
        end_time = time.time()
        
        # Display results
        print(f"\\n📊 Results:")
        print(f"   Final Answer: {result['final_answer']}")
        print(f"   Branches Used: {result['num_branches']}")
        print(f"   Strategy: {result['allocation_info']['strategy']}")
        print(f"   Backend: {result['backend_type']}")
        print(f"   Execution Time: {result['execution_time']:.2f}s")
        print(f"   Consensus Confidence: {result['consensus_info']['confidence']:.3f}")
        
        # Show prefill analysis if available
        if 'prefill_signals' in result:
            prefill = result['prefill_signals']
            print(f"\\n🔍 Prefill Analysis:")
            print(f"   Entropy: {prefill['entropy']:.3f}")
            print(f"   KL Divergence: {prefill['kl_divergence']:.3f}")
            print(f"   Confidence: {prefill['confidence']:.3f}")
            if 'reasoning' in result['allocation_info']:
                print(f"   Allocation Reasoning: {result['allocation_info']['reasoning']}")
        
        # Show reasoning paths (first 2)
        print(f"\\n🌿 Reasoning Paths ({len(result['reasoning_paths'])}):")
        for i, path in enumerate(result['reasoning_paths'][:2]):
            print(f"   Branch {i+1}: {path[:80]}...")
        if len(result['reasoning_paths']) > 2:
            print(f"   ... and {len(result['reasoning_paths']) - 2} more")
        
        # Show answer extraction
        print(f"\\n🎯 Answer Extraction:")
        for i, answer in enumerate(result['extracted_answers']):
            print(f"   Branch {i+1}: {answer}")
        
        print(f"\\n✅ 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: {cot.backend_type}")
            print(f"   - Adaptive Branching: {adaptive}")
            print(f"   - Static Branches: {args.branches}")
            print(f"   - Branch Range: {args.min_branches}-{args.max_branches}")
            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"❌ Test failed: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    main()