DeepXR/Helion-OSC

1"""2Helion-OSC Comprehensive Benchmark Suite3Performance benchmarking and comparison with other models4"""5 6import torch7import time8import psutil9import numpy as np10from typing import Dict, List, Any, Optional11from dataclasses import dataclass, asdict12import json13import logging14from transformers import AutoTokenizer, AutoModelForCausalLM15from tqdm import tqdm16import pandas as pd17import matplotlib.pyplot as plt18import seaborn as sns19 20logging.basicConfig(level=logging.INFO)21logger = logging.getLogger(__name__)22 23 24@dataclass25class BenchmarkResult:26    """Single benchmark result"""27    model_name: str28    task: str29    prompt_length: int30    generation_length: int31    temperature: float32    inference_time: float33    tokens_per_second: float34    memory_used_mb: float35    gpu_memory_mb: Optional[float]36    success: bool37    error: Optional[str] = None38 39 40@dataclass41class AggregatedResults:42    """Aggregated benchmark results"""43    model_name: str44    total_tests: int45    successful_tests: int46    failed_tests: int47    avg_inference_time: float48    avg_tokens_per_second: float49    avg_memory_mb: float50    min_inference_time: float51    max_inference_time: float52    std_inference_time: float53 54 55class PerformanceBenchmark:56    """Performance benchmarking utilities"""57    58    def __init__(self, model_name: str = "DeepXR/Helion-OSC"):59        self.model_name = model_name60        self.device = "cuda" if torch.cuda.is_available() else "cpu"61        62        logger.info(f"Loading model: {model_name}")63        self.tokenizer = AutoTokenizer.from_pretrained(model_name)64        self.model = AutoModelForCausalLM.from_pretrained(65            model_name,66            torch_dtype=torch.bfloat16 if self.device == "cuda" else torch.float32,67            device_map="auto" if self.device == "cuda" else None68        )69        70        if self.device == "cpu":71            self.model = self.model.to(self.device)72        73        self.model.eval()74        self.results: List[BenchmarkResult] = []75    76    def get_memory_usage(self) -> tuple:77        """Get current memory usage"""78        process = psutil.Process()79        ram_mb = process.memory_info().rss / 1024 / 102480        81        gpu_mb = None82        if torch.cuda.is_available():83            gpu_mb = torch.cuda.memory_allocated() / 1024 / 102484        85        return ram_mb, gpu_mb86    87    def benchmark_inference(88        self,89        prompt: str,90        task: str,91        max_length: int = 512,92        temperature: float = 0.7,93        num_runs: int = 194    ) -> List[BenchmarkResult]:95        """Benchmark inference performance"""96        run_results = []97        98        for run in range(num_runs):99            try:100                # Tokenize101                inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)102                prompt_length = inputs.input_ids.shape[1]103                104                # Warm up GPU105                if run == 0 and self.device == "cuda":106                    with torch.no_grad():107                        _ = self.model.generate(**inputs, max_length=prompt_length + 10)108                    torch.cuda.synchronize()109                110                # Clear cache111                if self.device == "cuda":112                    torch.cuda.empty_cache()113                114                # Measure memory before115                ram_before, gpu_before = self.get_memory_usage()116                117                # Generate118                start_time = time.time()119                120                with torch.no_grad():121                    outputs = self.model.generate(122                        **inputs,123                        max_length=max_length,124                        temperature=temperature,125                        do_sample=temperature > 0,126                        pad_token_id=self.tokenizer.eos_token_id127                    )128                129                if self.device == "cuda":130                    torch.cuda.synchronize()131                132                end_time = time.time()133                134                # Measure memory after135                ram_after, gpu_after = self.get_memory_usage()136                137                # Calculate metrics138                inference_time = end_time - start_time139                generation_length = outputs.shape[1] - prompt_length140                tokens_per_second = generation_length / inference_time if inference_time > 0 else 0141                memory_used = ram_after - ram_before142                gpu_memory = (gpu_after - gpu_before) if gpu_after and gpu_before else None143                144                result = BenchmarkResult(145                    model_name=self.model_name,146                    task=task,147                    prompt_length=prompt_length,148                    generation_length=generation_length,149                    temperature=temperature,150                    inference_time=inference_time,151                    tokens_per_second=tokens_per_second,152                    memory_used_mb=memory_used,153                    gpu_memory_mb=gpu_memory,154                    success=True155                )156                157                run_results.append(result)158                self.results.append(result)159                160            except Exception as e:161                logger.error(f"Benchmark failed: {e}")162                result = BenchmarkResult(163                    model_name=self.model_name,164                    task=task,165                    prompt_length=0,166                    generation_length=0,167                    temperature=temperature,168                    inference_time=0,169                    tokens_per_second=0,170                    memory_used_mb=0,171                    gpu_memory_mb=None,172                    success=False,173                    error=str(e)174                )175                run_results.append(result)176                self.results.append(result)177        178        return run_results179    180    def run_benchmark_suite(self) -> List[BenchmarkResult]:181        """Run comprehensive benchmark suite"""182        logger.info("Starting comprehensive benchmark suite...")183        184        test_cases = [185            {186                "prompt": "def fibonacci(n):",187                "task": "simple_function",188                "max_length": 256,189                "temperature": 0.7190            },191            {192                "prompt": "Write a Python class for a binary search tree with insert, search, and delete methods:",193                "task": "complex_class",194                "max_length": 1024,195                "temperature": 0.7196            },197            {198                "prompt": "Implement quicksort algorithm in Python with detailed comments:",199                "task": "algorithm",200                "max_length": 512,201                "temperature": 0.5202            },203            {204                "prompt": "Solve: What is the derivative of f(x) = x^3 + 2x^2 - 5x + 3?",205                "task": "math_simple",206                "max_length": 256,207                "temperature": 0.3208            },209            {210                "prompt": "Prove using mathematical induction that the sum of first n natural numbers is n(n+1)/2:",211                "task": "math_proof",212                "max_length": 1024,213                "temperature": 0.2214            },215            {216                "prompt": "Design a RESTful API for a todo list application with proper documentation:",217                "task": "system_design",218                "max_length": 2048,219                "temperature": 0.7220            },221        ]222        223        all_results = []224        225        for test_case in tqdm(test_cases, desc="Running benchmarks"):226            results = self.benchmark_inference(227                prompt=test_case["prompt"],228                task=test_case["task"],229                max_length=test_case["max_length"],230                temperature=test_case["temperature"],231                num_runs=3232            )233            all_results.extend(results)234        235        logger.info("Benchmark suite completed!")236        return all_results237    238    def aggregate_results(self) -> AggregatedResults:239        """Aggregate benchmark results"""240        if not self.results:241            raise ValueError("No benchmark results available")242        243        successful = [r for r in self.results if r.success]244        245        if not successful:246            raise ValueError("No successful benchmark runs")247        248        inference_times = [r.inference_time for r in successful]249        tokens_per_sec = [r.tokens_per_second for r in successful]250        memory_usage = [r.memory_used_mb for r in successful]251        252        return AggregatedResults(253            model_name=self.model_name,254            total_tests=len(self.results),255            successful_tests=len(successful),256            failed_tests=len(self.results) - len(successful),257            avg_inference_time=np.mean(inference_times),258            avg_tokens_per_second=np.mean(tokens_per_sec),259            avg_memory_mb=np.mean(memory_usage),260            min_inference_time=np.min(inference_times),261            max_inference_time=np.max(inference_times),262            std_inference_time=np.std(inference_times)263        )264    265    def save_results(self, output_file: str = "benchmark_results.json"):266        """Save benchmark results to file"""267        results_dict = [asdict(r) for r in self.results]268        269        with open(output_file, 'w') as f:270            json.dump(results_dict, f, indent=2)271        272        logger.info(f"Results saved to {output_file}")273    274    def generate_report(self, output_file: str = "benchmark_report.txt"):275        """Generate human-readable benchmark report"""276        agg = self.aggregate_results()277        278        report = f"""279{'='*80}280HELION-OSC BENCHMARK REPORT281{'='*80}282 283Model: {agg.model_name}284Device: {self.device}285 286OVERALL STATISTICS287{'='*80}288Total Tests: {agg.total_tests}289Successful: {agg.successful_tests}290Failed: {agg.failed_tests}291Success Rate: {(agg.successful_tests/agg.total_tests)*100:.2f}%292 293PERFORMANCE METRICS294{'='*80}295Average Inference Time: {agg.avg_inference_time:.4f} seconds296Min Inference Time: {agg.min_inference_time:.4f} seconds297Max Inference Time: {agg.max_inference_time:.4f} seconds298Std Inference Time: {agg.std_inference_time:.4f} seconds299 300Average Tokens/Second: {agg.avg_tokens_per_second:.2f}301Average Memory Usage: {agg.avg_memory_mb:.2f} MB302 303PER-TASK BREAKDOWN304{'='*80}305"""306        307        # Group by task308        df = pd.DataFrame([asdict(r) for r in self.results if r.success])309        if not df.empty:310            task_stats = df.groupby('task').agg({311                'inference_time': ['mean', 'min', 'max'],312                'tokens_per_second': 'mean',313                'memory_used_mb': 'mean'314            })315            316            report += task_stats.to_string()317        318        report += f"\n\n{'='*80}\n"319        320        with open(output_file, 'w') as f:321            f.write(report)322        323        logger.info(f"Report saved to {output_file}")324        print(report)325    326    def plot_results(self, output_dir: str = "./benchmark_plots"):327        """Generate visualization plots"""328        import os329        os.makedirs(output_dir, exist_ok=True)330        331        df = pd.DataFrame([asdict(r) for r in self.results if r.success])332        333        if df.empty:334            logger.warning("No data to plot")335            return336        337        # Set style338        sns.set_style("whitegrid")339        340        # Plot 1: Inference time by task341        plt.figure(figsize=(12, 6))342        sns.barplot(data=df, x='task', y='inference_time')343        plt.xticks(rotation=45, ha='right')344        plt.title('Inference Time by Task')345        plt.ylabel('Time (seconds)')346        plt.tight_layout()347        plt.savefig(f"{output_dir}/inference_time_by_task.png", dpi=300)348        plt.close()349        350        # Plot 2: Tokens per second by task351        plt.figure(figsize=(12, 6))352        sns.barplot(data=df, x='task', y='tokens_per_second')353        plt.xticks(rotation=45, ha='right')354        plt.title('Tokens Per Second by Task')355        plt.ylabel('Tokens/Second')356        plt.tight_layout()357        plt.savefig(f"{output_dir}/tokens_per_second_by_task.png", dpi=300)358        plt.close()359        360        # Plot 3: Memory usage by task361        plt.figure(figsize=(12, 6))362        sns.barplot(data=df, x='task', y='memory_used_mb')363        plt.xticks(rotation=45, ha='right')364        plt.title('Memory Usage by Task')365        plt.ylabel('Memory (MB)')366        plt.tight_layout()367        plt.savefig(f"{output_dir}/memory_usage_by_task.png", dpi=300)368        plt.close()369        370        # Plot 4: Scatter plot - generation length vs inference time371        plt.figure(figsize=(10, 6))372        sns.scatterplot(data=df, x='generation_length', y='inference_time', hue='task', s=100)373        plt.title('Generation Length vs Inference Time')374        plt.xlabel('Generation Length (tokens)')375        plt.ylabel('Inference Time (seconds)')376        plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')377        plt.tight_layout()378        plt.savefig(f"{output_dir}/length_vs_time.png", dpi=300)379        plt.close()380        381        logger.info(f"Plots saved to {output_dir}")382 383 384class ComparisonBenchmark:385    """Compare multiple models"""386    387    def __init__(self, model_names: List[str]):388        self.model_names = model_names389        self.benchmarks = {}390    391    def run_comparison(self):392        """Run benchmarks for all models"""393        for model_name in self.model_names:394            logger.info(f"\nBenchmarking {model_name}...")395            try:396                benchmark = PerformanceBenchmark(model_name)397                benchmark.run_benchmark_suite()398                self.benchmarks[model_name] = benchmark399            except Exception as e:400                logger.error(f"Failed to benchmark {model_name}: {e}")401    402    def generate_comparison_report(self, output_file: str = "comparison_report.txt"):403        """Generate comparison report"""404        report = f"""405{'='*80}406MODEL COMPARISON REPORT407{'='*80}408 409"""410        411        for model_name, benchmark in self.benchmarks.items():412            agg = benchmark.aggregate_results()413            report += f"""414Model: {model_name}415{'='*80}416Avg Inference Time: {agg.avg_inference_time:.4f}s417Avg Tokens/Second: {agg.avg_tokens_per_second:.2f}418Avg Memory Usage: {agg.avg_memory_mb:.2f} MB419Success Rate: {(agg.successful_tests/agg.total_tests)*100:.2f}%420 421"""422        423        with open(output_file, 'w') as f:424            f.write(report)425        426        print(report)427        logger.info(f"Comparison report saved to {output_file}")428 429 430def main():431    """Main benchmark script"""432    import argparse433    434    parser = argparse.ArgumentParser(description="Benchmark Helion-OSC model")435    parser.add_argument("--model", type=str, default="DeepXR/Helion-OSC")436    parser.add_argument("--output-dir", type=str, default="./benchmark_results")437    parser.add_argument("--compare", nargs='+', help="List of models to compare")438    parser.add_argument("--plot", action="store_true", help="Generate plots")439    440    args = parser.parse_args()441    442    import os443    os.makedirs(args.output_dir, exist_ok=True)444    445    if args.compare:446        # Comparison mode447        comparison = ComparisonBenchmark(args.compare)448        comparison.run_comparison()449        comparison.generate_comparison_report(450            os.path.join(args.output_dir, "comparison_report.txt")451        )452    else:453        # Single model benchmark454        benchmark = PerformanceBenchmark(args.model)455        benchmark.run_benchmark_suite()456        benchmark.save_results(os.path.join(args.output_dir, "benchmark_results.json"))457        benchmark.generate_report(os.path.join(args.output_dir, "benchmark_report.txt"))458        459        if args.plot:460            benchmark.plot_results(os.path.join(args.output_dir, "plots"))461 462 463if __name__ == "__main__":464    main()