chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
404
third_party/sglang/benchmark/asr/bench_sglang.py
vendored
Normal file
404
third_party/sglang/benchmark/asr/bench_sglang.py
vendored
Normal file
@@ -0,0 +1,404 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import time
|
||||
from statistics import mean, median
|
||||
|
||||
import httpx
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile
|
||||
from datasets import load_dataset
|
||||
from evaluate import load
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
||||
def to_bytes(y, sr):
|
||||
buffer = io.BytesIO()
|
||||
soundfile.write(buffer, y, sr, format="WAV")
|
||||
buffer.seek(0)
|
||||
return buffer
|
||||
|
||||
|
||||
async def run_asr_chat(client, model_name, y, sr):
|
||||
"""Use chat completions API with audio_url for ASR."""
|
||||
with to_bytes(y, sr) as f:
|
||||
audio_bytes = f.read()
|
||||
audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
|
||||
|
||||
start_time = time.perf_counter()
|
||||
response = await client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": f"data:audio/wav;base64,{audio_base64}"},
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
temperature=0.0,
|
||||
)
|
||||
end_time = time.perf_counter()
|
||||
|
||||
asr_text = response.choices[0].message.content
|
||||
latency = end_time - start_time
|
||||
return latency, asr_text
|
||||
|
||||
|
||||
def run_asr_transcription_sync(client, model_name, y, sr, language=None):
|
||||
"""Use audio transcriptions API for ASR (sync version)."""
|
||||
audio_buffer = to_bytes(y, sr)
|
||||
audio_buffer.name = "audio.wav" # OpenAI client needs a name attribute
|
||||
|
||||
start_time = time.perf_counter()
|
||||
kwargs = {
|
||||
"model": model_name,
|
||||
"file": audio_buffer,
|
||||
}
|
||||
if language:
|
||||
kwargs["language"] = language
|
||||
|
||||
transcription = client.audio.transcriptions.create(**kwargs)
|
||||
end_time = time.perf_counter()
|
||||
|
||||
latency = end_time - start_time
|
||||
return latency, transcription.text
|
||||
|
||||
|
||||
def run_asr_transcription_stream_sync(
|
||||
base_url, model_name, y, sr, language=None, show_stream=False
|
||||
):
|
||||
"""Use audio transcriptions API with streaming for ASR."""
|
||||
audio_buffer = to_bytes(y, sr)
|
||||
audio_bytes = audio_buffer.read()
|
||||
|
||||
data = {
|
||||
"model": model_name,
|
||||
"response_format": "json",
|
||||
"stream": "true",
|
||||
}
|
||||
if language:
|
||||
data["language"] = language
|
||||
|
||||
start_time = time.perf_counter()
|
||||
text_chunks = []
|
||||
|
||||
if show_stream:
|
||||
print("[STREAM] ", end="", flush=True)
|
||||
|
||||
with httpx.stream(
|
||||
"POST",
|
||||
f"{base_url}/v1/audio/transcriptions",
|
||||
data=data,
|
||||
files={"file": ("audio.wav", audio_bytes, "audio/wav")},
|
||||
timeout=60.0,
|
||||
) as response:
|
||||
for line in response.iter_lines():
|
||||
if line.startswith("data: ") and not line.startswith("data: [DONE]"):
|
||||
try:
|
||||
chunk = json.loads(line[6:])
|
||||
if "choices" in chunk and chunk["choices"]:
|
||||
delta = chunk["choices"][0].get("delta", {})
|
||||
content = delta.get("content", "")
|
||||
if content:
|
||||
text_chunks.append(content)
|
||||
if show_stream:
|
||||
print(content, end="", flush=True)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
if show_stream:
|
||||
print() # newline after stream
|
||||
|
||||
end_time = time.perf_counter()
|
||||
latency = end_time - start_time
|
||||
return latency, "".join(text_chunks)
|
||||
|
||||
|
||||
async def run_asr_transcription(
|
||||
client,
|
||||
model_name,
|
||||
y,
|
||||
sr,
|
||||
language=None,
|
||||
stream=False,
|
||||
base_url=None,
|
||||
show_stream=False,
|
||||
):
|
||||
"""Async wrapper for transcription API (runs sync call in executor)."""
|
||||
loop = asyncio.get_event_loop()
|
||||
if stream:
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
run_asr_transcription_stream_sync,
|
||||
base_url,
|
||||
model_name,
|
||||
y,
|
||||
sr,
|
||||
language,
|
||||
show_stream,
|
||||
)
|
||||
return await loop.run_in_executor(
|
||||
None, run_asr_transcription_sync, client, model_name, y, sr, language
|
||||
)
|
||||
|
||||
|
||||
async def bound_asr(
|
||||
sem,
|
||||
client,
|
||||
model_name,
|
||||
tokenizer,
|
||||
audio,
|
||||
reference,
|
||||
api_type="chat",
|
||||
language=None,
|
||||
stream=False,
|
||||
base_url=None,
|
||||
show_stream=False,
|
||||
):
|
||||
async with sem:
|
||||
try:
|
||||
if api_type == "transcription":
|
||||
latency, text = await run_asr_transcription(
|
||||
client,
|
||||
model_name,
|
||||
*audio,
|
||||
language=language,
|
||||
stream=stream,
|
||||
base_url=base_url,
|
||||
show_stream=show_stream,
|
||||
)
|
||||
else:
|
||||
latency, text = await run_asr_chat(client, model_name, *audio)
|
||||
|
||||
# Calculate tokens for throughput metrics
|
||||
num_output_tokens = len(tokenizer(text, add_special_tokens=False).input_ids)
|
||||
|
||||
# Normalize for WER evaluation
|
||||
# Whisper tokenizer has a normalize method
|
||||
if hasattr(tokenizer, "normalize"):
|
||||
out = tokenizer.normalize(text)
|
||||
ref = tokenizer.normalize(reference)
|
||||
else:
|
||||
out = text.lower().strip()
|
||||
ref = reference.lower().strip()
|
||||
|
||||
return latency, num_output_tokens, out, ref
|
||||
except Exception as e:
|
||||
print(f"Error during ASR: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def process_dataset(
|
||||
model_name,
|
||||
client,
|
||||
data,
|
||||
concurrent_request,
|
||||
api_type="chat",
|
||||
language=None,
|
||||
stream=False,
|
||||
base_url=None,
|
||||
show_predictions=False,
|
||||
):
|
||||
sem = asyncio.Semaphore(concurrent_request)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
|
||||
# Warmup
|
||||
print("Performing warmup...")
|
||||
audio_warmup, sr_warmup = (
|
||||
data[0]["audio"]["array"],
|
||||
data[0]["audio"]["sampling_rate"],
|
||||
)
|
||||
await bound_asr(
|
||||
sem,
|
||||
client,
|
||||
model_name,
|
||||
tokenizer,
|
||||
(audio_warmup, sr_warmup),
|
||||
"",
|
||||
api_type=api_type,
|
||||
language=language,
|
||||
stream=stream,
|
||||
base_url=base_url,
|
||||
show_stream=False, # Don't show stream during warmup
|
||||
)
|
||||
|
||||
tasks = []
|
||||
print(f"Processing {len(data)} samples...")
|
||||
for sample in data:
|
||||
audio, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
bound_asr(
|
||||
sem,
|
||||
client,
|
||||
model_name,
|
||||
tokenizer,
|
||||
(audio, sr),
|
||||
sample["text"],
|
||||
api_type=api_type,
|
||||
language=language,
|
||||
stream=stream,
|
||||
base_url=base_url,
|
||||
show_stream=show_predictions and stream,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
results = await asyncio.gather(*tasks)
|
||||
return [r for r in results if r is not None]
|
||||
|
||||
|
||||
def run_evaluation(args):
|
||||
# Use sync client for transcription API, async for chat API
|
||||
if args.api_type == "transcription":
|
||||
client = OpenAI(base_url=f"{args.base_url}/v1", api_key="None")
|
||||
else:
|
||||
client = AsyncOpenAI(base_url=f"{args.base_url}/v1", api_key="None")
|
||||
|
||||
print(f"Loading dataset: {args.dataset}...")
|
||||
print(f"Using API type: {args.api_type}" + (f" (streaming)" if args.stream else ""))
|
||||
dataset = load_dataset(args.dataset, split=args.split)
|
||||
|
||||
# Filter by duration if needed (Whisper max is 30s)
|
||||
def add_duration(sample):
|
||||
y, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"]
|
||||
sample["duration_ms"] = librosa.get_duration(y=y, sr=sr) * 1000
|
||||
return sample
|
||||
|
||||
if "duration_ms" not in dataset.column_names:
|
||||
dataset = dataset.map(add_duration)
|
||||
|
||||
dataset = dataset.filter(lambda x: x["duration_ms"] < 30000)
|
||||
|
||||
if args.n_examples > 0:
|
||||
dataset = dataset.select(range(min(args.n_examples, len(dataset))))
|
||||
|
||||
start = time.perf_counter()
|
||||
results = asyncio.run(
|
||||
process_dataset(
|
||||
args.model,
|
||||
client,
|
||||
dataset,
|
||||
args.concurrency,
|
||||
api_type=args.api_type,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
base_url=args.base_url,
|
||||
show_predictions=args.show_predictions,
|
||||
)
|
||||
)
|
||||
total_test_time = time.perf_counter() - start
|
||||
|
||||
if not results:
|
||||
print("No successful results to evaluate.")
|
||||
return
|
||||
|
||||
# Metrics
|
||||
latencies = [res[0] for res in results]
|
||||
total_tokens = sum([res[1] for res in results])
|
||||
predictions = [res[2] for res in results]
|
||||
references = [res[3] for res in results]
|
||||
|
||||
wer_metric = load("wer")
|
||||
wer_score = 100 * wer_metric.compute(references=references, predictions=predictions)
|
||||
|
||||
print("-" * 30)
|
||||
print(f"Results for {args.model}:")
|
||||
print(f"Total Requests: {len(results)}")
|
||||
print(f"WER: {wer_score:.4f}")
|
||||
print(f"Average Latency: {mean(latencies):.4f}s")
|
||||
print(f"Median Latency: {median(latencies):.4f}s")
|
||||
print(f"95th Latency: {np.percentile(latencies, 95):.4f}s")
|
||||
print(f"Throughput: {len(results) / total_test_time:.2f} req/s")
|
||||
print(f"Token Throughput: {total_tokens / total_test_time:.2f} tok/s")
|
||||
print(f"Total Test Time: {total_test_time:.4f}s")
|
||||
print("-" * 30)
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w") as f:
|
||||
import json
|
||||
|
||||
json.dump(
|
||||
{
|
||||
"model": args.model,
|
||||
"dataset": args.dataset,
|
||||
"wer": wer_score,
|
||||
"avg_latency": mean(latencies),
|
||||
"throughput": len(results) / total_test_time,
|
||||
"token_throughput": total_tokens / total_test_time,
|
||||
},
|
||||
f,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
if args.show_predictions:
|
||||
print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20)
|
||||
num_to_show = min(args.print_n, len(results))
|
||||
for i in range(num_to_show):
|
||||
print(f"Sample {i+1}:")
|
||||
print(f" REF: {references[i]}")
|
||||
print(f" PRED: {predictions[i]}")
|
||||
print("-" * 40)
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Benchmark sGLang ASR performance.")
|
||||
parser.add_argument(
|
||||
"--base-url", default="http://localhost:30000", help="sGLang server base URL"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model", default="openai/whisper-large-v3", help="Model name on the server"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
default="D4nt3/esb-datasets-earnings22-validation-tiny-filtered",
|
||||
help="HF dataset repo",
|
||||
)
|
||||
parser.add_argument("--split", default="validation", help="Dataset split")
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, default=4, help="Number of concurrent requests"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n-examples",
|
||||
"-n",
|
||||
type=int,
|
||||
default=-1,
|
||||
help="Number of examples to test (-1 for all)",
|
||||
)
|
||||
parser.add_argument("--output", help="Path to save results in JSON")
|
||||
parser.add_argument(
|
||||
"--show-predictions",
|
||||
action="store_true",
|
||||
help="Print sample predictions and references",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--print-n", type=int, default=5, help="Number of sample predictions to print"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-type",
|
||||
choices=["chat", "transcription"],
|
||||
default="chat",
|
||||
help="API type to use: 'chat' for chat completions with audio_url, 'transcription' for audio.transcriptions API",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--language",
|
||||
default=None,
|
||||
help="Language code for transcription API (e.g., 'en')",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
help="Use streaming mode for transcription API",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
run_evaluation(args)
|
||||
Reference in New Issue
Block a user