Create inference/evaluate.py

inference/evaluate.py

@@ -0,0 +1,401 @@
+#!/usr/bin/env python3
+"""
+Helion-2.5-Rnd Evaluation Script
+Comprehensive benchmark evaluation across multiple datasets
+"""
+
+import argparse
+import json
+import logging
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class HelionEvaluator:
+    """Evaluation framework for Helion model"""
+    
+    def __init__(
+        self,
+        model_path: str,
+        device: str = "cuda",
+        batch_size: int = 1,
+        max_length: int = 2048
+    ):
+        """
+        Initialize evaluator
+        
+        Args:
+            model_path: Path to model or HuggingFace model ID
+            device: Device to run evaluation on
+            batch_size: Batch size for evaluation
+            max_length: Maximum sequence length
+        """
+        logger.info(f"Loading model from {model_path}")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        self.device = device
+        self.batch_size = batch_size
+        self.max_length = max_length
+        
+        logger.info("Model loaded successfully")
+    
+    def generate(
+        self,
+        prompt: str,
+        max_new_tokens: int = 512,
+        temperature: float = 0.0,
+        **kwargs
+    ) -> str:
+        """Generate text from prompt"""
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=self.max_length
+        ).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature if temperature > 0 else 1.0,
+                do_sample=temperature > 0,
+                pad_token_id=self.tokenizer.pad_token_id,
+                **kwargs
+            )
+        
+        response = self.tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:],
+            skip_special_tokens=True
+        )
+        
+        return response.strip()
+    
+    def evaluate_mmlu(self, num_samples: Optional[int] = None) -> Dict:
+        """Evaluate on MMLU benchmark"""
+        logger.info("Evaluating on MMLU...")
+        
+        dataset = load_dataset("cais/mmlu", "all", split="test")
+        if num_samples:
+            dataset = dataset.select(range(min(num_samples, len(dataset))))
+        
+        correct = 0
+        total = 0
+        
+        for example in tqdm(dataset, desc="MMLU"):
+            question = example["question"]
+            choices = example["choices"]
+            answer = example["answer"]
+            
+            # Format prompt
+            prompt = f"Question: {question}\n\nChoices:\n"
+            for i, choice in enumerate(choices):
+                prompt += f"{chr(65+i)}. {choice}\n"
+            prompt += "\nAnswer: "
+            
+            # Generate response
+            response = self.generate(prompt, max_new_tokens=10, temperature=0.0)
+            
+            # Extract answer
+            pred = response.strip()[0].upper() if response else ""
+            correct_answer = chr(65 + answer)
+            
+            if pred == correct_answer:
+                correct += 1
+            total += 1
+        
+        accuracy = correct / total if total > 0 else 0
+        
+        return {
+            "benchmark": "MMLU",
+            "accuracy": accuracy,
+            "correct": correct,
+            "total": total
+        }
+    
+    def evaluate_gsm8k(self, num_samples: Optional[int] = None) -> Dict:
+        """Evaluate on GSM8K mathematical reasoning"""
+        logger.info("Evaluating on GSM8K...")
+        
+        dataset = load_dataset("gsm8k", "main", split="test")
+        if num_samples:
+            dataset = dataset.select(range(min(num_samples, len(dataset))))
+        
+        correct = 0
+        total = 0
+        
+        for example in tqdm(dataset, desc="GSM8K"):
+            question = example["question"]
+            answer = example["answer"]
+            
+            # Extract numerical answer
+            import re
+            match = re.search(r'####\s*(-?\d+(?:,\d+)*(?:\.\d+)?)', answer)
+            if not match:
+                continue
+            
+            correct_answer = match.group(1).replace(',', '')
+            
+            # Format prompt
+            prompt = f"Question: {question}\n\nLet's solve this step by step:\n"
+            
+            # Generate response
+            response = self.generate(prompt, max_new_tokens=512, temperature=0.0)
+            
+            # Extract predicted answer
+            pred_match = re.search(r'(?:answer is|=)\s*(-?\d+(?:,\d+)*(?:\.\d+)?)', response.lower())
+            if pred_match:
+                pred_answer = pred_match.group(1).replace(',', '')
+                if pred_answer == correct_answer:
+                    correct += 1
+            
+            total += 1
+        
+        accuracy = correct / total if total > 0 else 0
+        
+        return {
+            "benchmark": "GSM8K",
+            "accuracy": accuracy,
+            "correct": correct,
+            "total": total
+        }
+    
+    def evaluate_humaneval(self, num_samples: Optional[int] = None) -> Dict:
+        """Evaluate on HumanEval code generation"""
+        logger.info("Evaluating on HumanEval...")
+        
+        try:
+            dataset = load_dataset("openai_humaneval", split="test")
+        except:
+            logger.warning("HumanEval dataset not available")
+            return {"benchmark": "HumanEval", "error": "Dataset not available"}
+        
+        if num_samples:
+            dataset = dataset.select(range(min(num_samples, len(dataset))))
+        
+        results = []
+        
+        for example in tqdm(dataset, desc="HumanEval"):
+            prompt = example["prompt"]
+            
+            # Generate code
+            full_prompt = f"Complete the following Python function:\n\n{prompt}"
+            response = self.generate(
+                full_prompt,
+                max_new_tokens=512,
+                temperature=0.0
+            )
+            
+            # Extract code
+            code = prompt + response
+            
+            results.append({
+                "task_id": example["task_id"],
+                "completion": code,
+                "test": example["test"]
+            })
+        
+        # Note: Full evaluation requires executing code
+        # This is a simplified version
+        return {
+            "benchmark": "HumanEval",
+            "samples_generated": len(results),
+            "note": "Full evaluation requires code execution framework"
+        }
+    
+    def evaluate_truthfulqa(self, num_samples: Optional[int] = None) -> Dict:
+        """Evaluate on TruthfulQA"""
+        logger.info("Evaluating on TruthfulQA...")
+        
+        dataset = load_dataset("truthful_qa", "generation", split="validation")
+        if num_samples:
+            dataset = dataset.select(range(min(num_samples, len(dataset))))
+        
+        responses = []
+        
+        for example in tqdm(dataset, desc="TruthfulQA"):
+            question = example["question"]
+            
+            prompt = f"Question: {question}\n\nProvide a truthful and accurate answer:\nAnswer: "
+            
+            response = self.generate(prompt, max_new_tokens=256, temperature=0.0)
+            
+            responses.append({
+                "question": question,
+                "response": response,
+                "best_answer": example["best_answer"],
+                "correct_answers": example["correct_answers"],
+                "incorrect_answers": example["incorrect_answers"]
+            })
+        
+        return {
+            "benchmark": "TruthfulQA",
+            "samples_evaluated": len(responses),
+            "note": "Manual review required for truthfulness assessment"
+        }
+    
+    def evaluate_all(
+        self,
+        output_file: Optional[str] = None,
+        num_samples: Optional[int] = None
+    ) -> Dict:
+        """Run all evaluations"""
+        logger.info("Starting comprehensive evaluation...")
+        
+        results = {
+            "model": "DeepXR/Helion-2.5-Rnd",
+            "benchmarks": {}
+        }
+        
+        # Run evaluations
+        try:
+            results["benchmarks"]["mmlu"] = self.evaluate_mmlu(num_samples)
+        except Exception as e:
+            logger.error(f"MMLU evaluation failed: {e}")
+            results["benchmarks"]["mmlu"] = {"error": str(e)}
+        
+        try:
+            results["benchmarks"]["gsm8k"] = self.evaluate_gsm8k(num_samples)
+        except Exception as e:
+            logger.error(f"GSM8K evaluation failed: {e}")
+            results["benchmarks"]["gsm8k"] = {"error": str(e)}
+        
+        try:
+            results["benchmarks"]["humaneval"] = self.evaluate_humaneval(num_samples)
+        except Exception as e:
+            logger.error(f"HumanEval evaluation failed: {e}")
+            results["benchmarks"]["humaneval"] = {"error": str(e)}
+        
+        try:
+            results["benchmarks"]["truthfulqa"] = self.evaluate_truthfulqa(num_samples)
+        except Exception as e:
+            logger.error(f"TruthfulQA evaluation failed: {e}")
+            results["benchmarks"]["truthfulqa"] = {"error": str(e)}
+        
+        # Save results
+        if output_file:
+            output_path = Path(output_file)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            with open(output_path, 'w') as f:
+                json.dump(results, f, indent=2)
+            
+            logger.info(f"Results saved to {output_path}")
+        
+        # Print summary
+        logger.info("\n" + "="*50)
+        logger.info("EVALUATION SUMMARY")
+        logger.info("="*50)
+        
+        for benchmark, result in results["benchmarks"].items():
+            if "accuracy" in result:
+                logger.info(f"{benchmark.upper()}: {result['accuracy']:.2%}")
+            elif "error" in result:
+                logger.info(f"{benchmark.upper()}: ERROR - {result['error']}")
+            else:
+                logger.info(f"{benchmark.upper()}: {result.get('note', 'Completed')}")
+        
+        return results
+
+
+def main():
+    """Main evaluation entry point"""
+    parser = argparse.ArgumentParser(description="Evaluate Helion model")
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Model path or HuggingFace ID"
+    )
+    parser.add_argument(
+        "--benchmarks",
+        type=str,
+        nargs="+",
+        default=["all"],
+        choices=["all", "mmlu", "gsm8k", "humaneval", "truthfulqa"],
+        help="Benchmarks to run"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="evaluation_results.json",
+        help="Output file for results"
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=None,
+        help="Number of samples to evaluate (for quick testing)"
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device to use"
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Batch size"
+    )
+    
+    args = parser.parse_args()
+    
+    # Initialize evaluator
+    evaluator = HelionEvaluator(
+        model_path=args.model,
+        device=args.device,
+        batch_size=args.batch_size
+    )
+    
+    # Run evaluations
+    if "all" in args.benchmarks:
+        results = evaluator.evaluate_all(
+            output_file=args.output,
+            num_samples=args.num_samples
+        )
+    else:
+        results = {"model": args.model, "benchmarks": {}}
+        
+        if "mmlu" in args.benchmarks:
+            results["benchmarks"]["mmlu"] = evaluator.evaluate_mmlu(args.num_samples)
+        
+        if "gsm8k" in args.benchmarks:
+            results["benchmarks"]["gsm8k"] = evaluator.evaluate_gsm8k(args.num_samples)
+        
+        if "humaneval" in args.benchmarks:
+            results["benchmarks"]["humaneval"] = evaluator.evaluate_humaneval(args.num_samples)
+        
+        if "truthfulqa" in args.benchmarks:
+            results["benchmarks"]["truthfulqa"] = evaluator.evaluate_truthfulqa(args.num_samples)
+        
+        # Save results
+        with open(args.output, 'w') as f:
+            json.dump(results, f, indent=2)
+        
+        logger.info(f"Results saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()