Whisper model support & /v1/audio/transcriptions endpoint & benchmark (#16983)
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com> Co-authored-by: MahmoudAshraf97 <hassouna97.ma@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
MahmoudAshraf97
gemini-code-assist[bot]
parent
3a11e7dad9
commit
581bf53e03
@@ -0,0 +1,166 @@
|
|||||||
|
# ASR Benchmark
|
||||||
|
|
||||||
|
This benchmark evaluates the performance and accuracy (Word Error Rate - WER) of Automatic Speech Recognition (ASR) models served via SGLang.
|
||||||
|
|
||||||
|
## Supported Models
|
||||||
|
|
||||||
|
- `openai/whisper-large-v3`
|
||||||
|
- `openai/whisper-large-v3-turbo`
|
||||||
|
|
||||||
|
## Setup
|
||||||
|
|
||||||
|
Install the required dependencies:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
apt install ffmpeg
|
||||||
|
pip install librosa soundfile datasets evaluate jiwer transformers openai torchcodec torch
|
||||||
|
```
|
||||||
|
|
||||||
|
## Running the Benchmark
|
||||||
|
|
||||||
|
### 1. Start SGLang Server
|
||||||
|
|
||||||
|
Launch the SGLang server with a Whisper model:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m sglang.launch_server --model-path openai/whisper-large-v3 --port 30000
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2. Run the Benchmark Script
|
||||||
|
|
||||||
|
Basic usage (using chat completions API):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python bench_sglang.py --base-url http://localhost:30000 --model openai/whisper-large-v3 --n-examples 10
|
||||||
|
```
|
||||||
|
|
||||||
|
Using the OpenAI-compatible transcription API:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python bench_sglang.py \
|
||||||
|
--base-url http://localhost:30000 \
|
||||||
|
--model openai/whisper-large-v3 \
|
||||||
|
--api-type transcription \
|
||||||
|
--language English \
|
||||||
|
--n-examples 10
|
||||||
|
```
|
||||||
|
|
||||||
|
Run with streaming and show real-time output:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python bench_sglang.py \
|
||||||
|
--base-url http://localhost:30000 \
|
||||||
|
--model openai/whisper-large-v3 \
|
||||||
|
--api-type transcription \
|
||||||
|
--stream \
|
||||||
|
--show-predictions \
|
||||||
|
--concurrency 1
|
||||||
|
```
|
||||||
|
|
||||||
|
Run with higher concurrency and save results:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python bench_sglang.py \
|
||||||
|
--base-url http://localhost:30000 \
|
||||||
|
--model openai/whisper-large-v3 \
|
||||||
|
--concurrency 8 \
|
||||||
|
--n-examples 100 \
|
||||||
|
--output results.json \
|
||||||
|
--show-predictions
|
||||||
|
```
|
||||||
|
|
||||||
|
## Arguments
|
||||||
|
|
||||||
|
| Argument | Description | Default |
|
||||||
|
|----------|-------------|---------|
|
||||||
|
| `--base-url` | SGLang server URL | `http://localhost:30000` |
|
||||||
|
| `--model` | Model name on the server | `openai/whisper-large-v3` |
|
||||||
|
| `--dataset` | HuggingFace dataset for evaluation | `D4nt3/esb-datasets-earnings22-validation-tiny-filtered` |
|
||||||
|
| `--split` | Dataset split to use | `validation` |
|
||||||
|
| `--concurrency` | Number of concurrent requests | `4` |
|
||||||
|
| `--n-examples` | Number of examples to process (`-1` for all) | `-1` |
|
||||||
|
| `--output` | Path to save results as JSON | `None` |
|
||||||
|
| `--show-predictions` | Display sample predictions | `False` |
|
||||||
|
| `--print-n` | Number of samples to display | `5` |
|
||||||
|
| `--api-type` | API to use: `chat` (chat completions) or `transcription` (audio transcriptions) | `chat` |
|
||||||
|
| `--language` | Language for transcription API (e.g., `English`, `en`) | `None` |
|
||||||
|
| `--stream` | Enable streaming mode for transcription API | `False` |
|
||||||
|
|
||||||
|
## Metrics
|
||||||
|
|
||||||
|
The benchmark outputs:
|
||||||
|
|
||||||
|
| Metric | Description |
|
||||||
|
|--------|-------------|
|
||||||
|
| **Total Requests** | Number of successful ASR requests processed |
|
||||||
|
| **WER** | Word Error Rate (lower is better), computed using the `evaluate` library |
|
||||||
|
| **Average Latency** | Mean time per request (seconds) |
|
||||||
|
| **Median Latency** | 50th percentile latency (seconds) |
|
||||||
|
| **95th Latency** | 95th percentile latency (seconds) |
|
||||||
|
| **Throughput** | Requests processed per second |
|
||||||
|
| **Token Throughput** | Output tokens per second |
|
||||||
|
|
||||||
|
## Example Output
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python bench_sglang.py --api-type transcription --concurrency 128 --model openai/whisper-large-v3 --show-predictions
|
||||||
|
|
||||||
|
Loading dataset: D4nt3/esb-datasets-earnings22-validation-tiny-filtered...
|
||||||
|
Using API type: transcription
|
||||||
|
Repo card metadata block was not found. Setting CardData to empty.
|
||||||
|
WARNING:huggingface_hub.repocard:Repo card metadata block was not found. Setting CardData to empty.
|
||||||
|
Performing warmup...
|
||||||
|
Processing 511 samples...
|
||||||
|
------------------------------
|
||||||
|
Results for openai/whisper-large-v3:
|
||||||
|
Total Requests: 511
|
||||||
|
WER: 12.7690
|
||||||
|
Average Latency: 1.3602s
|
||||||
|
Median Latency: 1.2090s
|
||||||
|
95th Latency: 2.9986s
|
||||||
|
Throughput: 19.02 req/s
|
||||||
|
Token Throughput: 354.19 tok/s
|
||||||
|
Total Test Time: 26.8726s
|
||||||
|
------------------------------
|
||||||
|
|
||||||
|
==================== Sample Predictions ====================
|
||||||
|
Sample 1:
|
||||||
|
REF: on the use of taxonomy i you know i think it is it is early days for us to to make any clear indications to the market about the proportion that would fall under that requirement
|
||||||
|
PRED: on the eu taxonomy i think it is early days for us to make any clear indications to the market about the proportion that would fall under that requirement
|
||||||
|
----------------------------------------
|
||||||
|
Sample 2:
|
||||||
|
REF: so within fiscal year 2021 say 120 a 100 depending on what the micro will do and next year it is not necessarily payable in q one is we will look at what the cash flows for 2022 look like
|
||||||
|
PRED: so within fiscal year 2021 say $120000 $100000 depending on what the macro will do and next year it is not necessarily payable in q one is we will look at what the cash flows for 2022 look like
|
||||||
|
----------------------------------------
|
||||||
|
Sample 3:
|
||||||
|
REF: we talked about 4.7 gigawatts
|
||||||
|
PRED: we talked about 4.7 gigawatts
|
||||||
|
----------------------------------------
|
||||||
|
Sample 4:
|
||||||
|
REF: and you know depending on that working capital build we will we will see what that yields
|
||||||
|
PRED: and depending on that working capital build we will see what that yields what
|
||||||
|
----------------------------------------
|
||||||
|
Sample 5:
|
||||||
|
REF: so on on sinopec what we have agreed with sinopec way back then is that free cash flows after paying all capexs are distributed out 30 70%
|
||||||
|
PRED: so on sinopec what we have agreed with sinopec way back then is that free cash flows after paying all capexes are distributed out 30% 70%
|
||||||
|
----------------------------------------
|
||||||
|
============================================================
|
||||||
|
```
|
||||||
|
|
||||||
|
## Notes
|
||||||
|
|
||||||
|
- Audio samples longer than 30 seconds are automatically filtered out (Whisper limitation)
|
||||||
|
- The benchmark performs a warmup request before measuring performance
|
||||||
|
- Results are normalized using the model's tokenizer when available
|
||||||
|
- When using `--stream` with `--show-predictions`, use `--concurrency 1` for clean sequential output
|
||||||
|
- The `--language` option accepts both full names (e.g., `English`) and ISO 639-1 codes (e.g., `en`)
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
**Server connection refused**
|
||||||
|
- Ensure the SGLang server is running and accessible at the specified `--base-url`
|
||||||
|
- Check that the port is not blocked by a firewall
|
||||||
|
|
||||||
|
**Out of memory errors**
|
||||||
|
- Reduce `--concurrency` to lower GPU memory usage
|
||||||
|
- Use a smaller Whisper model variant
|
||||||
@@ -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)
|
||||||
@@ -545,6 +545,17 @@ class ModelConfig:
|
|||||||
if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures:
|
if "IQuestLoopCoderForCausalLM" in self.hf_config.architectures:
|
||||||
loop_num = getattr(self.hf_text_config, "loop_num", 1)
|
loop_num = getattr(self.hf_text_config, "loop_num", 1)
|
||||||
self.num_attention_layers = int(self.num_hidden_layers * int(loop_num))
|
self.num_attention_layers = int(self.num_hidden_layers * int(loop_num))
|
||||||
|
if "WhisperForConditionalGeneration" in self.hf_config.architectures:
|
||||||
|
# Whisper has unique layer ID scheme:
|
||||||
|
# - Encoder self-attention: 0 to encoder_layers-1 (no KV cache)
|
||||||
|
# - Decoder self-attention: encoder_layers to encoder_layers+decoder_layers-1 (uses KV cache)
|
||||||
|
# - Decoder cross-attention: encoder_layers+decoder_layers to encoder_layers+2*decoder_layers-1
|
||||||
|
# Even though cross-attention doesn't save KV cache, attention backend needs buffer to exist
|
||||||
|
encoder_layers = getattr(self.hf_text_config, "encoder_layers", 0)
|
||||||
|
decoder_layers = getattr(
|
||||||
|
self.hf_text_config, "decoder_layers", self.num_hidden_layers
|
||||||
|
)
|
||||||
|
self.num_attention_layers = encoder_layers + 2 * decoder_layers
|
||||||
self.num_nextn_predict_layers = getattr(
|
self.num_nextn_predict_layers = getattr(
|
||||||
self.hf_text_config, "num_nextn_predict_layers", None
|
self.hf_text_config, "num_nextn_predict_layers", None
|
||||||
)
|
)
|
||||||
@@ -1247,6 +1258,7 @@ multimodal_model_archs = [
|
|||||||
"InternS1ForConditionalGeneration",
|
"InternS1ForConditionalGeneration",
|
||||||
"InternS1ProForConditionalGeneration",
|
"InternS1ProForConditionalGeneration",
|
||||||
"Phi4MMForCausalLM",
|
"Phi4MMForCausalLM",
|
||||||
|
"WhisperForConditionalGeneration",
|
||||||
"Step3VLForConditionalGeneration",
|
"Step3VLForConditionalGeneration",
|
||||||
"POINTSV15ChatModel",
|
"POINTSV15ChatModel",
|
||||||
"DotsVLMForCausalLM",
|
"DotsVLMForCausalLM",
|
||||||
@@ -1285,11 +1297,17 @@ def is_image_gen_model(model_architectures: List[str]):
|
|||||||
|
|
||||||
|
|
||||||
def is_audio_model(model_architectures: List[str]):
|
def is_audio_model(model_architectures: List[str]):
|
||||||
return False
|
models = [
|
||||||
|
"WhisperForConditionalGeneration",
|
||||||
|
]
|
||||||
|
return any(model in model_architectures for model in models)
|
||||||
|
|
||||||
|
|
||||||
def is_encoder_decoder_model(model_architectures: List[str]):
|
def is_encoder_decoder_model(model_architectures: List[str]):
|
||||||
return "MllamaForConditionalGeneration" in model_architectures
|
models = [
|
||||||
|
"WhisperForConditionalGeneration",
|
||||||
|
]
|
||||||
|
return any(model in model_architectures for model in models)
|
||||||
|
|
||||||
|
|
||||||
def is_local_attention_model(model_architectures: List[str]):
|
def is_local_attention_model(model_architectures: List[str]):
|
||||||
|
|||||||
@@ -46,7 +46,7 @@ import orjson
|
|||||||
import requests
|
import requests
|
||||||
import uvicorn
|
import uvicorn
|
||||||
import uvloop
|
import uvloop
|
||||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
from fastapi import Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
|
||||||
from fastapi.exceptions import RequestValidationError
|
from fastapi.exceptions import RequestValidationError
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
||||||
@@ -93,6 +93,9 @@ from sglang.srt.entrypoints.openai.serving_tokenize import (
|
|||||||
OpenAIServingDetokenize,
|
OpenAIServingDetokenize,
|
||||||
OpenAIServingTokenize,
|
OpenAIServingTokenize,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.entrypoints.openai.serving_transcription import (
|
||||||
|
OpenAIServingTranscription,
|
||||||
|
)
|
||||||
from sglang.srt.entrypoints.warmup import execute_warmups
|
from sglang.srt.entrypoints.warmup import execute_warmups
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||||
@@ -298,6 +301,9 @@ async def lifespan(fast_api_app: FastAPI):
|
|||||||
fast_api_app.state.openai_serving_detokenize = OpenAIServingDetokenize(
|
fast_api_app.state.openai_serving_detokenize = OpenAIServingDetokenize(
|
||||||
_global_state.tokenizer_manager
|
_global_state.tokenizer_manager
|
||||||
)
|
)
|
||||||
|
fast_api_app.state.openai_serving_transcription = OpenAIServingTranscription(
|
||||||
|
_global_state.tokenizer_manager
|
||||||
|
)
|
||||||
|
|
||||||
# Initialize Ollama-compatible serving handler
|
# Initialize Ollama-compatible serving handler
|
||||||
fast_api_app.state.ollama_serving = OllamaServing(_global_state.tokenizer_manager)
|
fast_api_app.state.ollama_serving = OllamaServing(_global_state.tokenizer_manager)
|
||||||
@@ -1418,6 +1424,38 @@ async def openai_v1_detokenize(request: DetokenizeRequest, raw_request: Request)
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/v1/audio/transcriptions")
|
||||||
|
async def openai_v1_audio_transcriptions(
|
||||||
|
raw_request: Request,
|
||||||
|
file: UploadFile = File(...),
|
||||||
|
model: str = Form(default="default"),
|
||||||
|
language: Optional[str] = Form(default=None),
|
||||||
|
response_format: str = Form(default="json"),
|
||||||
|
temperature: float = Form(default=0.0),
|
||||||
|
stream: bool = Form(default=False),
|
||||||
|
):
|
||||||
|
"""OpenAI-compatible audio transcription endpoint."""
|
||||||
|
if response_format not in ["json", "text"]:
|
||||||
|
return ORJSONResponse(
|
||||||
|
content={"error": {"message": "Only 'json' and 'text' formats supported"}},
|
||||||
|
status_code=400,
|
||||||
|
)
|
||||||
|
|
||||||
|
audio_data = await file.read()
|
||||||
|
|
||||||
|
return (
|
||||||
|
await raw_request.app.state.openai_serving_transcription.create_transcription(
|
||||||
|
audio_data=audio_data,
|
||||||
|
model=model,
|
||||||
|
language=language,
|
||||||
|
response_format=response_format,
|
||||||
|
temperature=temperature,
|
||||||
|
stream=stream,
|
||||||
|
raw_request=raw_request,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@app.get("/v1/models", response_class=ORJSONResponse)
|
@app.get("/v1/models", response_class=ORJSONResponse)
|
||||||
async def available_models():
|
async def available_models():
|
||||||
"""Show available models. OpenAI-compatible endpoint."""
|
"""Show available models. OpenAI-compatible endpoint."""
|
||||||
|
|||||||
@@ -1389,3 +1389,51 @@ class ResponseReasoningTextContent(BaseModel):
|
|||||||
ResponseInputOutputItem: TypeAlias = Union[
|
ResponseInputOutputItem: TypeAlias = Union[
|
||||||
ResponseInputItemParam, "ResponseReasoningItem", ResponseFunctionToolCall
|
ResponseInputItemParam, "ResponseReasoningItem", ResponseFunctionToolCall
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
# ================== Transcription API Protocol Definitions ==================
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionRequest(BaseModel):
|
||||||
|
"""Request model for audio transcription (OpenAI-compatible)."""
|
||||||
|
|
||||||
|
model: str = DEFAULT_MODEL_NAME
|
||||||
|
language: Optional[str] = None
|
||||||
|
response_format: str = "json"
|
||||||
|
temperature: float = 0.0
|
||||||
|
stream: bool = False
|
||||||
|
# Internal fields (not from API)
|
||||||
|
audio_data: Optional[bytes] = None
|
||||||
|
audio_duration_s: float = 0.0
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionUsage(BaseModel):
|
||||||
|
"""Usage info for transcription response (duration-based)."""
|
||||||
|
|
||||||
|
type: Literal["duration"] = "duration"
|
||||||
|
seconds: int # Audio duration in seconds (rounded up)
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionResponse(BaseModel):
|
||||||
|
"""Non-streaming transcription response (OpenAI-compatible)."""
|
||||||
|
|
||||||
|
text: str
|
||||||
|
usage: Optional[TranscriptionUsage] = None
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionStreamChoice(BaseModel):
|
||||||
|
"""Delta content for streaming transcription."""
|
||||||
|
|
||||||
|
delta: DeltaMessage
|
||||||
|
finish_reason: Optional[str] = None
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionStreamResponse(BaseModel):
|
||||||
|
"""Streaming transcription chunk (OpenAI-compatible)."""
|
||||||
|
|
||||||
|
id: str = Field(default_factory=lambda: f"trsc-{uuid.uuid4().hex}")
|
||||||
|
object: Literal["transcription.chunk"] = "transcription.chunk"
|
||||||
|
created: int = Field(default_factory=lambda: int(time.time()))
|
||||||
|
model: str
|
||||||
|
choices: List[TranscriptionStreamChoice]
|
||||||
|
usage: Optional[UsageInfo] = None
|
||||||
|
|||||||
@@ -0,0 +1,220 @@
|
|||||||
|
# Copyright 2025 SGLang Team
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
# ==============================================================================
|
||||||
|
"""
|
||||||
|
OpenAI-compatible transcription endpoint handler for Whisper models.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import io
|
||||||
|
import logging
|
||||||
|
import math
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from typing import TYPE_CHECKING, AsyncGenerator, Optional, Union
|
||||||
|
|
||||||
|
from fastapi import Request
|
||||||
|
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
||||||
|
|
||||||
|
from sglang.srt.entrypoints.openai.protocol import (
|
||||||
|
DeltaMessage,
|
||||||
|
ErrorResponse,
|
||||||
|
TranscriptionRequest,
|
||||||
|
TranscriptionResponse,
|
||||||
|
TranscriptionStreamChoice,
|
||||||
|
TranscriptionStreamResponse,
|
||||||
|
TranscriptionUsage,
|
||||||
|
)
|
||||||
|
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
|
||||||
|
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class OpenAIServingTranscription(OpenAIServingBase):
|
||||||
|
"""Handler for /v1/audio/transcriptions requests"""
|
||||||
|
|
||||||
|
def __init__(self, tokenizer_manager: TokenizerManager):
|
||||||
|
super().__init__(tokenizer_manager)
|
||||||
|
|
||||||
|
def _request_id_prefix(self) -> str:
|
||||||
|
return "trsc-"
|
||||||
|
|
||||||
|
def _validate_request(self, request: TranscriptionRequest) -> Optional[str]:
|
||||||
|
"""Validate transcription request."""
|
||||||
|
# Validation is done in the route handler for form data
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _convert_to_internal_request(
|
||||||
|
self,
|
||||||
|
request: TranscriptionRequest,
|
||||||
|
raw_request: Request = None,
|
||||||
|
) -> tuple[GenerateReqInput, TranscriptionRequest]:
|
||||||
|
"""Convert transcription request to internal format."""
|
||||||
|
# Build sampling params - include language for WhisperProcessor
|
||||||
|
sampling_params = {
|
||||||
|
"temperature": request.temperature,
|
||||||
|
"max_new_tokens": 448, # Whisper default max tokens
|
||||||
|
"language": request.language, # Pass to WhisperProcessor for language-specific decoding
|
||||||
|
}
|
||||||
|
|
||||||
|
# For Whisper, we pass audio_data and let the processor handle it
|
||||||
|
adapted_request = GenerateReqInput(
|
||||||
|
text="", # Empty text - Whisper processor will set proper decoder tokens
|
||||||
|
audio_data=request.audio_data,
|
||||||
|
sampling_params=sampling_params,
|
||||||
|
stream=request.stream,
|
||||||
|
modalities=["audio"],
|
||||||
|
routing_key=self.extract_routing_key(raw_request),
|
||||||
|
)
|
||||||
|
|
||||||
|
return adapted_request, request
|
||||||
|
|
||||||
|
def _get_audio_duration(self, audio_data: bytes) -> float:
|
||||||
|
"""Calculate audio duration in seconds."""
|
||||||
|
try:
|
||||||
|
import soundfile as sf
|
||||||
|
|
||||||
|
audio_array, sr = sf.read(io.BytesIO(audio_data))
|
||||||
|
duration = len(audio_array) / sr
|
||||||
|
return duration
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Could not calculate audio duration: {e}")
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
async def create_transcription(
|
||||||
|
self,
|
||||||
|
audio_data: bytes,
|
||||||
|
model: str,
|
||||||
|
language: Optional[str],
|
||||||
|
response_format: str,
|
||||||
|
temperature: float,
|
||||||
|
stream: bool,
|
||||||
|
raw_request: Request,
|
||||||
|
) -> Union[TranscriptionResponse, StreamingResponse, Response, ORJSONResponse]:
|
||||||
|
"""Main entry point for transcription requests."""
|
||||||
|
# Calculate audio duration for usage reporting
|
||||||
|
audio_duration_s = self._get_audio_duration(audio_data)
|
||||||
|
|
||||||
|
# Build request
|
||||||
|
request = TranscriptionRequest(
|
||||||
|
audio_data=audio_data,
|
||||||
|
model=model,
|
||||||
|
language=language,
|
||||||
|
response_format=response_format,
|
||||||
|
temperature=temperature,
|
||||||
|
stream=stream,
|
||||||
|
audio_duration_s=audio_duration_s,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Use the base class handle_request pattern
|
||||||
|
return await self.handle_request(request, raw_request)
|
||||||
|
|
||||||
|
async def _handle_non_streaming_request(
|
||||||
|
self,
|
||||||
|
adapted_request: GenerateReqInput,
|
||||||
|
request: TranscriptionRequest,
|
||||||
|
raw_request: Request,
|
||||||
|
) -> Union[TranscriptionResponse, ErrorResponse, ORJSONResponse, Response]:
|
||||||
|
"""Handle non-streaming transcription request."""
|
||||||
|
try:
|
||||||
|
ret = await self.tokenizer_manager.generate_request(
|
||||||
|
adapted_request, raw_request
|
||||||
|
).__anext__()
|
||||||
|
except ValueError as e:
|
||||||
|
return self.create_error_response(str(e))
|
||||||
|
|
||||||
|
text = ret.get("text", "")
|
||||||
|
|
||||||
|
# Build response based on format
|
||||||
|
if request.response_format == "text":
|
||||||
|
return Response(content=text, media_type="text/plain")
|
||||||
|
|
||||||
|
# JSON format
|
||||||
|
usage = TranscriptionUsage(seconds=int(math.ceil(request.audio_duration_s)))
|
||||||
|
|
||||||
|
return TranscriptionResponse(text=text, usage=usage)
|
||||||
|
|
||||||
|
async def _handle_streaming_request(
|
||||||
|
self,
|
||||||
|
adapted_request: GenerateReqInput,
|
||||||
|
request: TranscriptionRequest,
|
||||||
|
raw_request: Request,
|
||||||
|
) -> StreamingResponse:
|
||||||
|
"""Handle streaming transcription request."""
|
||||||
|
return StreamingResponse(
|
||||||
|
self._generate_transcription_stream(adapted_request, request, raw_request),
|
||||||
|
media_type="text/event-stream",
|
||||||
|
background=self.tokenizer_manager.create_abort_task(adapted_request),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _generate_transcription_stream(
|
||||||
|
self,
|
||||||
|
adapted_request: GenerateReqInput,
|
||||||
|
request: TranscriptionRequest,
|
||||||
|
raw_request: Request,
|
||||||
|
) -> AsyncGenerator[str, None]:
|
||||||
|
"""Generate streaming transcription response."""
|
||||||
|
created_time = int(time.time())
|
||||||
|
request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
|
||||||
|
model = request.model
|
||||||
|
stream_buffer = ""
|
||||||
|
|
||||||
|
try:
|
||||||
|
async for content in self.tokenizer_manager.generate_request(
|
||||||
|
adapted_request, raw_request
|
||||||
|
):
|
||||||
|
finish_reason = content["meta_info"]["finish_reason"]
|
||||||
|
finish_reason_type = finish_reason["type"] if finish_reason else None
|
||||||
|
|
||||||
|
# Calculate delta (new text since last chunk)
|
||||||
|
current_text = content.get("text", "")
|
||||||
|
delta = current_text[len(stream_buffer) :]
|
||||||
|
stream_buffer = current_text
|
||||||
|
|
||||||
|
# Send content delta if there's new text
|
||||||
|
if delta:
|
||||||
|
choice_data = TranscriptionStreamChoice(
|
||||||
|
delta=DeltaMessage(content=delta),
|
||||||
|
finish_reason=None,
|
||||||
|
)
|
||||||
|
chunk = TranscriptionStreamResponse(
|
||||||
|
id=request_id,
|
||||||
|
created=created_time,
|
||||||
|
model=model,
|
||||||
|
choices=[choice_data],
|
||||||
|
)
|
||||||
|
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||||
|
|
||||||
|
# Send finish reason when done
|
||||||
|
if finish_reason_type:
|
||||||
|
choice_data = TranscriptionStreamChoice(
|
||||||
|
delta=DeltaMessage(),
|
||||||
|
finish_reason=finish_reason_type,
|
||||||
|
)
|
||||||
|
chunk = TranscriptionStreamResponse(
|
||||||
|
id=request_id,
|
||||||
|
created=created_time,
|
||||||
|
model=model,
|
||||||
|
choices=[choice_data],
|
||||||
|
)
|
||||||
|
yield f"data: {chunk.model_dump_json()}\n\n"
|
||||||
|
|
||||||
|
except ValueError as e:
|
||||||
|
error = self.create_streaming_error_response(str(e))
|
||||||
|
yield f"data: {error}\n\n"
|
||||||
|
|
||||||
|
yield "data: [DONE]\n\n"
|
||||||
@@ -694,9 +694,15 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
|
|||||||
"the engine with skip_tokenizer_init=False."
|
"the engine with skip_tokenizer_init=False."
|
||||||
)
|
)
|
||||||
|
|
||||||
input_ids, token_type_ids = await self._tokenize_texts(
|
# For audio-only requests (e.g., Whisper), text may be empty.
|
||||||
input_text, is_cross_encoder_request
|
# The multimodal processor will provide input_ids later.
|
||||||
)
|
if not input_text and self.mm_processor and obj.contains_mm_input():
|
||||||
|
# Use empty placeholder - multimodal processor will override
|
||||||
|
input_ids = []
|
||||||
|
else:
|
||||||
|
input_ids, token_type_ids = await self._tokenize_texts(
|
||||||
|
input_text, is_cross_encoder_request
|
||||||
|
)
|
||||||
|
|
||||||
if self.mm_processor and obj.contains_mm_input():
|
if self.mm_processor and obj.contains_mm_input():
|
||||||
if obj.image_data is not None and not isinstance(obj.image_data, list):
|
if obj.image_data is not None and not isinstance(obj.image_data, list):
|
||||||
|
|||||||
@@ -0,0 +1,543 @@
|
|||||||
|
from typing import Any, Iterable, List, Optional, Tuple
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from transformers import WhisperConfig
|
||||||
|
|
||||||
|
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
||||||
|
from sglang.srt.layers.activation import get_act_fn
|
||||||
|
from sglang.srt.layers.linear import (
|
||||||
|
ColumnParallelLinear,
|
||||||
|
QKVParallelLinear,
|
||||||
|
RowParallelLinear,
|
||||||
|
)
|
||||||
|
from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
|
||||||
|
from sglang.srt.layers.quantization import QuantizationConfig
|
||||||
|
from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
|
||||||
|
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
|
||||||
|
from sglang.srt.managers.schedule_batch import MultimodalInputs
|
||||||
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||||
|
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperAttention(torch.nn.Module):
|
||||||
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
embed_dim: int,
|
||||||
|
num_heads: int,
|
||||||
|
bias: bool = True,
|
||||||
|
layer_id: Optional[int] = None,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
is_cross_attention: bool = False,
|
||||||
|
is_encoder=False,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.total_num_heads = num_heads
|
||||||
|
head_dim = embed_dim // num_heads
|
||||||
|
self.is_cross_attention = is_cross_attention
|
||||||
|
self.is_encoder = is_encoder
|
||||||
|
|
||||||
|
tp_size = get_tensor_model_parallel_world_size()
|
||||||
|
assert (
|
||||||
|
num_heads % tp_size == 0
|
||||||
|
), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})"
|
||||||
|
self.num_heads = num_heads // tp_size
|
||||||
|
|
||||||
|
if (head_dim * num_heads) != embed_dim:
|
||||||
|
raise ValueError(
|
||||||
|
f"embed_dim must be divisible by num_heads (got `embed_dim`: {embed_dim}"
|
||||||
|
f" and `num_heads`: {num_heads})."
|
||||||
|
)
|
||||||
|
self.scaling = head_dim**-0.5
|
||||||
|
self.head_dim = head_dim
|
||||||
|
self.kv_size = self.num_heads * head_dim
|
||||||
|
|
||||||
|
if is_cross_attention:
|
||||||
|
self.q_proj = ColumnParallelLinear(
|
||||||
|
embed_dim, embed_dim, quant_config=quant_config
|
||||||
|
)
|
||||||
|
self.kv_proj = QKVParallelLinear(
|
||||||
|
hidden_size=embed_dim,
|
||||||
|
head_size=head_dim,
|
||||||
|
total_num_heads=0,
|
||||||
|
total_num_kv_heads=num_heads,
|
||||||
|
bias=bias,
|
||||||
|
quant_config=quant_config,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.qkv_proj = QKVParallelLinear(
|
||||||
|
embed_dim, head_dim, num_heads, quant_config=quant_config
|
||||||
|
)
|
||||||
|
self.out_proj = RowParallelLinear(
|
||||||
|
embed_dim, embed_dim, bias=bias, quant_config=quant_config
|
||||||
|
)
|
||||||
|
self.attn = RadixAttention(
|
||||||
|
self.num_heads,
|
||||||
|
head_dim,
|
||||||
|
scaling=1.0,
|
||||||
|
num_kv_heads=self.num_heads,
|
||||||
|
layer_id=layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
is_cross_attention=is_cross_attention,
|
||||||
|
attn_type=(
|
||||||
|
AttentionType.ENCODER_ONLY if is_encoder else AttentionType.DECODER
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
cross_hidden_states: Optional[torch.Tensor] = None,
|
||||||
|
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
|
||||||
|
"""Input shape: Batch x Time x Channel"""
|
||||||
|
|
||||||
|
if self.is_cross_attention:
|
||||||
|
q, _ = self.q_proj(hidden_states)
|
||||||
|
if cross_hidden_states is not None:
|
||||||
|
kv, _ = self.kv_proj(cross_hidden_states)
|
||||||
|
k, v = kv.split([self.kv_size, self.kv_size], dim=-1)
|
||||||
|
else:
|
||||||
|
k = torch.zeros_like(q)
|
||||||
|
v = torch.zeros_like(q)
|
||||||
|
|
||||||
|
q = q * self.scaling
|
||||||
|
num_heads = self.attn.tp_q_head_num
|
||||||
|
head_dim = self.attn.head_dim
|
||||||
|
|
||||||
|
q = q.view(-1, num_heads, head_dim)
|
||||||
|
k = k.view(-1, num_heads, head_dim)
|
||||||
|
v = v.view(-1, num_heads, head_dim)
|
||||||
|
|
||||||
|
q_len = q.shape[0]
|
||||||
|
kv_len = k.shape[0]
|
||||||
|
|
||||||
|
q = q.transpose(0, 1)
|
||||||
|
k = k.transpose(0, 1)
|
||||||
|
v = v.transpose(0, 1)
|
||||||
|
|
||||||
|
attn_weights = torch.bmm(q, k.transpose(1, 2))
|
||||||
|
|
||||||
|
# Apply block-diagonal mask for batched cross-attention
|
||||||
|
batch_size = forward_batch.batch_size if forward_batch else 1
|
||||||
|
if batch_size > 1 and kv_len > 0:
|
||||||
|
encoder_len_per_request = kv_len // batch_size
|
||||||
|
if encoder_len_per_request * batch_size == kv_len:
|
||||||
|
is_decode = forward_batch.forward_mode.is_decode()
|
||||||
|
if is_decode:
|
||||||
|
mask = torch.zeros(
|
||||||
|
(q_len, kv_len), device=q.device, dtype=torch.bool
|
||||||
|
)
|
||||||
|
for i in range(batch_size):
|
||||||
|
enc_start = i * encoder_len_per_request
|
||||||
|
enc_end = (i + 1) * encoder_len_per_request
|
||||||
|
mask[i, enc_start:enc_end] = True
|
||||||
|
attn_weights = attn_weights.masked_fill(
|
||||||
|
~mask.unsqueeze(0), float("-inf")
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
seq_lens = forward_batch.seq_lens
|
||||||
|
if seq_lens is not None and len(seq_lens) == batch_size:
|
||||||
|
seq_lens_list = seq_lens.tolist()
|
||||||
|
mask = torch.zeros(
|
||||||
|
(q_len, kv_len), device=q.device, dtype=torch.bool
|
||||||
|
)
|
||||||
|
q_start = 0
|
||||||
|
for i, dec_len in enumerate(seq_lens_list):
|
||||||
|
enc_start = i * encoder_len_per_request
|
||||||
|
enc_end = (i + 1) * encoder_len_per_request
|
||||||
|
q_end = q_start + dec_len
|
||||||
|
mask[q_start:q_end, enc_start:enc_end] = True
|
||||||
|
q_start = q_end
|
||||||
|
attn_weights = attn_weights.masked_fill(
|
||||||
|
~mask.unsqueeze(0), float("-inf")
|
||||||
|
)
|
||||||
|
|
||||||
|
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
|
||||||
|
attn_output = torch.bmm(attn_weights, v)
|
||||||
|
attn_output = attn_output.transpose(0, 1)
|
||||||
|
attn_output = attn_output.reshape(q_len, num_heads * head_dim)
|
||||||
|
else:
|
||||||
|
qkv, _ = self.qkv_proj(hidden_states)
|
||||||
|
q, k, v = qkv.chunk(chunks=3, dim=-1)
|
||||||
|
q = q * self.scaling
|
||||||
|
|
||||||
|
if self.is_encoder:
|
||||||
|
num_heads = self.attn.tp_q_head_num
|
||||||
|
head_dim = self.attn.head_dim
|
||||||
|
batch_size, seq_len, _ = hidden_states.shape
|
||||||
|
|
||||||
|
q = q.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||||
|
k = k.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||||
|
v = v.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
|
||||||
|
|
||||||
|
attn_output = torch.nn.functional.scaled_dot_product_attention(
|
||||||
|
q, k, v, scale=1.0
|
||||||
|
)
|
||||||
|
attn_output = attn_output.permute(0, 2, 1, 3).reshape(
|
||||||
|
batch_size, seq_len, num_heads * head_dim
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=True)
|
||||||
|
|
||||||
|
attn_output, _ = self.out_proj(attn_output)
|
||||||
|
|
||||||
|
return attn_output
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperEncoderLayer(torch.nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: WhisperConfig,
|
||||||
|
layer_id: Optional[int] = None,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.embed_dim = config.d_model
|
||||||
|
|
||||||
|
self.self_attn = WhisperAttention(
|
||||||
|
embed_dim=self.embed_dim,
|
||||||
|
num_heads=config.encoder_attention_heads,
|
||||||
|
layer_id=layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
is_encoder=True,
|
||||||
|
)
|
||||||
|
self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||||
|
|
||||||
|
self.activation_fn = get_act_fn(
|
||||||
|
config.activation_function, quant_config=quant_config
|
||||||
|
)
|
||||||
|
|
||||||
|
self.fc1 = ColumnParallelLinear(self.embed_dim, config.encoder_ffn_dim)
|
||||||
|
self.fc2 = RowParallelLinear(config.encoder_ffn_dim, self.embed_dim)
|
||||||
|
self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
hidden_states: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
||||||
|
hidden_states = self.self_attn(hidden_states, forward_batch)
|
||||||
|
|
||||||
|
hidden_states = residual + hidden_states
|
||||||
|
|
||||||
|
residual = hidden_states
|
||||||
|
hidden_states = self.final_layer_norm(hidden_states)
|
||||||
|
hidden_states, _ = self.fc1(hidden_states)
|
||||||
|
hidden_states = self.activation_fn(hidden_states)
|
||||||
|
|
||||||
|
hidden_states, _ = self.fc2(hidden_states)
|
||||||
|
|
||||||
|
hidden_states = residual + hidden_states
|
||||||
|
|
||||||
|
if hidden_states.dtype == torch.float16:
|
||||||
|
clamp_value = torch.finfo(hidden_states.dtype).max - 1000
|
||||||
|
hidden_states = torch.clamp(
|
||||||
|
hidden_states, min=-clamp_value, max=clamp_value
|
||||||
|
)
|
||||||
|
return hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperDecoderLayer(torch.nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: WhisperConfig,
|
||||||
|
layer_id: Optional[int] = None,
|
||||||
|
quant_config: Optional[QuantizationConfig] = None,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.embed_dim = config.d_model
|
||||||
|
|
||||||
|
# Offset decoder layer IDs to avoid overlap with encoder layers
|
||||||
|
decoder_self_attn_layer_id = config.encoder_layers + layer_id
|
||||||
|
decoder_cross_attn_layer_id = (
|
||||||
|
config.encoder_layers + config.decoder_layers + layer_id
|
||||||
|
)
|
||||||
|
|
||||||
|
self.self_attn = WhisperAttention(
|
||||||
|
embed_dim=self.embed_dim,
|
||||||
|
num_heads=config.decoder_attention_heads,
|
||||||
|
layer_id=decoder_self_attn_layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.activation_fn = get_act_fn(
|
||||||
|
config.activation_function, quant_config=quant_config
|
||||||
|
)
|
||||||
|
|
||||||
|
self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||||
|
self.encoder_attn = WhisperAttention(
|
||||||
|
embed_dim=self.embed_dim,
|
||||||
|
num_heads=config.decoder_attention_heads,
|
||||||
|
layer_id=decoder_cross_attn_layer_id,
|
||||||
|
quant_config=quant_config,
|
||||||
|
is_cross_attention=True,
|
||||||
|
)
|
||||||
|
self.encoder_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||||
|
self.fc1 = ColumnParallelLinear(self.embed_dim, config.decoder_ffn_dim)
|
||||||
|
self.fc2 = RowParallelLinear(config.decoder_ffn_dim, self.embed_dim)
|
||||||
|
self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
decoder_hidden_states: torch.Tensor,
|
||||||
|
encoder_hidden_states: Optional[torch.Tensor],
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
|
||||||
|
residual = decoder_hidden_states
|
||||||
|
decoder_hidden_states = self.self_attn_layer_norm(decoder_hidden_states)
|
||||||
|
decoder_hidden_states = self.self_attn(decoder_hidden_states, forward_batch)
|
||||||
|
decoder_hidden_states = residual + decoder_hidden_states
|
||||||
|
|
||||||
|
residual = decoder_hidden_states
|
||||||
|
decoder_hidden_states = self.encoder_attn_layer_norm(decoder_hidden_states)
|
||||||
|
decoder_hidden_states = self.encoder_attn(
|
||||||
|
decoder_hidden_states, forward_batch, encoder_hidden_states
|
||||||
|
)
|
||||||
|
decoder_hidden_states = residual + decoder_hidden_states
|
||||||
|
|
||||||
|
residual = decoder_hidden_states
|
||||||
|
decoder_hidden_states = self.final_layer_norm(decoder_hidden_states)
|
||||||
|
decoder_hidden_states, _ = self.fc1(decoder_hidden_states)
|
||||||
|
decoder_hidden_states = self.activation_fn(decoder_hidden_states)
|
||||||
|
decoder_hidden_states, _ = self.fc2(decoder_hidden_states)
|
||||||
|
|
||||||
|
decoder_hidden_states = residual + decoder_hidden_states
|
||||||
|
|
||||||
|
return decoder_hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperEncoder(torch.nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
embed_dim = config.d_model
|
||||||
|
self.embed_scale = embed_dim**-0.5 if config.scale_embedding else 1.0
|
||||||
|
|
||||||
|
self.conv1 = torch.nn.Conv1d(
|
||||||
|
config.num_mel_bins, embed_dim, kernel_size=3, padding=1
|
||||||
|
)
|
||||||
|
self.conv2 = torch.nn.Conv1d(
|
||||||
|
embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
|
||||||
|
)
|
||||||
|
self.embed_positions = torch.nn.Embedding(
|
||||||
|
config.max_source_positions, embed_dim
|
||||||
|
)
|
||||||
|
|
||||||
|
self.layers = torch.nn.ModuleList(
|
||||||
|
[
|
||||||
|
WhisperEncoderLayer(config, id, quant_config)
|
||||||
|
for id in range(config.encoder_layers)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
self.layer_norm = torch.nn.LayerNorm(config.d_model)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_features: torch.Tensor,
|
||||||
|
position_ids: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
):
|
||||||
|
inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
|
||||||
|
inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))
|
||||||
|
|
||||||
|
inputs_embeds = inputs_embeds.mT
|
||||||
|
|
||||||
|
hidden_states = inputs_embeds + self.embed_positions(position_ids)
|
||||||
|
|
||||||
|
for encoder_layer in self.layers:
|
||||||
|
hidden_states = encoder_layer(hidden_states, forward_batch)
|
||||||
|
|
||||||
|
hidden_states = self.layer_norm(hidden_states)
|
||||||
|
return hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperDecoder(torch.nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.max_target_positions = config.max_target_positions
|
||||||
|
self.max_source_positions = config.max_source_positions
|
||||||
|
self.embed_scale = config.d_model**-0.5 if config.scale_embedding else 1.0
|
||||||
|
|
||||||
|
self.embed_tokens = torch.nn.Embedding(
|
||||||
|
config.vocab_size, config.d_model, padding_idx=config.pad_token_id
|
||||||
|
)
|
||||||
|
self.embed_positions = torch.nn.Embedding(
|
||||||
|
self.max_target_positions, config.d_model
|
||||||
|
)
|
||||||
|
|
||||||
|
self.layers = torch.nn.ModuleList(
|
||||||
|
[
|
||||||
|
WhisperDecoderLayer(config, layer_idx, quant_config)
|
||||||
|
for layer_idx in range(config.decoder_layers)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.layer_norm = torch.nn.LayerNorm(config.d_model)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
encoder_hidden_states: Optional[torch.Tensor],
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
position_ids=None,
|
||||||
|
):
|
||||||
|
inputs_embeds = self.embed_tokens(input_ids)
|
||||||
|
positions = self.embed_positions(position_ids)
|
||||||
|
hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
|
||||||
|
|
||||||
|
for decoder_layer in self.layers:
|
||||||
|
hidden_states = decoder_layer(
|
||||||
|
hidden_states, encoder_hidden_states, forward_batch
|
||||||
|
)
|
||||||
|
|
||||||
|
hidden_states = self.layer_norm(hidden_states)
|
||||||
|
|
||||||
|
return hidden_states
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperForConditionalGeneration(torch.nn.Module):
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.encoder = WhisperEncoder(config, quant_config)
|
||||||
|
self.decoder = WhisperDecoder(config, quant_config)
|
||||||
|
self.proj_out = ParallelLMHead(
|
||||||
|
config.vocab_size, config.d_model, quant_config=quant_config
|
||||||
|
)
|
||||||
|
self.logits_processor = LogitsProcessor(config)
|
||||||
|
self.config = config
|
||||||
|
self._encoder_cache = {}
|
||||||
|
|
||||||
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
|
||||||
|
stacked_params_mapping = [
|
||||||
|
(".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
|
||||||
|
(".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
|
||||||
|
(".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
|
||||||
|
(".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
|
||||||
|
(".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
|
||||||
|
]
|
||||||
|
|
||||||
|
params_dict = dict(self.named_parameters())
|
||||||
|
weights_dict = dict(weights)
|
||||||
|
|
||||||
|
# Whisper has no k_proj bias, create zeros
|
||||||
|
for layer_idx in range(self.config.decoder_layers):
|
||||||
|
layer_prefix = f"model.decoder.layers.{layer_idx}.encoder_attn."
|
||||||
|
k_proj_key = layer_prefix + "k_proj.weight"
|
||||||
|
if k_proj_key in weights_dict:
|
||||||
|
k_proj_weight = weights_dict[k_proj_key]
|
||||||
|
bias_key = layer_prefix + "k_proj.bias"
|
||||||
|
if bias_key not in weights_dict:
|
||||||
|
weights_dict[bias_key] = torch.zeros(k_proj_weight.size(0))
|
||||||
|
|
||||||
|
weights_dict["proj_out.weight"] = weights_dict[
|
||||||
|
"model.decoder.embed_tokens.weight"
|
||||||
|
]
|
||||||
|
|
||||||
|
for name, loaded_weight in weights_dict.items():
|
||||||
|
name = name.replace("model.", "")
|
||||||
|
|
||||||
|
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||||
|
if weight_name not in name:
|
||||||
|
continue
|
||||||
|
name = name.replace(weight_name, param_name)
|
||||||
|
if name not in params_dict:
|
||||||
|
break
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = param.weight_loader
|
||||||
|
weight_loader(param, loaded_weight, shard_id)
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
if name not in params_dict:
|
||||||
|
continue
|
||||||
|
param = params_dict[name]
|
||||||
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||||
|
weight_loader(param, loaded_weight)
|
||||||
|
|
||||||
|
def pad_input_ids(self, input_ids: List[int], _mm_inputs: MultimodalInputs):
|
||||||
|
return input_ids
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
input_ids: torch.Tensor,
|
||||||
|
positions: torch.Tensor,
|
||||||
|
forward_batch: ForwardBatch,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> LogitsProcessorOutput:
|
||||||
|
dtype = self.encoder.conv1.weight.dtype
|
||||||
|
is_decode = forward_batch.forward_mode.is_decode()
|
||||||
|
|
||||||
|
if is_decode:
|
||||||
|
encoder_outputs = None
|
||||||
|
if forward_batch.req_pool_indices is not None:
|
||||||
|
req_indices = forward_batch.req_pool_indices.tolist()
|
||||||
|
encoder_list = []
|
||||||
|
for req_idx in req_indices:
|
||||||
|
if req_idx in self._encoder_cache:
|
||||||
|
encoder_list.append(self._encoder_cache[req_idx])
|
||||||
|
if encoder_list:
|
||||||
|
encoder_outputs = torch.cat(encoder_list, dim=0)
|
||||||
|
else:
|
||||||
|
encoder_list = []
|
||||||
|
mm_inputs_list = forward_batch.mm_inputs if forward_batch.mm_inputs else []
|
||||||
|
req_indices = (
|
||||||
|
forward_batch.req_pool_indices.tolist()
|
||||||
|
if forward_batch.req_pool_indices is not None
|
||||||
|
else []
|
||||||
|
)
|
||||||
|
|
||||||
|
for req_idx, mm_input in zip(req_indices, mm_inputs_list):
|
||||||
|
if mm_input is None or not mm_input.mm_items:
|
||||||
|
continue
|
||||||
|
|
||||||
|
features = mm_input.mm_items[0].feature
|
||||||
|
if features.ndim == 2:
|
||||||
|
features = features.unsqueeze(0)
|
||||||
|
|
||||||
|
encoder_len = features.shape[-1] // 2
|
||||||
|
encoder_position_ids = torch.arange(encoder_len).to(
|
||||||
|
features.device, non_blocking=True
|
||||||
|
)
|
||||||
|
|
||||||
|
req_encoder_outputs = self.encoder(
|
||||||
|
features.to(dtype), encoder_position_ids, forward_batch
|
||||||
|
)
|
||||||
|
req_encoder_outputs = req_encoder_outputs.squeeze(0)
|
||||||
|
|
||||||
|
self._encoder_cache[req_idx] = req_encoder_outputs
|
||||||
|
encoder_list.append(req_encoder_outputs)
|
||||||
|
|
||||||
|
if encoder_list:
|
||||||
|
encoder_outputs = torch.cat(encoder_list, dim=0)
|
||||||
|
else:
|
||||||
|
encoder_outputs = None
|
||||||
|
|
||||||
|
decoder_outputs = self.decoder(
|
||||||
|
input_ids, encoder_outputs, forward_batch, positions
|
||||||
|
)
|
||||||
|
|
||||||
|
logits = self.logits_processor(
|
||||||
|
input_ids=input_ids,
|
||||||
|
lm_head=self.proj_out,
|
||||||
|
hidden_states=decoder_outputs,
|
||||||
|
logits_metadata=forward_batch,
|
||||||
|
)
|
||||||
|
|
||||||
|
return logits
|
||||||
|
|
||||||
|
|
||||||
|
EntryClass = [WhisperForConditionalGeneration]
|
||||||
@@ -0,0 +1,191 @@
|
|||||||
|
import logging
|
||||||
|
from typing import Any, Dict, Optional
|
||||||
|
|
||||||
|
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||||
|
from sglang.srt.models.whisper import WhisperForConditionalGeneration
|
||||||
|
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||||
|
from sglang.srt.utils import load_audio
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# ISO 639-1 supported languages for Whisper
|
||||||
|
# From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
|
||||||
|
# Maps ISO 639-1 code -> Full language name
|
||||||
|
ISO639_1_SUPPORTED_LANGS = {
|
||||||
|
"af": "Afrikaans",
|
||||||
|
"ar": "Arabic",
|
||||||
|
"hy": "Armenian",
|
||||||
|
"az": "Azerbaijani",
|
||||||
|
"be": "Belarusian",
|
||||||
|
"bs": "Bosnian",
|
||||||
|
"bg": "Bulgarian",
|
||||||
|
"ca": "Catalan",
|
||||||
|
"zh": "Chinese",
|
||||||
|
"hr": "Croatian",
|
||||||
|
"cs": "Czech",
|
||||||
|
"da": "Danish",
|
||||||
|
"nl": "Dutch",
|
||||||
|
"en": "English",
|
||||||
|
"et": "Estonian",
|
||||||
|
"fi": "Finnish",
|
||||||
|
"fr": "French",
|
||||||
|
"gl": "Galician",
|
||||||
|
"de": "German",
|
||||||
|
"el": "Greek",
|
||||||
|
"he": "Hebrew",
|
||||||
|
"hi": "Hindi",
|
||||||
|
"hu": "Hungarian",
|
||||||
|
"is": "Icelandic",
|
||||||
|
"id": "Indonesian",
|
||||||
|
"it": "Italian",
|
||||||
|
"ja": "Japanese",
|
||||||
|
"kn": "Kannada",
|
||||||
|
"kk": "Kazakh",
|
||||||
|
"ko": "Korean",
|
||||||
|
"lv": "Latvian",
|
||||||
|
"lt": "Lithuanian",
|
||||||
|
"mk": "Macedonian",
|
||||||
|
"ms": "Malay",
|
||||||
|
"mr": "Marathi",
|
||||||
|
"mi": "Maori",
|
||||||
|
"ne": "Nepali",
|
||||||
|
"no": "Norwegian",
|
||||||
|
"fa": "Persian",
|
||||||
|
"pl": "Polish",
|
||||||
|
"pt": "Portuguese",
|
||||||
|
"ro": "Romanian",
|
||||||
|
"ru": "Russian",
|
||||||
|
"sr": "Serbian",
|
||||||
|
"sk": "Slovak",
|
||||||
|
"sl": "Slovenian",
|
||||||
|
"es": "Spanish",
|
||||||
|
"sw": "Swahili",
|
||||||
|
"sv": "Swedish",
|
||||||
|
"tl": "Tagalog",
|
||||||
|
"ta": "Tamil",
|
||||||
|
"th": "Thai",
|
||||||
|
"tr": "Turkish",
|
||||||
|
"uk": "Ukrainian",
|
||||||
|
"ur": "Urdu",
|
||||||
|
"vi": "Vietnamese",
|
||||||
|
"cy": "Welsh",
|
||||||
|
}
|
||||||
|
|
||||||
|
# Reverse mapping: Full language name (lowercase) -> ISO 639-1 code
|
||||||
|
LANG_NAME_TO_CODE = {
|
||||||
|
name.lower(): code for code, name in ISO639_1_SUPPORTED_LANGS.items()
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_language_to_code(language: Optional[str]) -> Optional[str]:
|
||||||
|
"""Convert a language input (full name or code) to ISO 639-1 code.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
language: Language as full name (e.g., 'English', 'Spanish') or
|
||||||
|
ISO 639-1 code (e.g., 'en', 'es')
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ISO 639-1 code or None if input is None
|
||||||
|
"""
|
||||||
|
if language is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
language_lower = language.lower().strip()
|
||||||
|
|
||||||
|
# Check if it's already a valid ISO code
|
||||||
|
if language_lower in ISO639_1_SUPPORTED_LANGS:
|
||||||
|
return language_lower
|
||||||
|
|
||||||
|
# Check if it's a full language name
|
||||||
|
if language_lower in LANG_NAME_TO_CODE:
|
||||||
|
return LANG_NAME_TO_CODE[language_lower]
|
||||||
|
|
||||||
|
# Not recognized
|
||||||
|
raise ValueError(
|
||||||
|
f"Language '{language}' not recognized. "
|
||||||
|
f"Use full name (e.g., 'English') or ISO 639-1 code (e.g., 'en')."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class WhisperProcessor(BaseMultimodalProcessor):
|
||||||
|
models = [WhisperForConditionalGeneration]
|
||||||
|
|
||||||
|
def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
|
||||||
|
super().__init__(hf_config, server_args, _processor, *args, **kwargs)
|
||||||
|
# Cache tokenizer for language token lookup
|
||||||
|
self._tokenizer = getattr(self._processor, "tokenizer", None)
|
||||||
|
|
||||||
|
def _extract_language_from_request(self, request_obj) -> Optional[str]:
|
||||||
|
sampling_params = getattr(request_obj, "sampling_params", None) or {}
|
||||||
|
language = sampling_params.pop("language", None)
|
||||||
|
return normalize_language_to_code(language)
|
||||||
|
|
||||||
|
def _get_language_token_id(self, language: Optional[str]) -> int:
|
||||||
|
# Default to English if not specified
|
||||||
|
if language is None:
|
||||||
|
language = "en" # Default to English
|
||||||
|
language_token = f"<|{language}|>"
|
||||||
|
return self._tokenizer.convert_tokens_to_ids(language_token)
|
||||||
|
|
||||||
|
async def process_mm_data_async(
|
||||||
|
self,
|
||||||
|
image_data,
|
||||||
|
audio_data,
|
||||||
|
input_text,
|
||||||
|
request_obj,
|
||||||
|
**kwargs,
|
||||||
|
) -> Optional[Dict[str, Any]]:
|
||||||
|
if not audio_data:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if len(audio_data) != 1:
|
||||||
|
raise ValueError(
|
||||||
|
f"Whisper expects exactly 1 audio input, got {len(audio_data)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
audios = [load_audio(audio) for audio in audio_data]
|
||||||
|
|
||||||
|
# For Whisper, ALWAYS use the proper transcription token sequence
|
||||||
|
# and IGNORE any text prompt - Whisper is a pure speech-to-text model
|
||||||
|
# The decoder_start_token_id and forced_decoder_ids from generation config
|
||||||
|
# set up: <|startoftranscript|> <|lang|> <|task|> [<|notimestamps|>]
|
||||||
|
|
||||||
|
# Extract language from request and get token ID
|
||||||
|
language = self._extract_language_from_request(request_obj)
|
||||||
|
language_token_id = self._get_language_token_id(language)
|
||||||
|
|
||||||
|
# Build decoder input tokens
|
||||||
|
# <|startoftranscript|> + <|lang|> + <|transcribe|> + <|notimestamps|>
|
||||||
|
decoder_start_token_id = getattr(
|
||||||
|
self.hf_config, "decoder_start_token_id", 50258
|
||||||
|
)
|
||||||
|
transcribe_token_id = self._tokenizer.convert_tokens_to_ids("<|transcribe|>")
|
||||||
|
notimestamps_token_id = self._tokenizer.convert_tokens_to_ids(
|
||||||
|
"<|notimestamps|>"
|
||||||
|
)
|
||||||
|
|
||||||
|
input_ids = [
|
||||||
|
decoder_start_token_id,
|
||||||
|
language_token_id,
|
||||||
|
transcribe_token_id,
|
||||||
|
notimestamps_token_id,
|
||||||
|
]
|
||||||
|
|
||||||
|
# Whisper expects input features padded to max_length (3000 frames = 30 seconds)
|
||||||
|
# This is the standard context length for Whisper
|
||||||
|
input_features = self._processor.feature_extractor(
|
||||||
|
audios[0],
|
||||||
|
sampling_rate=16000,
|
||||||
|
padding="max_length", # Pad to 3000 frames
|
||||||
|
return_tensors="pt",
|
||||||
|
)["input_features"][0]
|
||||||
|
|
||||||
|
return {
|
||||||
|
"input_ids": input_ids,
|
||||||
|
"mm_items": [
|
||||||
|
MultimodalDataItem(
|
||||||
|
feature=input_features,
|
||||||
|
modality=Modality.AUDIO,
|
||||||
|
)
|
||||||
|
],
|
||||||
|
}
|
||||||
@@ -1027,6 +1027,23 @@ register_conv_template(
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Whisper speech-to-text template
|
||||||
|
# Whisper uses special tokens: <|startoftranscript|>, <|en|>, <|transcribe|>, etc.
|
||||||
|
# Audio features are processed by encoder separately, not inserted into text
|
||||||
|
# The decoder start tokens (task, language) should be set via generation config
|
||||||
|
register_conv_template(
|
||||||
|
Conversation(
|
||||||
|
name="whisper",
|
||||||
|
system_template="",
|
||||||
|
system_message="",
|
||||||
|
roles=("", ""),
|
||||||
|
sep_style=SeparatorStyle.NO_COLON_SINGLE,
|
||||||
|
sep="",
|
||||||
|
stop_str=["<|endoftext|>"],
|
||||||
|
audio_token="", # Empty - audio is handled by encoder, not as text token
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
MODEL_TYPE_TO_TEMPLATE = {
|
MODEL_TYPE_TO_TEMPLATE = {
|
||||||
"internvl_chat": "internvl-2-5",
|
"internvl_chat": "internvl-2-5",
|
||||||
"deepseek_vl_v2": "deepseek-vl2",
|
"deepseek_vl_v2": "deepseek-vl2",
|
||||||
@@ -1036,6 +1053,7 @@ MODEL_TYPE_TO_TEMPLATE = {
|
|||||||
"minicpmo": "minicpmo",
|
"minicpmo": "minicpmo",
|
||||||
"deepseek-ocr": "deepseek-ocr",
|
"deepseek-ocr": "deepseek-ocr",
|
||||||
"paddleocr_vl": "paddle-ocr",
|
"paddleocr_vl": "paddle-ocr",
|
||||||
|
"whisper": "whisper",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -1129,3 +1147,11 @@ def match_paddle_ocr(model_path: str):
|
|||||||
return "paddle-ocr"
|
return "paddle-ocr"
|
||||||
model_type = get_model_type(model_path)
|
model_type = get_model_type(model_path)
|
||||||
return MODEL_TYPE_TO_TEMPLATE.get(model_type)
|
return MODEL_TYPE_TO_TEMPLATE.get(model_type)
|
||||||
|
|
||||||
|
|
||||||
|
@register_conv_template_matching_function
|
||||||
|
def match_whisper(model_path: str):
|
||||||
|
if "whisper" in model_path.lower():
|
||||||
|
return "whisper"
|
||||||
|
model_type = get_model_type(model_path)
|
||||||
|
return MODEL_TYPE_TO_TEMPLATE.get(model_type)
|
||||||
|
|||||||
@@ -1840,6 +1840,13 @@ class ServerArgs:
|
|||||||
self.speculative_algorithm is None
|
self.speculative_algorithm is None
|
||||||
), "Speculative decoding is currently not supported with Flex Attention backend"
|
), "Speculative decoding is currently not supported with Flex Attention backend"
|
||||||
|
|
||||||
|
# Encoder-decoder models (e.g., Whisper)
|
||||||
|
if model_config.is_encoder_decoder:
|
||||||
|
logger.warning(
|
||||||
|
"Cuda graph is disabled for encoder-decoder models (e.g., Whisper)"
|
||||||
|
)
|
||||||
|
self.disable_cuda_graph = True
|
||||||
|
|
||||||
# Major NVIDIA platforms backends
|
# Major NVIDIA platforms backends
|
||||||
if (
|
if (
|
||||||
self.attention_backend == "flashmla"
|
self.attention_backend == "flashmla"
|
||||||
|
|||||||
Reference in New Issue
Block a user