Spaces:
Configuration error
Configuration error
File size: 1,649 Bytes
9644d0b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | import torch
import time
import asyncio
from model_manager import ModelManager
from transformers import AutoTokenizer
async def run_benchmark():
manager = ModelManager()
model_id = "gemma3-270m-it"
print("Loading model...")
model_entry = await manager.get_model(model_id, px_subjective=True)
model = model_entry["model"]
tokenizer = model_entry["tokenizer"]
prompt = "Explain the concept of recursion in computer science using an analogy of nested boxes."
messages = [{"role": "user", "content": prompt}]
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
print("\nWarmup (1 token)...")
with torch.no_grad():
model.generate(**inputs, max_new_tokens=1)
print("\nBenchmarking generation (50 tokens)...")
start_time = time.time()
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=50)
torch.cuda.synchronize()
end_time = time.time()
duration = end_time - start_time
tokens_generated = 50
tps = tokens_generated / duration
metrics = manager.get_px_metrics(model_id)
steps = metrics.get("steps", 0)
print(f"\n--- Benchmark Results ---")
print(f"Time: {duration:.2f} s")
print(f"Tokens/sec: {tps:.2f}")
print(f"PX Recursion Steps per token (last token): {steps}")
print(f"Average time per token: {(duration/tokens_generated)*1000:.2f} ms")
print("-------------------------\n")
if __name__ == "__main__":
asyncio.run(run_benchmark())
|