Fix Whisper transcription for audio over 30 seconds (#33604)

This commit is contained in:
Shenxiu Liu
2026-08-15 23:51:05 +08:00
committed by GitHub
parent 4beb157e87
commit fb97be4359
9 changed files with 1279 additions and 46 deletions
@@ -0,0 +1,122 @@
"""Unit tests for energy-aware audio chunking (long-audio transcription).
Whisper's encoder window is 30 s; longer audio must be split before
prompting. The splitter must cut at low-energy points (pauses) inside the
tail search window of each stride — never blindly at the stride boundary,
which could land mid-word — and the chunks must be contiguous,
non-overlapping, and reproduce the full waveform when concatenated.
"""
from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel() # must precede any import that pulls in sgl_kernel
import io
import unittest
import numpy as np
import soundfile as sf
from sglang.srt.entrypoints.openai.audio_chunking import (
find_split_point,
split_audio_energy_aware,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=4, suite="base-a-test-cpu")
SR = 16000
def _tone_with_silences(duration_s: float, silences: list, sr: int = SR) -> np.ndarray:
"""A 440 Hz tone with zeroed-out gaps at the given (start_s, end_s) spans."""
t = np.arange(int(duration_s * sr)) / sr
wav = (0.5 * np.sin(2 * np.pi * 440.0 * t)).astype(np.float32)
for start_s, end_s in silences:
wav[int(start_s * sr) : int(end_s * sr)] = 0.0
return wav
def _wav_bytes(wav: np.ndarray, sr: int = SR) -> bytes:
buf = io.BytesIO()
sf.write(buf, wav, sr, format="WAV")
return buf.getvalue()
def _decode(chunk: bytes) -> np.ndarray:
data, sr = sf.read(io.BytesIO(chunk), dtype="float32")
assert sr == SR
return data
class TestFindSplitPoint(CustomTestCase):
def test_picks_quietest_window(self):
# Loud tone with one silent 200 ms gap; the split point must land
# inside the gap.
wav = _tone_with_silences(4.0, [(2.0, 2.2)])
idx = find_split_point(wav, int(1.0 * SR), int(3.0 * SR))
self.assertGreaterEqual(idx, int(2.0 * SR))
self.assertLess(idx, int(2.2 * SR))
def test_uniform_energy_returns_search_start(self):
wav = _tone_with_silences(4.0, [])
idx = find_split_point(wav, int(1.0 * SR), int(3.0 * SR))
# All windows are equally loud; the first one wins.
self.assertEqual(idx, int(1.0 * SR))
class TestSplitAudioEnergyAware(CustomTestCase):
def test_short_audio_single_chunk(self):
wav = _tone_with_silences(10.0, [])
chunks, offsets = split_audio_energy_aware(_wav_bytes(wav), max_clip_s=30.0)
self.assertEqual(len(chunks), 1)
self.assertEqual(offsets, [0.0])
self.assertEqual(len(_decode(chunks[0])), len(wav))
def test_long_audio_splits_at_silence(self):
# 70 s tone with silence gaps planted inside each stride's search
# window (the last 1 s before the 30 s boundary): the first cut must
# land in [29.2, 29.5], which puts the second search window at
# ~[58.2, 59.2] where the second gap [58.5, 58.8] lives.
gaps = [(29.2, 29.5), (58.5, 58.8)]
wav = _tone_with_silences(70.0, gaps)
chunks, offsets = split_audio_energy_aware(_wav_bytes(wav), max_clip_s=30.0)
self.assertEqual(len(chunks), 3)
self.assertEqual(offsets[0], 0.0)
# Each cut lands inside its silence gap, not at the blind 30 s mark.
for offset, (gap_start, gap_end) in zip(offsets[1:], gaps):
self.assertGreaterEqual(offset, gap_start)
self.assertLess(offset, gap_end)
decoded = [_decode(c) for c in chunks]
# No chunk exceeds the model's window.
for d in decoded:
self.assertLessEqual(len(d), 30 * SR)
# Chunks are contiguous and non-overlapping: offsets line up with
# cumulative chunk lengths and the total sample count is preserved.
cumulative = 0
for d, offset in zip(decoded, offsets):
self.assertEqual(cumulative, int(round(offset * SR)))
cumulative += len(d)
self.assertEqual(cumulative, len(wav))
# Concatenation reproduces the waveform (modulo PCM16 quantization).
stitched = np.concatenate(decoded)
np.testing.assert_allclose(stitched, wav, atol=2.0 / 32768)
def test_no_silence_still_makes_progress(self):
# A constant-energy tone has no preferred split point; the splitter
# must still terminate with bounded chunks covering all samples.
wav = _tone_with_silences(65.0, [])
chunks, offsets = split_audio_energy_aware(_wav_bytes(wav), max_clip_s=30.0)
self.assertGreaterEqual(len(chunks), 3)
decoded = [_decode(c) for c in chunks]
for d in decoded:
self.assertLessEqual(len(d), 30 * SR)
self.assertEqual(sum(len(d) for d in decoded), len(wav))
self.assertEqual(len(offsets), len(chunks))
if __name__ == "__main__":
unittest.main()
@@ -14,12 +14,20 @@ from sglang.test.test_utils import maybe_stub_sgl_kernel
maybe_stub_sgl_kernel() # must precede any import that pulls in sgl_kernel
import asyncio
import io
import json
import unittest
from typing import List
from unittest.mock import Mock
from unittest.mock import AsyncMock, Mock, patch
from sglang.srt.entrypoints.openai.protocol import TranscriptionRequest
import numpy as np
import soundfile as sf
from sglang.srt.entrypoints.openai.protocol import (
TranscriptionRequest,
TranscriptionResponse,
)
from sglang.srt.entrypoints.openai.serving_transcription import (
OpenAIServingTranscription,
)
@@ -241,6 +249,426 @@ class TestStreamingFusedAutodetect(CustomTestCase):
self.assertFalse(any("<|" in d for d in deltas))
class _MockChunkTokenizerManager:
"""Mock TM scripting one result list per dispatched request, in order."""
def __init__(self, results_per_request: List):
self.model_config = Mock()
self.model_config.hf_config = Mock()
self.model_config.hf_config.architectures = ["WhisperForConditionalGeneration"]
self.server_args = Mock(
incremental_streaming_output=False,
asr_max_concurrent_sessions=32,
)
self.request_logger = Mock(log_requests=False)
self.tokenizer = Mock()
self.requests: List[GenerateReqInput] = []
self.aborted: List[str] = []
self._results = results_per_request
self.active_dispatches = 0
self.max_active_dispatches = 0
def generate_request(self, adapted_request, raw_request):
idx = len(self.requests)
# Mimic the real generate_request assigning a rid (via
# normalize_batch_and_arguments) so abort-by-rid is exercisable.
adapted_request.rid = f"rid{idx}"
self.requests.append(adapted_request)
results = self._results[idx]
async def gen():
self.active_dispatches += 1
self.max_active_dispatches = max(
self.max_active_dispatches, self.active_dispatches
)
# Give concurrently scheduled generators a chance to overlap at
# dispatch, then record that this request reached the engine.
await asyncio.sleep(0)
self.active_dispatches -= 1
for r in results:
# A bare exception entry simulates a chunk request failing.
if isinstance(r, BaseException):
raise r
yield r
return gen()
def abort_request(self, rid: str = "", abort_all: bool = False):
self.aborted.append(rid)
def create_abort_task(self, adapted_request):
return None
def _long_wav_bytes(duration_s: float = 65.0) -> bytes:
"""A 16 kHz tone with silence gaps inside each 30 s stride's split-search
window, so the energy-aware splitter cuts at ~29.2 s and ~58.5 s."""
sr = 16000
t = np.arange(int(duration_s * sr)) / sr
wav = (0.5 * np.sin(2 * np.pi * 440.0 * t)).astype(np.float32)
for gap_start, gap_end in ((29.2, 29.5), (58.5, 58.8)):
wav[int(gap_start * sr) : int(gap_end * sr)] = 0.0
buf = io.BytesIO()
sf.write(buf, wav, sr, format="WAV")
return buf.getvalue()
class TestLongAudioChunkedNonStreaming(CustomTestCase):
"""Audio longer than Whisper's 30 s window must be split into chunk
requests and the transcripts stitched in order — without chunking the
feature extractor silently truncates everything past 30 s."""
def _create_transcription(self, tm, audio_bytes, language="en", **kwargs):
serving = OpenAIServingTranscription(tm)
loop = get_or_create_event_loop()
return loop.run_until_complete(
serving.create_transcription(
audio_data=audio_bytes,
model="whisper",
language=language,
response_format=kwargs.pop("response_format", "json"),
temperature=0.0,
stream=False,
raw_request=Mock(),
**kwargs,
)
)
def test_long_audio_is_chunked_and_stitched(self):
texts = [" part one.", " part two.", " part three."]
tm = _MockChunkTokenizerManager(
[
[{"text": t, "meta_info": {"finish_reason": {"type": "stop"}}}]
for t in texts
]
)
result = self._create_transcription(tm, _long_wav_bytes(65.0))
self.assertIsInstance(result, TranscriptionResponse, result)
# In-order plain concatenation (vLLM-parity stitching).
self.assertEqual(result.text, " part one. part two. part three.")
self.assertEqual(result.usage.seconds, 65)
self.assertEqual(len(tm.requests), 3)
self.assertEqual(tm.max_active_dispatches, 1)
# Each chunk request is independent: own audio payload, own
# sampling_params dict (the multimodal processor pops keys out of
# it per request), stream=False, audio modality.
params_ids = {id(req.sampling_params) for req in tm.requests}
self.assertEqual(len(params_ids), len(tm.requests))
total_samples = 0
for req in tm.requests:
self.assertFalse(req.stream)
self.assertEqual(req.modalities, ["audio"])
data, sr = sf.read(io.BytesIO(req.audio_data), dtype="float32")
self.assertEqual(sr, 16000)
self.assertLessEqual(len(data), 30 * 16000)
total_samples += len(data)
self.assertEqual(total_samples, 65 * 16000)
def test_chunk_failure_returns_error_and_stops_dispatch(self):
# When one chunk request fails, create_transcription must return an
# error response (not a partial transcript), abort the in-flight
# request, and leave later chunks undispatched.
tm = _MockChunkTokenizerManager(
[
[
{
"text": " part one.",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
[ValueError("chunk boom")],
[
{
"text": " part three.",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
]
)
result = self._create_transcription(tm, _long_wav_bytes(65.0))
# Error response, not a TranscriptionResponse.
self.assertNotIsInstance(result, TranscriptionResponse)
# Sequential dispatch means the third chunk never reaches the engine.
self.assertEqual(len(tm.requests), 2)
self.assertIn(tm.requests[-1].rid, tm.aborted)
self.assertEqual(tm.max_active_dispatches, 1)
def test_split_failure_returns_error_without_dispatch(self):
tm = _MockChunkTokenizerManager([])
with patch(
"sglang.srt.entrypoints.openai.serving_transcription."
"split_audio_energy_aware",
side_effect=RuntimeError("decode failed"),
):
result = self._create_transcription(tm, _long_wav_bytes(65.0))
self.assertEqual(result.status_code, 400)
self.assertIn("Failed to split audio", json.loads(result.body)["message"])
self.assertEqual(tm.requests, [])
def test_short_audio_stays_unchunked(self):
tm = _MockChunkTokenizerManager(
[[{"text": " short.", "meta_info": {"finish_reason": {"type": "stop"}}}]]
)
result = self._create_transcription(tm, _long_wav_bytes(10.0))
self.assertIsInstance(result, TranscriptionResponse, result)
self.assertEqual(result.text, " short.")
self.assertEqual(len(tm.requests), 1)
def test_chunked_fused_autodetect_first_chunk_language_wins(self):
# language=None → fused auto-detect per chunk. Each chunk carries
# its own forced prefix; the stitched text must strip all of them,
# and the reported language comes from the first chunk.
tm = _MockChunkTokenizerManager(
[
[
{
"text": "<|en|><|transcribe|><|notimestamps|> part one.",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
[
{
"text": "<|fr|><|transcribe|><|notimestamps|> part two.",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
[
{
"text": "<|en|><|transcribe|><|notimestamps|> part three.",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
]
)
result = self._create_transcription(tm, _long_wav_bytes(65.0), language=None)
self.assertIsInstance(result, TranscriptionResponse, result)
self.assertEqual(result.text, "part one. part two. part three.")
self.assertNotIn("<|", result.text)
self.assertEqual(len(tm.requests), 3)
# Every chunk request kept the fused regex constraint.
for req in tm.requests:
self.assertIn("regex", req.sampling_params)
def test_chunked_fused_spaceless_script_not_space_joined(self):
# zh/ja/th chunk texts carry no boundary whitespace; stitching must
# not inject an ASCII space the model never emitted.
tm = _MockChunkTokenizerManager(
[
[
{
"text": "<|zh|><|transcribe|><|notimestamps|>你好",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
[
{
"text": "<|zh|><|transcribe|><|notimestamps|>世界",
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
]
)
result = self._create_transcription(tm, _long_wav_bytes(40.0), language=None)
self.assertIsInstance(result, TranscriptionResponse, result)
self.assertEqual(result.text, "你好世界")
self.assertEqual(len(tm.requests), 2)
def test_chunked_fused_language_uses_first_nonempty_chunk(self):
tm = _MockChunkTokenizerManager(
[
[
{
"text": "<|fr|><|transcribe|><|notimestamps|>",
"output_ids": [],
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
[
{
"text": "<|en|><|transcribe|><|notimestamps|> Hello",
"output_ids": [],
"meta_info": {"finish_reason": {"type": "stop"}},
}
],
]
)
result = self._create_transcription(
tm,
_long_wav_bytes(40.0),
language=None,
response_format="verbose_json",
)
self.assertEqual(result.text, "Hello")
self.assertEqual(result.language, "en")
class TestLongAudioChunkedStreaming(CustomTestCase):
"""_generate_long_audio_stream: chunks transcribed sequentially, deltas
emitted in audio order, exactly one finish frame."""
def _run_stream(self, results_per_request, fused=False, n_chunks=2):
tm = _MockChunkTokenizerManager(results_per_request)
serving = OpenAIServingTranscription(tm)
request = TranscriptionRequest(model="whisper", stream=True)
if fused:
request._fused_autodetect = True
request._fused_ts_variant = False
request._audio_chunks = [b"chunk%d" % i for i in range(n_chunks)]
# Streaming has no segment timing, so the offsets are intentionally
# not set here — only the non-streaming verbose_json path reads them.
adapted = GenerateReqInput(
text="", modalities=["audio"], sampling_params={"temperature": 0.0}
)
raw_request = Mock()
raw_request.is_disconnected = AsyncMock(return_value=False)
async def drive():
frames = []
async for frame in serving._generate_long_audio_stream(
adapted, request, raw_request
):
frames.append(frame)
return frames
loop = get_or_create_event_loop()
return tm, request, loop.run_until_complete(drive())
@staticmethod
def _finish_reasons(frames: List[str]) -> List[str]:
out = []
for line in frames:
if not line.startswith("data: ") or line.strip() == "data: [DONE]":
continue
obj = json.loads(line[len("data: ") :])
for choice in obj.get("choices", []):
if choice.get("finish_reason"):
out.append(choice["finish_reason"])
return out
def test_chunks_streamed_in_order_with_single_finish(self):
tm, _, frames = self._run_stream(
[
[_chunk(" Hello"), _chunk(" Hello world", finish="stop")],
[_chunk(" Again"), _chunk(" Again done", finish="stop")],
]
)
self.assertEqual(
_deltas_from_sse(frames), [" Hello", " world", " Again", " done"]
)
self.assertEqual(self._finish_reasons(frames), ["stop"])
self.assertEqual(frames[-1], "data: [DONE]\n\n")
# Both chunk requests were dispatched and streamed.
self.assertEqual(len(tm.requests), 2)
self.assertTrue(all(req.stream for req in tm.requests))
def test_disconnect_between_chunks_stops_and_aborts(self):
# Client disconnects after the first chunk: the second chunk is
# never dispatched, and the in-flight request from chunk 0 is
# already done so nothing is left decoding.
tm = _MockChunkTokenizerManager(
[
[_chunk(" Hello", finish="stop")],
[_chunk(" world", finish="stop")],
]
)
serving = OpenAIServingTranscription(tm)
request = TranscriptionRequest(model="whisper", stream=True)
request._audio_chunks = [b"chunk0", b"chunk1"]
adapted = GenerateReqInput(
text="", modalities=["audio"], sampling_params={"temperature": 0.0}
)
raw_request = Mock()
# Connected for the first chunk, disconnected before the second.
raw_request.is_disconnected = AsyncMock(side_effect=[False, True])
async def drive():
return [
f
async for f in serving._generate_long_audio_stream(
adapted, request, raw_request
)
]
frames = get_or_create_event_loop().run_until_complete(drive())
self.assertEqual(_deltas_from_sse(frames), [" Hello"])
# Only the first chunk was dispatched.
self.assertEqual(len(tm.requests), 1)
def test_abnormal_chunk_finish_reason_not_masked(self):
# A non-final chunk truncated at the token cap (finish="length")
# must surface in the single final frame even though later chunks
# stop cleanly — otherwise silently missing transcript content
# reads as success.
_, _, frames = self._run_stream(
[
[_chunk(" Hello", finish="length")],
[_chunk(" world", finish="stop")],
]
)
self.assertEqual(_deltas_from_sse(frames), [" Hello", " world"])
self.assertEqual(self._finish_reasons(frames), ["length"])
def test_fused_chunks_strip_prefixes_and_preserve_boundary_space(self):
tm, request, frames = self._run_stream(
[
[
_chunk("<|fr|><|transcribe|>"),
_chunk(
"<|fr|><|transcribe|><|notimestamps|> Bonjour", finish="stop"
),
],
[
_chunk(
"<|fr|><|transcribe|><|notimestamps|> le monde",
finish="stop",
)
],
],
fused=True,
)
deltas = _deltas_from_sse(frames)
self.assertFalse(any("<|" in d for d in deltas))
# No leading space at stream start; the later chunk's own leading
# space is the seam separator.
self.assertFalse(deltas[0].startswith(" "))
self.assertEqual("".join(deltas), "Bonjour le monde")
self.assertEqual(request.language, "fr")
self.assertEqual(self._finish_reasons(frames), ["stop"])
def test_fused_spaceless_script_chunks_not_space_joined(self):
# zh/ja/th transcripts carry no boundary whitespace; the seam must
# not inject an ASCII space the model never emitted.
_, request, frames = self._run_stream(
[
[_chunk("<|zh|><|transcribe|><|notimestamps|>你好", finish="stop")],
[_chunk("<|zh|><|transcribe|><|notimestamps|>世界", finish="stop")],
],
fused=True,
)
self.assertEqual("".join(_deltas_from_sse(frames)), "你好世界")
self.assertEqual(request.language, "zh")
def test_fused_leading_silence_uses_first_nonempty_chunk_language(self):
_, request, frames = self._run_stream(
[
[_chunk("<|fr|><|transcribe|><|notimestamps|>", finish="stop")],
[
_chunk(
"<|en|><|transcribe|><|notimestamps|> Hello",
finish="stop",
)
],
],
fused=True,
)
self.assertEqual("".join(_deltas_from_sse(frames)), "Hello")
self.assertEqual(request.language, "en")
class TestStreamingIncrementalOutputMode(CustomTestCase):
"""Server runs with ``incremental_streaming_output=True``.
@@ -314,5 +314,74 @@ class TestWhisperBuildFusedAutodetectParams(CustomTestCase):
SamplingParams(**params)
class _FakeTokenizer:
"""Decodes token ids to a deterministic placeholder string."""
eos_token_id = 50257
def decode(self, ids, skip_special_tokens=True):
return " " + " ".join(f"tok{i}" for i in ids)
class TestWhisperChunkedVerboseResponse(CustomTestCase):
"""build_verbose_response_chunked for audio split into >30 s chunks.
Whisper timestamp tokens are relative to each chunk's own 30 s window,
so every chunk's segments must be shifted by the chunk's start offset
in the original audio and segment ids must keep counting across chunks.
"""
TS = WhisperAdapter.TIMESTAMP_BASE_TOKEN_ID # <|0.00|>
def test_segments_offset_by_chunk_start(self):
request = TranscriptionRequest(
model="whisper", language="en", audio_duration_s=57.06
)
from sglang.srt.entrypoints.openai.protocol import TranscriptionUsage
usage = TranscriptionUsage(seconds=58)
rets = [
# chunk 0: one closed segment [0.00 → 5.00]
{"output_ids": [100, 101, self.TS + 250]},
# chunk 1: closed segment [0.00 → 2.00] + trailing unclosed text
{"output_ids": [200, self.TS + 100, 300]},
]
resp = WhisperAdapter().build_verbose_response_chunked(
request, "你好世界", rets, [0.0, 29.3], _FakeTokenizer(), usage
)
self.assertEqual(resp.language, "en")
self.assertEqual(resp.duration, 57.06)
self.assertEqual(resp.text, "你好世界")
self.assertEqual([s.id for s in resp.segments], [0, 1, 2])
self.assertEqual(
[(s.start, s.end) for s in resp.segments],
[(0.0, 5.0), (29.3, 31.3), (31.3, 31.3)],
)
def test_single_chunk_at_zero_offset_matches_unchunked(self):
request = TranscriptionRequest(
model="whisper", language="en", audio_duration_s=10.0
)
from sglang.srt.entrypoints.openai.protocol import TranscriptionUsage
usage = TranscriptionUsage(seconds=10)
output_ids = [100, self.TS + 50]
text, segments = WhisperAdapter._parse_segments(output_ids, _FakeTokenizer())
chunked = WhisperAdapter().build_verbose_response_chunked(
request,
text,
[{"output_ids": output_ids}],
[0.0],
_FakeTokenizer(),
usage,
)
self.assertEqual(chunked.text, text)
self.assertEqual(
[(s.id, s.start, s.end) for s in chunked.segments],
[(s.id, s.start, s.end) for s in segments],
)
if __name__ == "__main__":
unittest.main()