Spaces:
Sleeping
Sleeping
| """Direct SentenceTransformers baseline benchmark. | |
| This benchmark intentionally avoids HTTP, FastAPI, and DynamicBatcher overhead. | |
| It measures one direct model.encode([text]) call per input text. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import time | |
| from pathlib import Path | |
| from typing import Any | |
| import torch | |
| from sentence_transformers import SentenceTransformer | |
| from utils import base_metadata, build_texts, format_ms, percentile, seconds_to_ms, write_json | |
| DEFAULT_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" | |
| DEFAULT_DEVICE = "mps" if torch.backends.mps.is_available() else "cpu" | |
| DEFAULT_REQUESTS = 40 | |
| DEFAULT_TEXTS_PER_REQUEST = 2 | |
| DEFAULT_WARMUP = 0 | |
| def encode_one(model: SentenceTransformer, text: str) -> float: | |
| start = time.perf_counter() | |
| model.encode([text]) | |
| return time.perf_counter() - start | |
| def run_naive( | |
| model: SentenceTransformer, | |
| requests: int, | |
| texts_per_request: int, | |
| ) -> tuple[list[float], float]: | |
| texts = build_texts(texts_per_request) | |
| latencies: list[float] = [] | |
| start = time.perf_counter() | |
| for _ in range(requests): | |
| for text in texts: | |
| latencies.append(encode_one(model, text)) | |
| wall_time = time.perf_counter() - start | |
| return latencies, wall_time | |
| def summarize( | |
| latencies: list[float], | |
| wall_time_seconds: float, | |
| metadata: dict[str, Any], | |
| ) -> dict[str, Any]: | |
| total_sequences = len(latencies) | |
| return { | |
| "metadata": metadata, | |
| "summary": { | |
| "total_requests": metadata["requests"], | |
| "total_sequences": total_sequences, | |
| "wall_time_seconds": wall_time_seconds, | |
| "throughput_sequences_per_sec": total_sequences / wall_time_seconds if wall_time_seconds > 0 else 0.0, | |
| "avg_latency_ms": seconds_to_ms(sum(latencies) / total_sequences) if total_sequences else None, | |
| "p50_latency_ms": seconds_to_ms(percentile(latencies, 50)), | |
| "p95_latency_ms": seconds_to_ms(percentile(latencies, 95)), | |
| "p99_latency_ms": seconds_to_ms(percentile(latencies, 99)), | |
| }, | |
| } | |
| def print_summary(report: dict[str, Any]) -> None: | |
| metadata = report["metadata"] | |
| summary = report["summary"] | |
| print("\nNaive Benchmark") | |
| print(f"Label: {metadata['label']}") | |
| print(f"Model: {metadata['model_name']}") | |
| print(f"Device: {metadata['device']}") | |
| print(f"Requests: {metadata['requests']}") | |
| print(f"Texts/request: {metadata['texts_per_request']}") | |
| print(f"Warmup requests: {metadata['warmup']}") | |
| print("\nResults") | |
| print(f"Total sequences: {summary['total_sequences']}") | |
| print(f"Wall time: {summary['wall_time_seconds']:.3f}s") | |
| print(f"Sequence throughput: {summary['throughput_sequences_per_sec']:.2f} seq/s") | |
| print(f"Avg latency: {format_ms(summary['avg_latency_ms'])}") | |
| print(f"p50 latency: {format_ms(summary['p50_latency_ms'])}") | |
| print(f"p95 latency: {format_ms(summary['p95_latency_ms'])}") | |
| print(f"p99 latency: {format_ms(summary['p99_latency_ms'])}") | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Benchmark naive direct SentenceTransformers encoding.") | |
| parser.add_argument("--label", default="naive-direct", help="Human-readable run label stored in output metadata.") | |
| parser.add_argument("--model-name", default=DEFAULT_MODEL_NAME, help=f"SentenceTransformer model. Default: {DEFAULT_MODEL_NAME}") | |
| parser.add_argument("--device", default=DEFAULT_DEVICE, help=f"Torch device passed to SentenceTransformer. Default: {DEFAULT_DEVICE}") | |
| parser.add_argument("--requests", type=int, default=DEFAULT_REQUESTS, help="Request-equivalent loop count.") | |
| parser.add_argument("--texts-per-request", type=int, default=DEFAULT_TEXTS_PER_REQUEST, help="Texts per request-equivalent loop.") | |
| parser.add_argument("--warmup", type=int, default=DEFAULT_WARMUP, help="Warmup request-equivalent loops excluded from results.") | |
| parser.add_argument("--output", type=Path, help="Optional JSON output path.") | |
| return parser.parse_args() | |
| def validate_args(args: argparse.Namespace) -> None: | |
| if args.requests < 1: | |
| raise SystemExit("--requests must be >= 1") | |
| if args.texts_per_request < 1: | |
| raise SystemExit("--texts-per-request must be >= 1") | |
| if args.warmup < 0: | |
| raise SystemExit("--warmup must be >= 0") | |
| def main() -> None: | |
| args = parse_args() | |
| validate_args(args) | |
| model = SentenceTransformer(args.model_name, device=args.device) | |
| if args.warmup: | |
| run_naive(model, args.warmup, args.texts_per_request) | |
| latencies, wall_time = run_naive(model, args.requests, args.texts_per_request) | |
| metadata = { | |
| **base_metadata(), | |
| "label": args.label, | |
| "model_name": args.model_name, | |
| "device": args.device, | |
| "requests": args.requests, | |
| "texts_per_request": args.texts_per_request, | |
| "warmup": args.warmup, | |
| } | |
| report = summarize(latencies, wall_time, metadata) | |
| print_summary(report) | |
| if args.output: | |
| write_json(report, args.output) | |
| print(f"\nWrote JSON results to {args.output}") | |
| if __name__ == "__main__": | |
| main() | |