fix eval

scripts/eval.py

@@ -1,271 +1,489 @@
 #!/usr/bin/env python3
 """
-eval.py Evaluation script for modified OLMoE model using lm-evaluation-harness
+eval.py - Evaluation script for OLMoE models using lm-evaluation-harness
+
+This script supports evaluation of both:
+1. Standard Transformers OLMoE models
+2. Custom MyOLMoE models with modified routing
+
+Usage Examples:
+    # Evaluate standard OLMoE model
+    python eval.py --model_type transformers --tasks mmlu hellaswag
+
+    # Evaluate custom MyOLMoE model with non-deterministic routing
+    python eval.py --model_type custom --routing_type non_deterministic --tasks mmlu
 """
+
 import argparse
 import json
 import os
-from typing import Dict, List, Optional
+import sys
+import logging
+from typing import Dict, List, Optional, Any
+import numpy as np
 import torch
-from transformers import AutoConfig, AutoTokenizer
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
+
+# lm-eval imports
 from lm_eval import evaluator
-# Remove the problematic import - we don't need get_model
-import logging
+from lm_eval.models.huggingface import HFLM
 
 # Set up logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
 
 def parse_args():
     """Parse command line arguments."""
-    parser = argparse.ArgumentParser(description="Evaluate myolmoe model")
+    parser = argparse.ArgumentParser(
+        description="Evaluate OLMoE models using lm-evaluation-harness",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Standard OLMoE evaluation
+  python eval.py --model_type transformers --tasks mmlu arc_easy
+
+  # Custom MyOLMoE with non-deterministic routing
+  python eval.py --model_type custom --routing_type non_deterministic \\
+    --router_temperature 0.8 --tasks mmlu hellaswag
+
+  # Dense routing evaluation
+  python eval.py --model_type custom --routing_type dense --tasks gsm8k
+        """
+    )
     
     # Model arguments
-    parser.add_argument("--model_path", type=str, default="/root/.cache/huggingface/hub/models--allenai--OLMoE-7B/snapshots/6d84c48581ece794365f2b8e9cfb043c68ade9c5",
-                       help="Path to the pretrained model")
-    parser.add_argument("--model_type", type=str, default="hf-auto",
-                       help="Model type for lm-eval")
-    
-    # Routing configuration
-    parser.add_argument("--routing_type", type=str, default="non_deterministic",
-                       choices=["dense", "sparse", "non_deterministic"],
-                       help="Type of routing to use")
-    parser.add_argument("--router_temperature", type=float, default=1.0,
-                       help="Temperature for non-deterministic routing")
-    parser.add_argument("--num_experts_per_tok", type=int, default=8,
-                       help="Number of experts per token")
+    parser.add_argument(
+        "--model_path", 
+        type=str, 
+        default="allenai/OLMoE-1B-7B-0924",
+        help="Path or name of the pretrained model"
+    )
+    parser.add_argument(
+        "--model_type", 
+        type=str, 
+        default="transformers",
+        choices=["transformers", "custom"],
+        help="Model type: 'transformers' for standard OLMoE, 'custom' for MyOLMoE"
+    )
+    parser.add_argument(
+        "--custom_model_path",
+        type=str,
+        default="./myolmoe_model",
+        help="Path to custom MyOLMoE model code (when using --model_type custom)"
+    )
+    
+    # Routing configuration (only for custom models)
+    parser.add_argument(
+        "--routing_type", 
+        type=str, 
+        default="sparse",
+        choices=["dense", "sparse", "non_deterministic"],
+        help="Routing type (only used with custom models)"
+    )
+    parser.add_argument(
+        "--router_temperature", 
+        type=float, 
+        default=1.0,
+        help="Temperature for non-deterministic routing"
+    )
+    parser.add_argument(
+        "--num_experts_per_tok", 
+        type=int, 
+        default=8,
+        help="Number of experts per token"
+    )
     
     # Evaluation arguments
-    parser.add_argument("--tasks", type=str, nargs="+", 
-                       default=['mmlu'],
-                       # , 'gsm8k'
-                    #    default=["hellaswag", "arc_easy", "arc_challenge", "winogrande"],
-                       help="Tasks to evaluate on")
-    parser.add_argument("--num_fewshot", type=int, default=0,
-                       help="Number of few-shot examples")
-    parser.add_argument("--batch_size", type=int, default=64,
-                       help="Batch size for evaluation")
-    parser.add_argument("--max_batch_size", type=int, default=None,
-                       help="Maximum batch size")
-    parser.add_argument("--device", type=str, default="cuda",
-                       help="Device to use for evaluation")
-    parser.add_argument("--dtype", type=str, default="float16",
-                       choices=["float16", "bfloat16", "float32"],
-                       help="Data type for model weights")
+    parser.add_argument(
+        "--tasks", 
+        type=str, 
+        nargs="+", 
+        default=["mmlu"],
+        help="Tasks to evaluate on (e.g., mmlu, hellaswag, arc_easy, gsm8k)"
+    )
+    parser.add_argument(
+        "--num_fewshot", 
+        type=int, 
+        default=0,
+        help="Number of few-shot examples"
+    )
+    parser.add_argument(
+        "--batch_size", 
+        type=int, 
+        default=8,
+        help="Batch size for evaluation"
+    )
+    parser.add_argument(
+        "--max_batch_size", 
+        type=int, 
+        default=None,
+        help="Maximum batch size (auto if None)"
+    )
+    parser.add_argument(
+        "--device", 
+        type=str, 
+        default="auto",
+        help="Device to use ('auto', 'cuda', 'cpu')"
+    )
+    parser.add_argument(
+        "--dtype", 
+        type=str, 
+        default="auto",
+        choices=["auto", "float16", "bfloat16", "float32"],
+        help="Data type for model weights"
+    )
     
     # Output arguments
-    parser.add_argument("--output_dir", type=str, default="./eval_results",
-                       help="Directory to save evaluation results")
-    parser.add_argument("--output_filename", type=str, default=None,
-                       help="Filename for results (auto-generated if not provided)")
+    parser.add_argument(
+        "--output_dir", 
+        type=str, 
+        default="./eval_results",
+        help="Directory to save evaluation results"
+    )
+    parser.add_argument(
+        "--output_filename", 
+        type=str, 
+        default=None,
+        help="Custom filename for results (auto-generated if not provided)"
+    )
     
     # Additional arguments
-    parser.add_argument("--limit", type=int, default=None,
-                       help="Limit number of examples per task")
-    parser.add_argument("--write_out", action="store_true",
-                       help="Write out individual predictions")
-    parser.add_argument("--trust_remote_code", action="store_true",
-                       help="Trust remote code when loading model")
-    parser.add_argument("--verbosity", type=str, default="INFO",
-                       choices=["DEBUG", "INFO", "WARNING", "ERROR"],
-                       help="Logging verbosity level")
+    parser.add_argument(
+        "--limit", 
+        type=int, 
+        default=None,
+        help="Limit number of examples per task (for testing)"
+    )
+    parser.add_argument(
+        "--write_out", 
+        action="store_true",
+        help="Write out individual predictions to files"
+    )
+    parser.add_argument(
+        "--trust_remote_code", 
+        action="store_true",
+        help="Trust remote code when loading model"
+    )
+    parser.add_argument(
+        "--verbosity", 
+        type=str, 
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Logging verbosity level"
+    )
     
     return parser.parse_args()
 
 
-def setup_model_config(model_path: str, routing_config: Dict) -> None:
+def load_transformers_model(args) -> HFLM:
     """
-    Update model configuration with routing settings.
-    """
-    config_path = os.path.join(model_path, "config.json")
+    Load standard Transformers OLMoE model.
     
-    if os.path.exists(config_path):
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-        
-        # Update routing configuration
-        config.update(routing_config)
+    Args:
+        args: Parsed command line arguments
         
-        # Save updated config
-        with open(config_path, 'w') as f:
-            json.dump(config, f, indent=2)
-        
-        logger.info(f"Updated model config with routing settings: {routing_config}")
-    else:
-        logger.warning(f"Config file not found at {config_path}")
+    Returns:
+        HFLM: Wrapped model ready for evaluation
+    """
+    logger.info(f"Loading Transformers OLMoE model: {args.model_path}")
+    
+    # Create HFLM model directly
+    model = HFLM(
+        pretrained=args.model_path,
+        device=args.device,
+        batch_size=args.batch_size,
+        max_batch_size=args.max_batch_size,
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code
+    )
+    
+    logger.info("Transformers model loaded successfully")
+    return model
 
 
-def validate_model_setup(model_path: str) -> bool:
+def load_custom_model(args) -> HFLM:
     """
-    Validate that the model can be loaded with the current configuration.
+    Load custom MyOLMoE model with routing configuration.
+    
+    Args:
+        args: Parsed command line arguments
+        
+    Returns:
+        HFLM: Wrapped model ready for evaluation
     """
+    logger.info(f"Loading custom MyOLMoE model: {args.model_path}")
+    logger.info(f"Routing configuration: {args.routing_type}")
+    
+    # Add custom model path to Python path
+    if os.path.exists(args.custom_model_path):
+        sys.path.insert(0, args.custom_model_path)
+        logger.info(f"Added {args.custom_model_path} to Python path")
+    else:
+        logger.warning(f"Custom model path not found: {args.custom_model_path}")
+    
     try:
-        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-        
-        logger.info(f"Model validation successful:")
-        logger.info(f"  - Model type: {config.model_type}")
-        logger.info(f"  - Routing type: {getattr(config, 'routing_type', 'not specified')}")
-        logger.info(f"  - Vocab size: {config.vocab_size}")
-        logger.info(f"  - Hidden size: {config.hidden_size}")
-        logger.info(f"  - Num layers: {config.num_hidden_layers}")
-        logger.info(f"  - Num experts: {getattr(config, 'num_experts', 'not specified')}")
-        
-        return True
-    except Exception as e:
-        logger.error(f"Model validation failed: {e}")
-        return False
-
-def run_evaluation(args) -> Dict:
-    """Run evaluation with properly wrapped model."""
-    from transformers import AutoModelForCausalLM
-    import sys, os
-    sys.path.insert(0, os.path.join(os.path.dirname(__file__), "myolmoe_model"))
-
-    # 1. Load config and override routing parameters
+        # Import custom model class
+        from modeling_myolmoe import MyOLMoEForCausalLM
+        logger.info("Successfully imported MyOLMoEForCausalLM")
+    except ImportError as e:
+        logger.error(f"Failed to import custom model: {e}")
+        logger.error("Make sure the custom model code is available in the specified path")
+        raise
+    
+    # Load and configure model
     config = AutoConfig.from_pretrained(
         args.model_path,
-        trust_remote_code=True
+        trust_remote_code=args.trust_remote_code
     )
+    
+    # Override routing configuration
     config.routing_type = args.routing_type
     config.router_temperature = args.router_temperature
     config.num_experts_per_tok = args.num_experts_per_tok
-
-
-    # 2. Load model with updated config
-    torch_dtype = {
-        "float16": torch.float16,
-        "bfloat16": torch.bfloat16,
-        "float32": torch.float32
-    }[args.dtype]
-
-    from modeling_myolmoe import MyOLMoEForCausalLM
-
+    
+    logger.info(f"Model config updated:")
+    logger.info(f"  - routing_type: {config.routing_type}")
+    logger.info(f"  - router_temperature: {config.router_temperature}")
+    logger.info(f"  - num_experts_per_tok: {config.num_experts_per_tok}")
+    
+    # Determine torch dtype
+    if args.dtype == "auto":
+        torch_dtype = "auto"
+    else:
+        torch_dtype = {
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "float32": torch.float32
+        }[args.dtype]
+    
+    # Load the custom model
     hf_model = MyOLMoEForCausalLM.from_pretrained(
         args.model_path,
         config=config,
         torch_dtype=torch_dtype,
-        device_map="auto"
+        device_map="auto" if args.device == "auto" else None,
+        trust_remote_code=args.trust_remote_code
     ).eval()
-
-
-    # 3. Wrap the Hugging Face model in HFLM
-    eval_model = HFLM(
-        pretrained=hf_model,  # Pass the initialized model
+    
+    # Wrap in HFLM
+    model = HFLM(
+        pretrained=hf_model,
         device=args.device,
         batch_size=args.batch_size,
         max_batch_size=args.max_batch_size,
         dtype=args.dtype
     )
+    
+    logger.info("Custom model loaded successfully")
+    return model
 
-    # 4. Run evaluation with the wrapped model
-    results = evaluator.simple_evaluate(
-        model=eval_model,  # Pass the wrapped model
-        tasks=args.tasks,
-        num_fewshot=args.num_fewshot,
-        limit=args.limit,
-        write_out=args.write_out,
-        verbosity=args.verbosity,
-    )
 
-    return results
+def validate_model_config(model_path: str, trust_remote_code: bool = False) -> Dict[str, Any]:
+    """
+    Validate model configuration and return key information.
+    
+    Args:
+        model_path: Path to the model
+        trust_remote_code: Whether to trust remote code
+        
+    Returns:
+        Dict containing model configuration information
+    """
+    try:
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote_code)
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code)
+        
+        model_info = {
+            "model_type": getattr(config, "model_type", "unknown"),
+            "vocab_size": getattr(config, "vocab_size", "unknown"),
+            "hidden_size": getattr(config, "hidden_size", "unknown"),
+            "num_layers": getattr(config, "num_hidden_layers", "unknown"),
+            "num_experts": getattr(config, "num_experts", "not specified"),
+            "routing_type": getattr(config, "routing_type", "default"),
+        }
+        
+        logger.info("Model validation successful:")
+        for key, value in model_info.items():
+            logger.info(f"  {key}: {value}")
+            
+        return model_info
+        
+    except Exception as e:
+        logger.error(f"Model validation failed: {e}")
+        raise
 
-import numpy as np
-import torch
 
-def make_serializable(obj):
+def make_serializable(obj: Any) -> Any:
+    """
+    Convert objects to JSON-serializable format.
+    
+    Args:
+        obj: Object to convert
+        
+    Returns:
+        JSON-serializable version of the object
+    """
     if isinstance(obj, dict):
         return {k: make_serializable(v) for k, v in obj.items()}
     elif isinstance(obj, list):
         return [make_serializable(v) for v in obj]
     elif isinstance(obj, tuple):
         return tuple(make_serializable(v) for v in obj)
-    # NumPy scalars
     elif isinstance(obj, (np.integer, np.floating)):
         return obj.item()
-    # NumPy dtypes
     elif isinstance(obj, np.dtype):
         return str(obj)
-    # PyTorch tensor → list
     elif isinstance(obj, torch.Tensor):
         return obj.tolist()
-    # PyTorch dtype (e.g. torch.float16)
     elif isinstance(obj, torch.dtype):
         return str(obj)
-    # Anything else leave alone
     else:
         return obj
 
-def save_results(results: Dict, args) -> str:
-    """Save evaluation results to file, after converting to JSON-safe types, and print them."""
-    os.makedirs(args.output_dir, exist_ok=True)
 
-    # build filename exactly as before…
+def run_evaluation(args) -> Dict[str, Any]:
+    """
+    Run evaluation on the specified model.
+    
+    Args:
+        args: Parsed command line arguments
+        
+    Returns:
+        Dict containing evaluation results
+    """
+    logger.info("Starting evaluation...")
+    
+    # Validate model first
+    validate_model_config(args.model_path, args.trust_remote_code)
+    
+    # Load appropriate model
+    if args.model_type == "transformers":
+        model = load_transformers_model(args)
+    elif args.model_type == "custom":
+        model = load_custom_model(args)
+    else:
+        raise ValueError(f"Unknown model type: {args.model_type}")
+    
+    # Run evaluation
+    logger.info(f"Running evaluation on tasks: {args.tasks}")
+    logger.info(f"Few-shot examples: {args.num_fewshot}")
+    logger.info(f"Batch size: {args.batch_size}")
+    
+    results = evaluator.simple_evaluate(
+        model=model,
+        tasks=args.tasks,
+        num_fewshot=args.num_fewshot,
+        limit=args.limit,
+        write_out=args.write_out,
+        verbosity=args.verbosity,
+    )
+    
+    logger.info("Evaluation completed successfully")
+    return results
+
+
+def save_results(results: Dict[str, Any], args) -> str:
+    """
+    Save evaluation results to file.
+    
+    Args:
+        results: Evaluation results
+        args: Parsed command line arguments
+        
+    Returns:
+        str: Path to saved results file
+    """
+    os.makedirs(args.output_dir, exist_ok=True)
+    
+    # Generate filename if not provided
     if args.output_filename is None:
         model_name = os.path.basename(args.model_path.rstrip('/'))
         tasks_str = "_".join(args.tasks[:3])
         if len(args.tasks) > 3:
             tasks_str += f"_and_{len(args.tasks)-3}_more"
-        filename = f"{model_name}_{args.routing_type}_{tasks_str}_results.json"
+        
+        if args.model_type == "custom":
+            filename = f"{model_name}_{args.routing_type}_{tasks_str}_results.json"
+        else:
+            filename = f"{model_name}_transformers_{tasks_str}_results.json"
     else:
         filename = args.output_filename
+        
     if not filename.endswith('.json'):
         filename += '.json'
+    
     output_path = os.path.join(args.output_dir, filename)
-
+    
+    # Prepare metadata
     metadata = {
         "model_path": args.model_path,
-        "routing_type": args.routing_type,
-        "router_temperature": args.router_temperature,
-        "num_experts_per_tok": args.num_experts_per_tok,
+        "model_type": args.model_type,
         "tasks": args.tasks,
         "num_fewshot": args.num_fewshot,
         "batch_size": args.batch_size,
         "device": args.device,
         "dtype": args.dtype,
+        "limit": args.limit,
     }
+    
+    # Add routing-specific metadata for custom models
+    if args.model_type == "custom":
+        metadata.update({
+            "routing_type": args.routing_type,
+            "router_temperature": args.router_temperature,
+            "num_experts_per_tok": args.num_experts_per_tok,
+        })
+    
     results_with_metadata = {
         "metadata": metadata,
         "results": results
     }
-
-    # convert everything
-    serializable = make_serializable(results_with_metadata)
-
-    # write to disk
+    
+    # Convert to JSON-serializable format
+    serializable_results = make_serializable(results_with_metadata)
+    
+    # Save to file
     with open(output_path, 'w') as f:
-        json.dump(serializable, f, indent=2)
-
+        json.dump(serializable_results, f, indent=2)
+    
     logger.info(f"Results saved to {output_path}")
     return output_path
 
 
-
-
-def print_summary(results: Dict, routing_type: str) -> None:
+def print_summary(results: Dict[str, Any], args) -> None:
     """
-    Print a summary of evaluation results.
+    Print a formatted summary of evaluation results.
+    
+    Args:
+        results: Evaluation results
+        args: Parsed command line arguments
     """
-    print(f"\n{'='*60}")
-    print(f"EVALUATION SUMMARY - Routing: {routing_type.upper()}")
-    print(f"{'='*60}")
+    print(f"\n{'='*80}")
+    print(f"EVALUATION SUMMARY")
+    print(f"Model: {args.model_path}")
+    print(f"Type: {args.model_type.upper()}")
+    if args.model_type == "custom":
+        print(f"Routing: {args.routing_type.upper()}")
+    print(f"Tasks: {', '.join(args.tasks)}")
+    print(f"{'='*80}")
     
     if "results" in results:
         for task, metrics in results["results"].items():
             if isinstance(metrics, dict):
-                print(f"\n{task.upper()}:")
+                print(f"\n📊 {task.upper()}:")
                 for metric, value in metrics.items():
-                    if isinstance(value, (int, float)):
-                        if metric.endswith('_stderr'):
-                            continue  # Skip stderr for summary
+                    if isinstance(value, (int, float)) and not metric.endswith('_stderr'):
                         stderr_key = f"{metric}_stderr"
                         stderr = metrics.get(stderr_key, 0)
-                        print(f"  {metric}: {value:.4f} (±{stderr:.4f})")
-                    else:
-                        print(f"  {metric}: {value}")
+                        print(f"   {metric:.<20} {value:.4f} (±{stderr:.4f})")
+    else:
+        print("\n⚠️  No results found in evaluation output")
     
-    print(f"\n{'='*60}")
+    print(f"\n{'='*80}")
 
 
 def main():
@@ -279,6 +497,10 @@ def main():
         logger.setLevel(numeric_level)
     
     try:
+        logger.info("="*80)
+        logger.info("Starting OLMoE Model Evaluation")
+        logger.info("="*80)
+        
         # Run evaluation
         results = run_evaluation(args)
         
@@ -286,13 +508,18 @@ def main():
         output_path = save_results(results, args)
         
         # Print summary
-        print_summary(results, args.routing_type)
+        print_summary(results, args)
         
-        logger.info("Evaluation completed successfully!")
+        logger.info(f"✅ Evaluation completed successfully!")
+        logger.info(f"📁 Results saved to: {output_path}")
         
+    except KeyboardInterrupt:
+        logger.info("Evaluation interrupted by user")
+        sys.exit(1)
     except Exception as e:
-        logger.error(f"Evaluation failed: {e}")
-        raise
+        logger.error(f"❌ Evaluation failed: {e}")
+        logger.debug("Full traceback:", exc_info=True)
+        sys.exit(1)
 
 
 if __name__ == "__main__":