| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import time |
|
|
| from server import ( |
| END_SOLUTION, |
| THINKING_TOKEN_BUDGETS, |
| ChatMessage, |
| ChatRequest, |
| default_model_id, |
| inference_config, |
| stream_model_once, |
| ) |
|
|
|
|
| DEFAULT_PROMPT = "Find the value of x if 3x + 7 = 31. Show the reasoning briefly." |
|
|
|
|
| async def benchmark(args: argparse.Namespace) -> None: |
| request = ChatRequest( |
| messages=[ChatMessage(role="user", content=args.prompt)], |
| model=args.model or default_model_id(), |
| temperature=0.0 if args.fast_mode else args.temperature, |
| max_tokens=args.max_tokens, |
| thinking_effort=args.thinking_effort, |
| fast_mode=args.fast_mode, |
| ) |
| generation_state: dict[str, object] = {} |
| max_tokens = args.max_tokens |
| if args.fast_mode: |
| max_tokens += THINKING_TOKEN_BUDGETS[request.thinking_effort] |
|
|
| started = time.monotonic() |
| chunks: list[str] = [] |
| async for chunk in stream_model_once( |
| request, |
| max_tokens=max_tokens, |
| state=generation_state, |
| stop_sequences=(END_SOLUTION,), |
| temperature_override=0.0 if args.fast_mode else None, |
| ): |
| chunks.append(chunk) |
| elapsed = time.monotonic() - started |
|
|
| generated_tokens = int(generation_state.get("generated_tokens") or 0) |
| tokens_per_second = generated_tokens / elapsed if elapsed > 0 else 0.0 |
| config = inference_config() |
| print(f"backend={config.backend}") |
| print(f"model={request.model}") |
| print(f"fast_mode={request.fast_mode}") |
| print(f"generated_tokens={generated_tokens}") |
| print(f"seconds={elapsed:.3f}") |
| print(f"tokens_per_second={tokens_per_second:.2f}") |
| print(f"stop_reason={generation_state.get('stop_reason')}") |
| print(f"text_chars={sum(len(chunk) for chunk in chunks)}") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Benchmark one Transformers streaming pass.") |
| parser.add_argument("--model", default=None, help="Model id from /api/models. Defaults to server default.") |
| parser.add_argument("--prompt", default=DEFAULT_PROMPT) |
| parser.add_argument("--max-tokens", type=int, default=512) |
| parser.add_argument("--temperature", type=float, default=0.6) |
| parser.add_argument("--thinking-effort", default="speed", choices=sorted(THINKING_TOKEN_BUDGETS)) |
| parser.add_argument("--fast-mode", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(benchmark(parse_args())) |
|
|
