Spaces:
Running on Zero
Running on Zero
| """ | |
| Benchmarking & Performance Profiler for Nova AI - ChatGPT Astra Edition. | |
| Measures Time-To-First-Token (TTFT), Tokens-Per-Second (TPS), | |
| and memory footprint across Nova Astra models. | |
| """ | |
| import time | |
| import psutil | |
| from typing import List, Dict, Any | |
| from rich.table import Table | |
| from rich.panel import Panel | |
| from src.config import AVAILABLE_MODELS, GenerationConfig | |
| from src.model_engine import NovaAstraEngine | |
| from src.utils import console, get_system_hardware_info | |
| def run_benchmark( | |
| models_to_test: List[str] = None, | |
| test_prompt: str = "Explain the fundamentals of quantum computing in three structured bullet points." | |
| ) -> List[Dict[str, Any]]: | |
| """Profile latency, generation throughput, and RAM usage for Nova Astra models.""" | |
| if not models_to_test: | |
| from src.config import config | |
| models_to_test = [config.default_model_id] | |
| hw = get_system_hardware_info() | |
| console.print(Panel( | |
| f"[bold]System Hardware:[/bold] {hw['device_name']} | " | |
| f"[bold]Logical Cores:[/bold] {hw['cpu_count_logical']} | " | |
| f"[bold]Available RAM:[/bold] {hw['available_ram_gb']} GB", | |
| title="🌟 Nova Astra Hardware Profiler", | |
| border_style="cyan" | |
| )) | |
| results = [] | |
| for model_id in models_to_test: | |
| console.print(f"\n[bold yellow]⏱️ Profiling Astra Model:[/bold yellow] [bold cyan]{model_id}[/bold cyan]") | |
| engine = NovaAstraEngine(model_id=model_id) | |
| # Measure Load Time & RAM | |
| ram_before = psutil.virtual_memory().used / (1024 ** 2) | |
| load_start = time.perf_counter() | |
| engine.load_model() | |
| load_time = time.perf_counter() - load_start | |
| ram_after = psutil.virtual_memory().used / (1024 ** 2) | |
| model_ram_mb = max(ram_after - ram_before, 0.0) | |
| # Warm-up run | |
| messages = [ | |
| {"role": "system", "content": "You are Nova Astra, a concise assistant."}, | |
| {"role": "user", "content": "Hi"} | |
| ] | |
| engine.generate_response(messages, GenerationConfig(max_new_tokens=10)) | |
| # Actual Benchmark Run | |
| messages = [ | |
| {"role": "system", "content": "You are Nova Astra, an expert assistant."}, | |
| {"role": "user", "content": test_prompt} | |
| ] | |
| cfg = GenerationConfig(max_new_tokens=128, temperature=0.7) | |
| response, metrics = engine.generate_response(messages, cfg) | |
| record = { | |
| "model_id": model_id.split("/")[-1] if "/" in model_id else model_id, | |
| "load_time_sec": round(load_time, 2), | |
| "ram_mb": round(model_ram_mb, 1), | |
| "ttft_ms": metrics.get("time_to_first_token_ms", 0.0), | |
| "tokens_per_sec": metrics.get("tokens_per_second", 0.0), | |
| "tokens_generated": metrics.get("tokens_generated", 0), | |
| "total_latency_sec": metrics.get("total_latency_seconds", 0.0) | |
| } | |
| results.append(record) | |
| # Display Rich Comparison Table | |
| table = Table(title="🌟 Nova AI (ChatGPT Astra Edition) Benchmark Results", border_style="green") | |
| table.add_column("Model Name", style="bold cyan") | |
| table.add_column("Load Time (s)", justify="right") | |
| table.add_column("Model RAM (MB)", justify="right") | |
| table.add_column("TTFT (ms)", justify="right") | |
| table.add_column("Throughput (Tokens/s)", justify="right", style="bold green") | |
| table.add_column("Total Latency (s)", justify="right") | |
| for r in results: | |
| table.add_row( | |
| r["model_id"], | |
| str(r["load_time_sec"]), | |
| str(r["ram_mb"]), | |
| str(r["ttft_ms"]), | |
| str(r["tokens_per_sec"]), | |
| str(r["total_latency_sec"]) | |
| ) | |
| console.print(table) | |
| return results | |