File size: 2,522 Bytes
fc9409e | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | 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()))
|