evaluate_orbit.py

#!/usr/bin/env python3
"""
Comprehensive evaluation script for ORBIT neural reasoner.

This script provides a complete evaluation pipeline including:
- Task-specific accuracy computation
- Anytime performance curves
- Calibration metrics (ECE, Brier score)
- Out-of-distribution robustness testing
- Option usage analysis

Usage:
    python evaluate_orbit.py --model_path checkpoints/best_model.pt --data_path data/test_dataset.pt
    python evaluate_orbit.py --config config/evaluation.json --output_dir results/
    python evaluate_orbit.py --model_path model.pt --quick_eval  # Fast evaluation with subset

Requirements: 8.1, 8.2, 8.3, 8.4, 8.5
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Dict, Any, Optional

import torch
import numpy as np

# Add orbit to path
sys.path.append(str(Path(__file__).parent))

from orbit.models.orbit_model import ORBITModel
from orbit.models.data_models import ModelConfig
from orbit.training.data_loader import MultiTaskDataLoader, MultiTaskDataset
from orbit.evaluation import (
    ORBITEvaluator, EvaluationConfig, EvaluationResults,
    AccuracyComputer, AnytimeCurveGenerator, CalibrationMetrics,
    OODEvaluator, CorruptionConfig, CorruptionType
)

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def parse_arguments() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Comprehensive evaluation for ORBIT neural reasoner",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    # Model and data paths
    parser.add_argument(
        "--model_path", 
        type=str, 
        required=True,
        help="Path to trained ORBIT model checkpoint"
    )
    parser.add_argument(
        "--data_path", 
        type=str,
        help="Path to test dataset (if not provided, will use default test sets)"
    )
    parser.add_argument(
        "--config_path", 
        type=str,
        help="Path to model configuration file"
    )
    
    # Evaluation configuration
    parser.add_argument(
        "--eval_config", 
        type=str,
        help="Path to evaluation configuration JSON file"
    )
    parser.add_argument(
        "--output_dir", 
        type=str, 
        default="./evaluation_results",
        help="Directory to save evaluation results"
    )
    
    # Evaluation options
    parser.add_argument(
        "--quick_eval", 
        action="store_true",
        help="Run quick evaluation on subset of data"
    )
    parser.add_argument(
        "--skip_ood", 
        action="store_true",
        help="Skip OOD robustness evaluation"
    )
    parser.add_argument(
        "--skip_options", 
        action="store_true",
        help="Skip option usage analysis"
    )
    parser.add_argument(
        "--batch_size", 
        type=int, 
        default=32,
        help="Batch size for evaluation"
    )
    parser.add_argument(
        "--max_samples", 
        type=int,
        help="Maximum number of samples to evaluate (for quick testing)"
    )
    
    # Device settings
    parser.add_argument(
        "--device", 
        type=str, 
        default="auto",
        choices=["auto", "cuda", "cpu"],
        help="Device for evaluation"
    )
    
    # Output options
    parser.add_argument(
        "--save_detailed", 
        action="store_true",
        help="Save detailed evaluation outputs"
    )
    parser.add_argument(
        "--no_save", 
        action="store_true",
        help="Don't save results to file"
    )
    parser.add_argument(
        "--verbose", 
        action="store_true",
        help="Enable verbose logging"
    )
    
    return parser.parse_args()


def load_model_and_config(model_path: str, config_path: Optional[str] = None) -> tuple[ORBITModel, ModelConfig]:
    """
    Load ORBIT model and configuration from checkpoint.
    
    Args:
        model_path: Path to model checkpoint
        config_path: Optional path to config file
        
    Returns:
        Tuple of (model, config)
    """
    logger.info(f"Loading model from {model_path}")
    
    # Load checkpoint
    checkpoint = torch.load(model_path, map_location='cpu')
    
    # Load config
    if config_path:
        logger.info(f"Loading config from {config_path}")
        with open(config_path, 'r') as f:
            config_dict = json.load(f)
        config = ModelConfig(**config_dict)
    elif 'config' in checkpoint:
        logger.info("Using config from checkpoint")
        config = checkpoint['config']
    else:
        logger.warning("No config found, using default")
        config = ModelConfig()
    
    # Create model
    model = ORBITModel(config)
    
    # Load state dict
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
    else:
        # Assume checkpoint is the state dict
        model.load_state_dict(checkpoint)
    
    logger.info(f"Model loaded successfully. Parameters: {model.count_parameters():,}")
    
    return model, config


def load_evaluation_config(eval_config_path: Optional[str], args: argparse.Namespace) -> EvaluationConfig:
    """
    Load evaluation configuration from file or create from arguments.
    
    Args:
        eval_config_path: Path to evaluation config file
        args: Command line arguments
        
    Returns:
        EvaluationConfig object
    """
    if eval_config_path and Path(eval_config_path).exists():
        logger.info(f"Loading evaluation config from {eval_config_path}")
        with open(eval_config_path, 'r') as f:
            config_dict = json.load(f)
        eval_config = EvaluationConfig(**config_dict)
    else:
        logger.info("Creating evaluation config from arguments")
        eval_config = EvaluationConfig()
    
    # Override with command line arguments
    eval_config.batch_size = args.batch_size
    eval_config.device = args.device
    eval_config.results_dir = args.output_dir
    eval_config.save_results = not args.no_save
    eval_config.save_detailed_outputs = args.save_detailed
    eval_config.enable_ood_evaluation = not args.skip_ood
    eval_config.enable_option_analysis = not args.skip_options
    
    # Quick evaluation settings
    if args.quick_eval:
        eval_config.halting_thresholds = [0.3, 0.5, 0.7, 1.0]  # Fewer thresholds
        eval_config.calibration_bins = 5  # Fewer bins
        eval_config.max_iterations = 4  # Fewer iterations
        
        # Simpler OOD corruptions for quick eval
        if eval_config.enable_ood_evaluation:
            eval_config.ood_corruption_configs = [
                CorruptionConfig(CorruptionType.PALETTE_CHANGE, strength=1.0),
                CorruptionConfig(CorruptionType.ROTATION, strength=1.0, parameters={'angle': 90}),
                CorruptionConfig(CorruptionType.NOISE_INJECTION, strength=0.1)
            ]
    
    return eval_config


def load_test_data(data_path: Optional[str], 
                  batch_size: int, 
                  max_samples: Optional[int] = None) -> MultiTaskDataLoader:
    """
    Load test dataset for evaluation.
    
    Args:
        data_path: Path to test data
        batch_size: Batch size for data loader
        max_samples: Maximum number of samples to load
        
    Returns:
        MultiTaskDataLoader for test data
    """
    if data_path and Path(data_path).exists():
        logger.info(f"Loading test data from {data_path}")
        # Load custom test dataset
        test_data = torch.load(data_path)
        if isinstance(test_data, dict):
            dataset = MultiTaskDataset.from_dict(test_data)
        else:
            dataset = test_data
    else:
        logger.info("Creating synthetic test dataset")
        # Create synthetic test data for demonstration
        dataset = create_synthetic_test_dataset(max_samples or 1000)
    
    # Limit dataset size if requested
    if max_samples and len(dataset) > max_samples:
        logger.info(f"Limiting dataset to {max_samples} samples")
        indices = torch.randperm(len(dataset))[:max_samples]
        dataset = torch.utils.data.Subset(dataset, indices)
    
    # Create data loader
    dataloader = MultiTaskDataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=False,  # Don't shuffle for evaluation
        num_workers=0   # Single-threaded for evaluation
    )
    
    logger.info(f"Test dataset loaded. Size: {len(dataset)}, Batches: {len(dataloader)}")
    
    return dataloader


def create_synthetic_test_dataset(num_samples: int) -> MultiTaskDataset:
    """
    Create synthetic test dataset for demonstration.
    
    Args:
        num_samples: Number of samples to generate
        
    Returns:
        MultiTaskDataset with synthetic data
    """
    logger.info(f"Creating synthetic test dataset with {num_samples} samples")
    
    # This is a placeholder - in practice, you would load real test data
    # For now, create minimal synthetic data to demonstrate the evaluation pipeline
    
    grid_tasks = []
    text_tasks = []
    
    # Generate synthetic grid tasks (ARC-like)
    for i in range(num_samples // 2):
        # Simple 5x5 grids with random patterns
        input_grid = torch.randint(0, 10, (5, 5, 1))
        target_grid = torch.randint(0, 10, (5, 5))  # Target is 2D
        
        grid_tasks.append({
            'input_grid': input_grid,
            'target_grid': target_grid,
            'task_type': 'ARC'
        })
    
    # Generate synthetic text tasks (NLI-like)
    for i in range(num_samples // 2):
        # Simple token sequences
        input_tokens = torch.randint(4, 1000, (20,))  # Avoid special tokens 0-3
        target_tokens = torch.randint(4, 1000, (10,))
        
        text_tasks.append({
            'input_tokens': input_tokens,
            'target_tokens': target_tokens,
            'task_type': 'NLI'
        })
    
    # Combine into dataset
    all_tasks = grid_tasks + text_tasks
    
    # Create MultiTaskDataset
    dataset = MultiTaskDataset(all_tasks)
    
    return dataset


def run_evaluation(model: ORBITModel, 
                  eval_config: EvaluationConfig,
                  test_dataloader: MultiTaskDataLoader) -> EvaluationResults:
    """
    Run comprehensive evaluation.
    
    Args:
        model: ORBIT model to evaluate
        eval_config: Evaluation configuration
        test_dataloader: Test data loader
        
    Returns:
        EvaluationResults with all metrics
    """
    logger.info("Starting comprehensive ORBIT evaluation...")
    
    # Create evaluator
    evaluator = ORBITEvaluator(model, eval_config)
    
    # Run evaluation
    results = evaluator.evaluate(test_dataloader)
    
    return results


def print_results_summary(results: EvaluationResults) -> None:
    """Print a formatted summary of evaluation results."""
    print("\n" + "="*60)
    print("ORBIT EVALUATION RESULTS SUMMARY")
    print("="*60)
    
    # Overall metrics
    print(f"\nOverall Performance:")
    print(f"  Accuracy: {results.overall_accuracy:.4f}")
    print(f"  Energy at Halt: {results.energy_at_halt:.4f}")
    print(f"  Average Halting Mass: {results.average_halting_mass:.4f}")
    
    # Task-specific accuracies
    print(f"\nTask-Specific Accuracies:")
    for task_type, accuracy in results.task_accuracies.items():
        print(f"  {task_type}: {accuracy.accuracy:.4f} ({accuracy.correct_predictions}/{accuracy.num_samples})")
    
    # Calibration metrics
    print(f"\nCalibration Metrics:")
    print(f"  ECE: {results.calibration_results.ece:.4f}")
    print(f"  Brier Score: {results.calibration_results.brier_score:.4f}")
    print(f"  Reliability: {results.calibration_results.reliability:.4f}")
    print(f"  Confidence: {results.calibration_results.confidence:.4f}")
    
    # Anytime performance
    if results.anytime_curve:
        print(f"\nAnytime Performance (sample points):")
        for i, point in enumerate(results.anytime_curve[::2]):  # Show every other point
            print(f"  H={point.halting_mass:.2f}: Acc={point.accuracy:.4f}, E={point.energy:.4f}")
    
    # OOD robustness
    if results.ood_results:
        print(f"\nOOD Robustness:")
        print(f"  Average Robustness: {results.average_robustness:.4f}")
        for ood_result in results.ood_results:
            print(f"  {ood_result.corruption_type}: {ood_result.robustness_score:.4f}")
    
    # Option analysis
    if results.option_analysis:
        effectiveness = results.option_analysis.get('effectiveness_analysis', {})
        print(f"\nOption Analysis:")
        if 'most_effective_option' in effectiveness:
            option_id, improvement = effectiveness['most_effective_option']
            print(f"  Most Effective Option: #{option_id} (ΔE={improvement:.4f})")
        if 'average_energy_improvement' in effectiveness:
            print(f"  Average Energy Improvement: {effectiveness['average_energy_improvement']:.4f}")
        if 'usage_entropy' in effectiveness:
            print(f"  Usage Diversity (entropy): {effectiveness['usage_entropy']:.4f}")
    
    # Performance info
    print(f"\nEvaluation Info:")
    print(f"  Samples: {results.num_samples}")
    print(f"  Time: {results.evaluation_time:.2f}s")
    print(f"  Samples/sec: {results.num_samples / results.evaluation_time:.1f}")
    
    print("="*60)


def main():
    """Main evaluation function."""
    args = parse_arguments()
    
    # Setup verbose logging if requested
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    try:
        # Load model and configuration
        model, model_config = load_model_and_config(args.model_path, args.config_path)
        
        # Load evaluation configuration
        eval_config = load_evaluation_config(args.eval_config, args)
        
        # Load test data
        test_dataloader = load_test_data(args.data_path, args.batch_size, args.max_samples)
        
        # Run evaluation
        results = run_evaluation(model, eval_config, test_dataloader)
        
        # Print results summary
        print_results_summary(results)
        
        # Save results if requested
        if not args.no_save:
            logger.info(f"Results saved to {eval_config.results_dir}")
        
        logger.info("Evaluation completed successfully!")
        
    except Exception as e:
        logger.error(f"Evaluation failed: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()