#!/usr/bin/env python """ Evaluation script for reasoning gym datasets. This script evaluates LLM performance on reasoning gym datasets using the OpenRouter API. Usage: python eval.py --config config.yaml [options] Options: --model MODEL Override model specified in config --output-dir DIR Override output directory specified in config --category CATEGORY Evaluate only datasets from this category --max-concurrent NUM Maximum number of concurrent API calls --n NUM Number of completions to generate per prompt (default: 1, each completion is a separate API call) --base-url URL API base URL (default: https://openrouter.ai/api/v1) --save-metadata Save entry metadata in results --full-results Save the full results file --verbose Print detailed model responses --debug Enable debug logging --resume DIR Resume evaluation from the specified directory Environment variables: OPENROUTER_API_KEY Required API key for OpenRouter """ import argparse import asyncio import json import logging import os import subprocess import sys from datetime import datetime from pathlib import Path from typing import Any, Optional from eval_config import CategoryConfig, DatasetConfig, EvalConfig from openai import AsyncOpenAI from tqdm.asyncio import tqdm_asyncio import reasoning_gym from reasoning_gym.utils import extract_answer class CheckpointManager: """Manages checkpoints for resumable evaluation.""" def __init__(self, output_dir: Path): """Initialize the checkpoint manager. Args: output_dir: Directory where checkpoints and results are stored """ self.output_dir = output_dir self.checkpoint_path = output_dir / "checkpoint.json" self.completed_datasets = set() self.previous_category_results = {} # Store previously completed category results self.load_checkpoint() def load_checkpoint(self) -> None: """Load existing checkpoint and previous results if available.""" # Load checkpoint file if self.checkpoint_path.exists(): with open(self.checkpoint_path, "r") as f: checkpoint_data = json.load(f) self.completed_datasets = set(checkpoint_data.get("completed_datasets", [])) # Load previous category results for category_dir in self.output_dir.iterdir(): if category_dir.is_dir(): category_name = category_dir.name self.previous_category_results[category_name] = [] # Load each dataset result file in this category for dataset_file in category_dir.glob("*.json"): try: with open(dataset_file, "r") as f: dataset_result = json.load(f) self.previous_category_results[category_name].append(dataset_result) except Exception as e: logging.warning(f"Error loading previous result {dataset_file}: {str(e)}") def is_dataset_completed(self, category: str, dataset: str) -> bool: """Check if a dataset has been completed. Args: category: Category name dataset: Dataset name Returns: True if the dataset has been completed, False otherwise """ return f"{category}/{dataset}" in self.completed_datasets def mark_dataset_completed(self, category: str, dataset: str) -> None: """Mark a dataset as completed and update checkpoint file. Args: category: Category name dataset: Dataset name """ self.completed_datasets.add(f"{category}/{dataset}") self._save_checkpoint() def get_dataset_result(self, category: str, dataset: str) -> Optional[dict[str, Any]]: """Get previously completed dataset result if available. Args: category: Category name dataset: Dataset name Returns: Dataset result dict if found, None otherwise """ # Try to find the dataset in previously loaded results first if category in self.previous_category_results: for dataset_result in self.previous_category_results[category]: if dataset_result["name"] == dataset: return dataset_result # If not found in memory, try to load from file dataset_path = self.output_dir / category / f"{dataset}.json" if dataset_path.exists(): try: with open(dataset_path, "r") as f: return json.load(f) except Exception as e: logging.error(f"Error loading dataset result from {dataset_path}: {str(e)}") return None def _save_checkpoint(self) -> None: """Save checkpoint to disk.""" checkpoint_data = { "completed_datasets": list(self.completed_datasets), "last_updated": datetime.now().isoformat(), } with open(self.checkpoint_path, "w") as f: json.dump(checkpoint_data, f, indent=2) # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", handlers=[logging.StreamHandler()], ) # httpx logging will be configured in the AsyncModelEvaluator class # based on the debug flag def get_git_hash() -> str: """Get current git hash for reproducibility.""" cmd = ["git", "rev-parse", "HEAD"] try: return subprocess.check_output(cmd, text=True, stderr=subprocess.PIPE).strip() except Exception: return "unknown" class AsyncModelEvaluator: """Evaluates models on reasoning datasets with async API calls via OpenRouter.""" def __init__( self, config: EvalConfig, api_key: Optional[str] = None, base_url: str = "https://openrouter.ai/api/v1", verbose: bool = False, debug: bool = False, ): """Initialize the evaluator with configuration. Args: config: Evaluation configuration api_key: API key for the service (optional for some APIs) base_url: API base URL verbose: Whether to print detailed model responses debug: Whether to enable debug logging """ self.config = config self.base_url = base_url self.verbose = verbose self.debug = debug # Set up logging self.logger = logging.getLogger("AsyncModelEvaluator") if debug: self.logger.setLevel(logging.DEBUG) # Enable httpx logs in debug mode logging.getLogger("httpx").setLevel(logging.INFO) else: # Suppress httpx logs in normal mode logging.getLogger("httpx").setLevel(logging.WARNING) # Set up API client self.client = AsyncOpenAI(base_url=self.base_url, api_key=api_key) # Concurrency control self.semaphore = asyncio.Semaphore(config.max_concurrent) # Metadata self.git_hash = get_git_hash() self.start_time = datetime.now() # Checkpoint and resume related attributes self.resume_dir = None self.output_dir = None self.checkpoint_manager = None def create_output_dir(self) -> Path: """Create output directory or use existing one for resuming. Returns: Path to the output directory """ # Check if we're resuming from a previous run if self.resume_dir: output_dir = Path(self.resume_dir) if not output_dir.exists(): raise ValueError(f"Resume directory {output_dir} does not exist") self.logger.info(f"Resuming evaluation from {output_dir}") return output_dir # Create new output directory timestamp = self.start_time.strftime("%Y%m%d_%H%M%S") model_name = self.config.model.replace("/", "_") if len(self.config.categories) == 1: # Include category name in the output directory when evaluating a single category category_name = self.config.categories[0].category output_dir = Path(self.config.output_dir) / f"{model_name}_{category_name}_{timestamp}" else: # Original format for multiple categories output_dir = Path(self.config.output_dir) / f"{model_name}_{timestamp}" output_dir.mkdir(parents=True, exist_ok=True) return output_dir def _save_dataset_results(self, category_name: str, dataset_name: str, results: dict[str, Any]) -> None: """Save individual dataset results to file. Args: category_name: Category name dataset_name: Dataset name results: Dataset evaluation results """ category_dir = self.output_dir / category_name category_dir.mkdir(exist_ok=True) dataset_path = category_dir / f"{dataset_name}.json" with open(dataset_path, "w") as f: json.dump(results, f, indent=2) def _update_partial_summary(self, category_results: list[dict[str, Any]]) -> None: """Update partial summary after each category completes. Args: category_results: List of category results completed so far """ partial_results = { "metadata": { "timestamp": self.start_time.isoformat(), "model": self.config.model, "provider": self.config.provider, "git_hash": self.git_hash, "duration_seconds": (datetime.now() - self.start_time).total_seconds(), "max_tokens": self.config.max_tokens, "temperature": self.config.temperature, "top_p": self.config.top_p, "partial": True, }, "categories": category_results, } # Generate partial summary partial_results["summary"] = self.generate_summary(partial_results) # Save partial summary summary_path = self.output_dir / "summary.json" with open(summary_path, "w") as f: json.dump(partial_results["summary"], f, indent=2) async def get_single_response(self, prompt: str) -> str: """Get a single response from model with retry logic via OpenRouter. Args: prompt: The prompt to send to the model Returns: The model's response text Raises: Exception: If all retries fail """ max_retries = 10 base_delay = 1.0 max_delay = 60.0 backoff_factor = 2.0 for attempt in range(max_retries): try: async with self.semaphore: # Prepare API call parameters params = { "model": self.config.model, "messages": [ {"role": self.config.system_role, "content": self.config.get_system_prompt()}, {"role": "user", "content": prompt}, ], } # Add sampling parameters if specified if self.config.max_tokens is not None: params["max_tokens"] = self.config.max_tokens if self.config.temperature is not None: params["temperature"] = self.config.temperature if self.config.top_p is not None: params["top_p"] = self.config.top_p # Add provider configuration if specified if self.config.provider: params["extra_body"] = {"provider": {"order": [self.config.provider], "allow_fallbacks": False}} completion = await self.client.chat.completions.create(**params) response = completion.choices[0].message.content if self.verbose: self.logger.info(f"Response: {response}") return response except Exception as e: delay = min(max_delay, base_delay * (backoff_factor**attempt)) self.logger.warning(f"Attempt {attempt+1}/{max_retries} failed: {str(e)}") self.logger.warning(f"Retrying in {delay:.2f} seconds...") await asyncio.sleep(delay) raise Exception(f"Failed to get model response after {max_retries} attempts") async def get_model_response(self, prompt: str) -> list[str]: """Get multiple responses from model by making multiple API calls. Args: prompt: The prompt to send to the model Returns: A list of model response texts Raises: Exception: If all attempts fail """ if self.verbose: self.logger.info(f"Prompt: {prompt}") self.logger.info(f"Generating {self.config.completions_per_prompt} completions...") # Create tasks for multiple completions tasks = [] for i in range(self.config.completions_per_prompt): tasks.append(self.get_single_response(prompt)) # Execute all tasks concurrently responses = await asyncio.gather(*tasks, return_exceptions=True) # Handle any exceptions valid_responses = [] for i, response in enumerate(responses): if isinstance(response, Exception): self.logger.error(f"Completion {i+1} failed: {str(response)}") else: valid_responses.append(response) if self.verbose: self.logger.info(f"Response {len(valid_responses)}: {response}") if not valid_responses: raise Exception("All completion attempts failed") return valid_responses async def process_entry( self, dataset: reasoning_gym.dataset.ProceduralDataset, entry: dict[str, Any], entry_index: int, dataset_name: str, ) -> dict[str, Any]: """Process a single dataset entry. Args: dataset: The dataset instance entry: The entry to process entry_index: Index of the entry in the dataset dataset_name: Name of the dataset Returns: Dict with processing results """ responses = None completion_results = [] best_score = 0.0 total_score = 0.0 best_answer = None best_response = None try: # Get multiple model responses responses = await self.get_model_response(entry["question"]) # Count total completions for mean score calculation total_completions = len(responses) for i, response in enumerate(responses): try: # Try to extract answer and score it model_answer = extract_answer(response) score = dataset.score_answer(answer=model_answer, entry=entry) completion_result = { "model_answer": model_answer, "full_model_response": response, "score": score, } # Track scores if score > best_score: best_score = score best_answer = model_answer best_response = response # If we don't have a best answer yet, use the first non-None answer elif best_answer is None and model_answer is not None: best_answer = model_answer best_response = response best_score = score total_score += score completion_results.append(completion_result) if self.verbose: print(f"Question: {entry['question']}") print(f"Expected: {entry['answer']}") print(f"Completion {i+1} Answer: {model_answer}") print(f"Completion {i+1} Score: {score}") print("-" * 40) except Exception as e: self.logger.error(f"Error processing completion {i+1}: {str(e)}") # Add failed completion with score 0.0 (already counted in total_completions) completion_results.append( { "model_answer": "ERROR", "full_model_response": response, "score": 0.0, "error": str(e), } ) # If we have no valid completions, log a warning instead of raising an exception if not best_answer: self.logger.warning( f"Failed to extract a valid answer from model responses for dataset '{dataset_name}', entry index {entry_index}" ) # Use None instead of empty string as the best answer best_answer = None best_response = responses[0] if responses and len(responses) > 0 else None best_score = 0.0 # Calculate mean score - count all completions including failures mean_score = total_score / total_completions if total_completions > 0 else 0.0 result = { "question": entry["question"], "expected_answer": str(entry["answer"]), "best_model_answer": best_answer, "best_full_model_response": best_response, "best_score": best_score, "mean_score": mean_score, "completions": completion_results, } # Only include metadata if configured to do so if self.config.save_metadata: result["metadata"] = entry["metadata"] return result except Exception as e: self.logger.error(f"Error processing entry: {str(e)}") result = { "question": entry["question"], "expected_answer": str(entry["answer"]), "best_model_answer": None, # First check if we already have a best_response from partial processing # If not, then fall back to the first response or None "best_full_model_response": ( best_response if best_response is not None else (responses[0] if responses and len(responses) > 0 else None) ), "best_score": best_score if best_score > 0 else 0.0, "mean_score": total_score / total_completions if total_completions > 0 else 0.0, "error": str(e), "completions": completion_results if "completion_results" in locals() else [], } # Only include metadata if configured to do so if self.config.save_metadata: result["metadata"] = entry["metadata"] return result async def evaluate_dataset(self, category_name: str, dataset_config: DatasetConfig) -> dict[str, Any]: """Evaluate a single dataset. Args: category_name: Name of the category dataset_config: Configuration for the dataset Returns: Dict with evaluation results """ dataset_name = dataset_config.dataset # Check if this dataset has already been completed if self.checkpoint_manager.is_dataset_completed(category_name, dataset_name): # Get the dataset result from checkpoint manager dataset_result = self.checkpoint_manager.get_dataset_result(category_name, dataset_name) if dataset_result: self.logger.info(f"Skipping already completed dataset: {dataset_name}") return dataset_result # If we can't load the result, we'll need to re-evaluate the dataset self.logger.info(f"Re-evaluating dataset: {dataset_name}") # Remove from completed datasets so it will be processed self.checkpoint_manager.completed_datasets.discard(f"{category_name}/{dataset_name}") self.logger.info(f"Evaluating dataset: {dataset_name}") try: # Create dataset with all parameters dataset_params = {} # Add all parameters from the config params dictionary # Make sure we don't have a nested 'params' dictionary for k, v in dataset_config.params.items(): if k != "params": dataset_params[k] = v elif isinstance(v, dict): # If there's a nested params dict, flatten it dataset_params.update(v) # Add size and seed if they're not None if dataset_config.size is not None: dataset_params["size"] = dataset_config.size if dataset_config.seed is not None: dataset_params["seed"] = dataset_config.seed dataset = reasoning_gym.create_dataset(dataset_name, **dataset_params) # Get all entries all_entries = list(dataset) # Process entries with progress bar, passing the entry index and dataset name tasks = [self.process_entry(dataset, entry, idx, dataset_name) for idx, entry in enumerate(all_entries)] results = await tqdm_asyncio.gather(*tasks, desc=f"Processing {dataset_name}", leave=True) # Calculate metrics total_best_score = sum(r["best_score"] for r in results) total_mean_score = sum(r["mean_score"] for r in results) average_best_score = total_best_score / len(results) if results else 0 average_mean_score = total_mean_score / len(results) if results else 0 dataset_results = { "name": dataset_name, "category": category_name, "average_best_score": average_best_score, "average_mean_score": average_mean_score, "total_examples": len(results), "config": {"size": dataset_config.size, "seed": dataset_config.seed, **dataset_config.params}, "system_prompt": self.config.get_system_prompt(), "completions_per_prompt": self.config.completions_per_prompt, "results": results, } # Mark dataset as completed and save results self.checkpoint_manager.mark_dataset_completed(category_name, dataset_name) self._save_dataset_results(category_name, dataset_name, dataset_results) return dataset_results except Exception as e: self.logger.error(f"Error evaluating dataset {dataset_name}: {str(e)}") return { "name": dataset_name, "category": category_name, "average_best_score": 0.0, "average_mean_score": 0.0, "total_examples": 0, "config": {"size": dataset_config.size, "seed": dataset_config.seed, **dataset_config.params}, "system_prompt": self.config.get_system_prompt(), "error": str(e), "results": [], } async def evaluate_category(self, category_config: CategoryConfig) -> dict[str, Any]: """Evaluate all datasets in a category. Args: category_config: Configuration for the category Returns: Dict with category evaluation results """ category_name = category_config.category self.logger.info(f"Evaluating category: {category_name}") # Check if all datasets in this category are already completed all_completed = True for dataset_config in category_config.datasets: if not self.checkpoint_manager.is_dataset_completed(category_name, dataset_config.dataset): all_completed = False break # If all datasets are completed and we have previous results, use them if all_completed and category_name in self.checkpoint_manager.previous_category_results: self.logger.info(f"Using previously completed results for category: {category_name}") return { "name": category_name, "datasets": self.checkpoint_manager.previous_category_results[category_name], } # Process datasets sequentially to ensure proper checkpointing dataset_results = [] for dataset_config in category_config.datasets: result = await self.evaluate_dataset(category_name, dataset_config) dataset_results.append(result) return { "name": category_name, "datasets": dataset_results, } async def evaluate_all(self) -> dict[str, Any]: """Evaluate all categories and datasets, resuming from checkpoint if available. Returns: Dict with all evaluation results and summary """ self.logger.info(f"Starting evaluation of {len(self.config.categories)} categories") # Initialize output directory and checkpoint manager self.output_dir = self.create_output_dir() self.checkpoint_manager = CheckpointManager(self.output_dir) # Process each category sequentially to ensure proper checkpointing category_results = [] for category in self.config.categories: category_result = await self.evaluate_category(category) category_results.append(category_result) # Update partial summary after each category self._update_partial_summary(category_results) # Generate results structure results = { "metadata": { "timestamp": self.start_time.isoformat(), "model": self.config.model, "provider": self.config.provider, "git_hash": self.git_hash, "duration_seconds": (datetime.now() - self.start_time).total_seconds(), "max_tokens": self.config.max_tokens, "temperature": self.config.temperature, "top_p": self.config.top_p, "partial": False, # Mark as complete }, "categories": category_results, } # Generate summary results["summary"] = self.generate_summary(results) return results def generate_summary(self, results: dict[str, Any]) -> dict[str, Any]: """Generate a summary of evaluation results in the original configuration order. Args: results: The full evaluation results Returns: Dict with summary information """ summary = { "total_datasets": 0, "total_examples": 0, "dataset_best_scores": {}, "dataset_mean_scores": {}, } # Iterate through categories and datasets in the original order from config for category_config in self.config.categories: for dataset_config in category_config.datasets: dataset_name = dataset_config.dataset dataset_found = False # Find corresponding results for category in results["categories"]: if category["name"] == category_config.category: for dataset in category["datasets"]: if dataset["name"] == dataset_name: # Add to summary in original order summary["dataset_best_scores"][dataset_name] = dataset["average_best_score"] summary["dataset_mean_scores"][dataset_name] = dataset["average_mean_score"] summary["total_datasets"] += 1 summary["total_examples"] += dataset["total_examples"] dataset_found = True break # If dataset wasn't found in results (error), add with score 0 if not dataset_found: summary["dataset_best_scores"][dataset_name] = 0.0 summary["dataset_mean_scores"][dataset_name] = 0.0 summary["total_datasets"] += 1 return summary def save_results(self, results: dict[str, Any]) -> tuple[str, str]: """Save evaluation results to files. Args: results: The evaluation results to save Returns: Tuple of (results_path, summary_path) """ # Output directory is already created during evaluation results_path = None # Save full results if configured to do so if self.config.save_full_results: results_path = self.output_dir / "results.json" with open(results_path, "w") as f: json.dump(results, f, indent=2) # Add timestamp, git hash, model, provider, sampling parameters, and duration to summary summary_data = results["summary"].copy() summary_data["timestamp"] = self.start_time.isoformat() summary_data["git_hash"] = self.git_hash summary_data["model"] = self.config.model summary_data["provider"] = self.config.provider summary_data["system_prompt"] = self.config.get_system_prompt() if self.config.system_prompt_id: summary_data["system_prompt_id"] = self.config.system_prompt_id summary_data["max_tokens"] = self.config.max_tokens summary_data["temperature"] = self.config.temperature summary_data["top_p"] = self.config.top_p summary_data["completions_per_prompt"] = self.config.completions_per_prompt summary_data["duration_seconds"] = results["metadata"]["duration_seconds"] summary_data["partial"] = False # Mark as complete # Save summary summary_path = self.output_dir / "summary.json" with open(summary_path, "w") as f: json.dump(summary_data, f, indent=2) # Individual dataset results are already saved during evaluation return str(results_path) if results_path else None, str(summary_path) def print_summary(self, results: dict[str, Any]) -> None: """Print a summary of evaluation results to the console. Args: results: The evaluation results """ summary = results["summary"] print("\nEvaluation Summary:") print("------------------") print(f"Model: {self.config.model}") print(f"Provider: {self.config.provider}") system_prompt = self.config.get_system_prompt() print(f"System Prompt: {system_prompt[:50]}..." if len(system_prompt) > 50 else system_prompt) print(f"Max Tokens: {self.config.max_tokens}") print(f"Temperature: {self.config.temperature}") print(f"Top-p: {self.config.top_p}") print(f"Completions per prompt: {self.config.completions_per_prompt}") print(f"Git Hash: {self.git_hash}") print(f"Duration: {results['metadata']['duration_seconds']:.2f} seconds") print() print("Dataset Scores (in configuration order):") print(" Dataset Name Best Score Mean Score Examples") print(" ------------------------------------------------------------------") for dataset_name in summary["dataset_best_scores"].keys(): best_score = summary["dataset_best_scores"][dataset_name] mean_score = summary["dataset_mean_scores"][dataset_name] # Find the number of examples for this dataset examples = 0 for category in results["categories"]: for dataset in category["datasets"]: if dataset["name"] == dataset_name: examples = dataset["total_examples"] break # Use fixed-width formatting for better alignment print(f" {dataset_name:<30} {best_score:>8.1%} {mean_score:>8.1%} {examples:>8}") print() print(f"Total datasets: {summary['total_datasets']}") print(f"Total examples: {summary['total_examples']}") async def main_async(): """Main async function.""" parser = argparse.ArgumentParser(description="Evaluate models on reasoning datasets") parser.add_argument("--config", required=True, help="Path to configuration file (YAML or JSON)") parser.add_argument("--model", help="Override model specified in config") parser.add_argument("--output-dir", help="Override output directory specified in config") parser.add_argument("--category", help="Evaluate only datasets from this category") parser.add_argument("--max-concurrent", type=int, help="Maximum number of concurrent API calls") parser.add_argument("--n", type=int, default=1, help="Number of completions to generate per prompt") parser.add_argument("--base-url", default="https://openrouter.ai/api/v1", help="API base URL") parser.add_argument( "--api-key", help="API key for the service (optional for some APIs, defaults to OPENROUTER_API_KEY env var for OpenRouter URLs)", ) parser.add_argument("--save-metadata", action="store_true", help="Save entry metadata in results") parser.add_argument("--full-results", action="store_true", help="Save the full results file") parser.add_argument("--verbose", action="store_true", help="Print detailed model responses") parser.add_argument("--debug", action="store_true", help="Enable debug logging") parser.add_argument("--resume", help="Resume evaluation from the specified directory") args = parser.parse_args() # Get API key from command line or environment variable api_key = args.api_key if api_key is None: # If base_url is OpenRouter, try to get API key from environment if args.base_url.startswith("https://openrouter.ai/api"): api_key = os.getenv("OPENROUTER_API_KEY") if not api_key: print("Warning: OPENROUTER_API_KEY environment variable is not set") print("Please set it using: export OPENROUTER_API_KEY=your-api-key") print("Or provide it directly with --api-key") return 1 # Load configuration config_path = args.config if config_path.endswith(".yaml") or config_path.endswith(".yml"): config = EvalConfig.from_yaml(config_path) elif config_path.endswith(".json"): config = EvalConfig.from_json(config_path) else: print("Error: Configuration file must be YAML or JSON") return 1 # Apply command line overrides if args.model: config.model = args.model if args.output_dir: config.output_dir = args.output_dir if args.max_concurrent: config.max_concurrent = args.max_concurrent if args.n: config.completions_per_prompt = args.n if args.save_metadata: config.save_metadata = True if args.full_results: config.save_full_results = True # Filter categories if --category is specified if args.category: # Keep only the specified category filtered_categories = [cat for cat in config.categories if cat.category == args.category] if not filtered_categories: print(f"Error: Category '{args.category}' not found in configuration") return 1 config.categories = filtered_categories # Create evaluator evaluator = AsyncModelEvaluator( config=config, api_key=api_key, base_url=args.base_url, verbose=args.verbose, debug=args.debug ) # Set resume directory if specified if args.resume: evaluator.resume_dir = args.resume # Run evaluation try: results = await evaluator.evaluate_all() # Save and print results results_path, summary_path = evaluator.save_results(results) evaluator.print_summary(results) if results_path: print(f"\nResults saved to: {results_path}") print(f"Summary saved to: {summary_path}") return 0 except Exception as e: print(f"Error during evaluation: {str(e)}") if args.debug: import traceback traceback.print_exc() return 1 def main(): """Entry point.""" exit_code = asyncio.run(main_async()) sys.exit(exit_code) if __name__ == "__main__": main()