LLM/WhisperLiveKit
1
0
Fork 0
mirror of https://github.com/QuentinFuxa/WhisperLiveKit.git synced 2026-03-27 06:16:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Compare commits

base: LLM:v0.2.20
Branches Tags
LLM:main LLM:benchmarks-h100 LLM:qwen3 LLM:voxtral_tests LLM:feature/voxtral LLM:api_live LLM:VAD-evolutions LLM:rework_state_light LLM:new_api LLM:language-detection LLM:ScriptProcessorNode-to-AudioWorklet LLM:translation-improvements LLM:windows_audio_loopback LLM:new-api LLM:regularfry-ayo-logging-fixes LLM:seamless-streaming
LLM:v0.2.20 LLM:v0.2.19 LLM:0.2.18 LLM:0.2.17.post1 LLM:0.2.17 LLM:0.2.16 LLM:0.2.15 LLM:0.2.14 LLM:0.2.13 LLM:0.2.12 LLM:0.2.11 LLM:0.2.10 LLM:0.2.9 LLM:0.2.8 LLM:0.2.7 LLM:0.2.6 LLM:0.2.5 LLM:0.2.4.dev0 LLM:0.2.2 LLM:0.2.1 LLM:0.1.9 LLM:0.1.8 LLM:0.1.7 LLM:0.1.6 LLM:0.1.5 LLM:0.1.4 LLM:0.1.3 LLM:0.1.1 LLM:0.0.1
..
compare: LLM:benchmarks-h100
Branches Tags
LLM:benchmarks-h100 LLM:main LLM:qwen3 LLM:voxtral_tests LLM:feature/voxtral LLM:api_live LLM:VAD-evolutions LLM:rework_state_light LLM:new_api LLM:language-detection LLM:ScriptProcessorNode-to-AudioWorklet LLM:translation-improvements LLM:windows_audio_loopback LLM:new-api LLM:regularfry-ayo-logging-fixes LLM:seamless-streaming
LLM:v0.2.20 LLM:v0.2.19 LLM:0.2.18 LLM:0.2.17.post1 LLM:0.2.17 LLM:0.2.16 LLM:0.2.15 LLM:0.2.14 LLM:0.2.13 LLM:0.2.12 LLM:0.2.11 LLM:0.2.10 LLM:0.2.9 LLM:0.2.8 LLM:0.2.7 LLM:0.2.6 LLM:0.2.5 LLM:0.2.4.dev0 LLM:0.2.2 LLM:0.2.1 LLM:0.1.9 LLM:0.1.8 LLM:0.1.7 LLM:0.1.6 LLM:0.1.5 LLM:0.1.4 LLM:0.1.3 LLM:0.1.1 LLM:0.0.1

20 Commits
v0.2.20 ... benchmarks

Author SHA1 Message Date
Quentin Fuxa
47d4cbeecc reorganize benchmarks: move H100 results to benchmarks/h100/ 2026-03-15 23:59:00 +01:00
Quentin Fuxa
f75dfb386d final benchmark: Voxtral vLLM realtime streaming 2026-03-15 23:59:00 +01:00
Quentin Fuxa
276ba84d02 update figures with Voxtral vLLM results 2026-03-15 23:55:00 +01:00
Quentin Fuxa
36b3885cf2 add Voxtral 4B to benchmark figures 2026-03-15 23:30:00 +01:00
Quentin Fuxa
a29e799ba5 update H100 benchmark figures with ACL6060 results 2026-03-15 22:30:00 +01:00
Quentin Fuxa
22325ba326 tune simul-kv: 2s inference interval, configurable min_new_seconds 2026-03-15 21:30:00 +01:00
Quentin Fuxa
a540a5fd10 fix simul-kv audio trim bug, add 1.7B v2 alignment heads 2026-03-15 20:45:00 +01:00
Quentin Fuxa
7b08ea74ab add H100 benchmark figures 2026-03-15 19:15:00 +01:00
Quentin Fuxa
b69eaf82be qwen3 simul+kv: optimized streaming with kv cache reuse 2026-03-15 18:30:00 +01:00
Quentin Fuxa
ed503be140 qwen 2026-01-02 23:52:00 +01:00
Quentin Fuxa
a6a85431f6 update benchmark with qwen3 which reuses kv cache 2026-03-15 22:32:01 +01:00
Quentin Fuxa
dd48997674 qwen3: reuse encoder kv cache 2026-03-15 22:31:39 +01:00
Quentin Fuxa
f24481dc29 update archi 2026-03-15 11:36:45 +01:00
Quentin Fuxa
ed76f40ee5 Merge branch 'main' of https://github.com/QuentinFuxa/WhisperLiveKit 2026-03-15 11:16:38 +01:00
Quentin Fuxa
5330b3fac5 update benchmark part 2026-03-15 11:16:26 +01:00
Quentin Fuxa
0c73a73aa3 update benchmark results and procedure 2026-03-15 11:16:15 +01:00
Quentin Fuxa
2d6bc4f572 Add '*.c' to .dockerignore 2026-03-14 00:18:10 +01:00
Quentin Fuxa
dfd5bf417c voxtral mlx : improved chunking 2026-03-14 00:13:29 +01:00
Quentin Fuxa
9d8db7ab38 add qwen3 simul in tests 2026-03-14 00:13:09 +01:00
Quentin Fuxa
fa15115163 qwen3 alignment heads 2026-03-14 00:12:50 +01:00
53 changed files with 17199 additions and 2386 deletions
Expand all files Collapse all files

1
.dockerignore
View File

@@ -11,3 +11,4 @@ __pycache__
.secrets
dist
build
*.c

205
BENCHMARK.md
View File

@@ -1,205 +0,0 @@
# WhisperLiveKit Benchmark Report
Benchmark comparing all supported ASR backends, streaming policies, and model sizes on Apple Silicon.
All tests run through the full AudioProcessor pipeline (same code path as production WebSocket).
## Test Environment
| Property | Value |
|----------|-------|
| Hardware | Apple M4, 32 GB RAM |
| OS | macOS 25.3.0 (arm64) |
| Python | 3.13 |
| faster-whisper | 1.2.1 |
| mlx-whisper | installed (via mlx) |
| Voxtral MLX | native MLX backend |
| Voxtral (HF) | transformers-based |
| VAC (Silero VAD) | enabled unless noted |
| Chunk size | 100 ms |
| Pacing | no-realtime (as fast as possible) |
## Audio Test Files
| File | Duration | Language | Speakers | Description |
|------|----------|----------|----------|-------------|
| `00_00_07_english_1_speaker.wav` | 7.2 s | English | 1 | Short dictation with pauses |
| `00_00_16_french_1_speaker.wav` | 16.3 s | French | 1 | French speech with intentional silence gaps |
| `00_00_30_english_3_speakers.wav` | 30.0 s | English | 3 | Multi-speaker conversation |
Ground truth transcripts (`.transcript.json`) with per-word timestamps are hand-verified.
---
## Results
### English -- Short (7.2 s, 1 speaker)
| Backend | Policy | Model | RTF | WER | Timestamp MAE |
|---------|--------|-------|-----|-----|---------------|
| faster-whisper | LocalAgreement | base | 0.20x | 21.1% | 0.080 s |
| faster-whisper | SimulStreaming | base | 0.14x | 0.0% | 0.239 s |
| faster-whisper | LocalAgreement | small | 0.59x | 21.1% | 0.089 s |
| faster-whisper | SimulStreaming | small | 0.39x | 0.0% | 0.221 s |
| mlx-whisper | LocalAgreement | base | 0.05x | 21.1% | 0.080 s |
| mlx-whisper | SimulStreaming | base | 0.14x | 10.5% | 0.245 s |
| mlx-whisper | LocalAgreement | small | 0.16x | 21.1% | 0.089 s |
| mlx-whisper | SimulStreaming | small | 0.20x | 10.5% | 0.226 s |
| voxtral-mlx | voxtral | 4B | 0.32x | 0.0% | 0.254 s |
| voxtral (HF) | voxtral | 4B | 1.29x | 0.0% | 1.876 s |
### English -- Multi-speaker (30.0 s, 3 speakers)
| Backend | Policy | Model | RTF | WER | Timestamp MAE |
|---------|--------|-------|-----|-----|---------------|
| faster-whisper | LocalAgreement | base | 0.24x | 44.7% | 0.235 s |
| faster-whisper | SimulStreaming | base | 0.10x | 5.3% | 0.398 s |
| faster-whisper | LocalAgreement | small | 0.59x | 25.0% | 0.226 s |
| faster-whisper | SimulStreaming | small | 0.26x | 5.3% | 0.387 s |
| mlx-whisper | LocalAgreement | base | 0.06x | 23.7% | 0.237 s |
| mlx-whisper | SimulStreaming | base | 0.11x | 5.3% | 0.395 s |
| mlx-whisper | LocalAgreement | small | 0.13x | 25.0% | 0.226 s |
| mlx-whisper | SimulStreaming | small | 0.20x | 5.3% | 0.394 s |
| voxtral-mlx | voxtral | 4B | 0.31x | 9.2% | 0.176 s |
| voxtral (HF) | voxtral | 4B | 1.00x | 32.9% | 1.034 s |
<p align="center">
<img src="benchmark_chart.png" alt="Benchmark comparison on 30s English" width="800">
</p>
<p align="center">
<img src="benchmark_scatter.png" alt="Speed vs Accuracy tradeoff" width="700">
</p>
### French (16.3 s, 1 speaker, `--language fr`)
| Backend | Policy | Model | RTF | WER | Timestamp MAE |
|---------|--------|-------|-----|-----|---------------|
| faster-whisper | LocalAgreement | base | 0.22x | 25.7% | 3.460 s |
| faster-whisper | SimulStreaming | base | 0.10x | 31.4% | 3.660 s |
| faster-whisper | LocalAgreement | small | 0.76x | 42.9% | 0.051 s |
| faster-whisper | SimulStreaming | small | 0.29x | 25.7% | 0.219 s |
| mlx-whisper | LocalAgreement | base | 0.09x | ~45%* | ~5.0 s* |
| mlx-whisper | SimulStreaming | base | 0.09x | 40.0% | 3.540 s |
| mlx-whisper | LocalAgreement | small | 0.14x | 25.7% | 0.083 s |
| mlx-whisper | SimulStreaming | small | 0.17x | 31.4% | 0.203 s |
| voxtral-mlx | voxtral | 4B | 0.18x | 37.1% | 3.422 s |
| voxtral (HF) | voxtral | 4B | 0.63x | 28.6% | 4.040 s |
\* mlx-whisper + LocalAgreement + base is unstable on this French file (WER fluctuates 34-1037% across runs due to hallucination loops). The `small` model does not have this problem.
**Timestamp note:** The base model produces very high timestamp MAE (3.4-3.7s) on this French file because it misaligns words around the silence gaps. The small model handles this much better (0.05-0.22s MAE). Voxtral also drifts on the silence gaps.
---
## Model Size Comparison (base vs small)
| | base | small | Observation |
|--|------|-------|-------------|
| **RTF** | 0.05-0.24x | 0.13-0.76x | small is 2-3x slower |
| **English WER (SS)** | 0-5.3% | 0-5.3% | No improvement: SimulStreaming already saturates on base |
| **English WER (LA)** | 21-44.7% | 21-25% | small reduces LA errors on longer audio |
| **French WER** | 25-40% | 25-43% | Mixed: depends on backend/policy combo |
| **French timestamps** | 3.4-5.0s MAE | 0.05-0.22s MAE | small is dramatically better for French timestamps |
In short: **base + SimulStreaming** gives the best speed/accuracy tradeoff for English. The small model only helps if you need LocalAgreement (for subtitle-grade timestamps) or non-English languages.
---
## Key Findings
### Speed (RTF = processing time / audio duration, lower is better)
1. **mlx-whisper + LocalAgreement + base** is the fastest combo on Apple Silicon: 0.05-0.06x RTF on English. 30 seconds of audio in under 2 seconds.
2. For **faster-whisper**, SimulStreaming is faster than LocalAgreement. For **mlx-whisper**, it is the opposite: LocalAgreement (0.05-0.06x) outperforms SimulStreaming (0.11-0.14x) on speed.
3. **voxtral-mlx** runs at 0.18-0.32x RTF -- 3-5x slower than mlx-whisper base, but well within real-time.
4. **voxtral (HF transformers)** hits 1.0-1.3x RTF. At the real-time boundary on Apple Silicon. Use the MLX variant instead.
5. The **small** model is 2-3x slower than base across all backends.
### Accuracy (WER = Word Error Rate, lower is better)
1. **SimulStreaming** gives dramatically lower WER than LocalAgreement on the whisper backends. On the 30s English file: 5.3% vs 23-44%.
2. **voxtral-mlx** hits 0% on short English and 9.2% on multi-speaker. It auto-detects language natively. Whisper also supports `--language auto`, but tends to bias towards English on short segments.
3. **LocalAgreement** tends to repeat the last sentence at end-of-stream (a known LCP artifact), inflating WER. This is visible in the 21% WER on the 7s file -- the same 4 extra words appear in every LA run.
4. On **French** with the correct `--language fr`, whisper base achieves 25-40% WER -- comparable to Voxtral's 28-37%. The small model does not consistently improve French WER.
### Timestamps (MAE = Mean Absolute Error on word start times)
1. **LocalAgreement** gives the best timestamps on English (0.08-0.09s MAE).
2. **SimulStreaming** is less precise (0.22-0.40s MAE) but good enough for most applications.
3. On French with silence gaps, **base model timestamps are unreliable** (3.4-5s MAE). The **small model fixes this** (0.05-0.22s MAE). This is the strongest argument for using `small` over `base`.
4. **voxtral-mlx** has good timestamps on English (0.18-0.25s MAE) but drifts on audio with long silence gaps (3.4s MAE on the French file).
### VAC (Voice Activity Classification) Impact
| Backend | Policy | VAC | 7s English WER | 30s English WER |
|---------|--------|-----|----------------|-----------------|
| faster-whisper | LocalAgreement | on | 21.1% | 44.7% |
| faster-whisper | LocalAgreement | off | 100.0% | 100.0% |
| voxtral-mlx | voxtral | on | 0.0% | 9.2% |
| voxtral-mlx | voxtral | off | 0.0% | 9.2% |
- **Whisper backends need VAC** to work in streaming mode. Without it the buffer logic breaks down and you get empty or garbage output.
- **Voxtral is unaffected by VAC** since it handles its own internal chunking. Identical results with or without. VAC still saves compute on silent segments.
---
## Recommendations
| Use Case | Backend | Policy | Model | Notes |
|----------|---------|--------|-------|-------|
| Fastest English (Apple Silicon) | mlx-whisper | SimulStreaming | base | 0.11x RTF, 5.3% WER |
| Fastest English (Linux/GPU) | faster-whisper | SimulStreaming | base | 0.10x RTF, 5.3% WER |
| Best accuracy, English | faster-whisper | SimulStreaming | small | 0.26x RTF, 5.3% WER, still fast |
| Multilingual / auto-detect | voxtral-mlx | voxtral | 4B | 100+ languages, 0.18-0.32x RTF |
| Best timestamps | any | LocalAgreement | small | 0.05-0.09s MAE, good for subtitles |
| Low memory / embedded | mlx-whisper | SimulStreaming | base | Smallest footprint, fastest response |
---
## Caveats
- **3 test files, ~53 seconds total.** Results give relative rankings between backends but should not be taken as definitive WER numbers. Run on your own data for production decisions.
- **RTF varies between runs** (up to +/-30%) depending on thermal state, background processes, and model caching. The numbers above are single sequential runs on a warm machine.
- **Only base and small tested.** Medium and large-v3 would likely improve WER at the cost of higher RTF. We did not test them here because they are slow on Apple Silicon without GPU.
---
## Reproducing These Benchmarks
```bash
# Install test dependencies
pip install -e ".[test]"
# Single backend test
python test_backend_offline.py --backend faster-whisper --policy simulstreaming --model base --no-realtime
# With a specific language
python test_backend_offline.py --backend mlx-whisper --policy simulstreaming --model small --lan fr --no-realtime
# Multi-backend auto-detect benchmark
python test_backend_offline.py --benchmark --no-realtime
# Export to JSON
python test_backend_offline.py --benchmark --no-realtime --json results.json
# Test with your own audio
python test_backend_offline.py --backend voxtral-mlx --audio your_file.wav --no-realtime
```
The benchmark harness computes WER and timestamp accuracy automatically when ground truth
`.transcript.json` files exist alongside the audio files. See `audio_tests/` for the format.
---
## Help Us Benchmark on More Hardware
These results are from a single Apple M4 machine. We'd love to see numbers from other setups: Linux with CUDA GPUs, older Macs, different CPU architectures, cloud instances, etc.
If you run the benchmark on your hardware, please open an issue or PR with your results and we will add them here. The more data points we have, the better the recommendations get.
What we are especially interested in:
- **NVIDIA GPUs** (RTX 3090, 4090, A100, T4, etc.) with faster-whisper
- **Older Apple Silicon** (M1, M2, M3) with mlx-whisper and voxtral-mlx
- **Medium and large-v3 models** (we only tested base and small so far)
- **Longer audio files** or domain-specific audio (medical, legal, call center)
- **Other languages** beyond English and French

37
README.md
View File

@@ -95,14 +95,6 @@ See [docs/API.md](docs/API.md) for the complete API reference.
> - For HTTPS requirements, see the **Parameters** section for SSL configuration options.
#### Use it to capture audio from web pages.
Go to `chrome-extension` for instructions.
<p align="center">
<img src="https://raw.githubusercontent.com/QuentinFuxa/WhisperLiveKit/refs/heads/main/chrome-extension/demo-extension.png" alt="WhisperLiveKit Demo" width="600">
</p>
#### Optional Dependencies
@@ -134,13 +126,24 @@ uv sync --extra cu129 --extra voxtral-hf --extra translation
See **Parameters & Configuration** below on how to use them.
<p align="center">
<img src="benchmark_scatter.png" alt="Speed vs Accuracy tradeoff" width="700">
<img src="benchmark_scatter_en_aware.png" alt="Speed vs Accuracy — English" width="700">
</p>
<p align="center">
<img src="benchmark_scatter_fr_aware.png" alt="Speed vs Accuracy — French" width="700">
</p>
See **[BENCHMARK.md](BENCHMARK.md)** for the full benchmark with tables, model size comparison, and more.
Benchmarks use 6 minutes of public [LibriVox](https://librivox.org/) audiobook recordings per language (30s + 60s + 120s + 180s), with ground truth from [Project Gutenberg](https://www.gutenberg.org/). Fully reproducible with `python scripts/run_scatter_benchmark.py`.
We are actively looking for benchmark results on other hardware (NVIDIA GPUs, different Apple Silicon chips, cloud instances). If you run the benchmarks on your machine, please share your results via an issue or PR!
#### Use it to capture audio from web pages.
Go to `chrome-extension` for instructions.
<p align="center">
<img src="https://raw.githubusercontent.com/QuentinFuxa/WhisperLiveKit/refs/heads/main/chrome-extension/demo-extension.png" alt="WhisperLiveKit Demo" width="600">
</p>
### Voxtral Backend
@@ -259,7 +262,7 @@ async def websocket_endpoint(websocket: WebSocket):
| SimulStreaming backend options | Description | Default |
|-----------|-------------|---------|
| `--disable-fast-encoder` | Disable Faster Whisper or MLX Whisper backends for the encoder (if installed). Inference can be slower but helpful when GPU memory is limited | `False` |
| `--custom-alignment-heads` | Use your own alignment heads, useful when `--model-dir` is used. Use `scripts/determine_alignment_heads.py` to extract them. <img src="scripts/alignment_heads.png" alt="WhisperLiveKit Demo" width="300">
| `--custom-alignment-heads` | Use your own alignment heads, useful when `--model-dir` is used. Use `scripts/determine_alignment_heads.py` to extract them. <img src="scripts/alignment_heads_qwen3_asr_1.7B.png" alt="WhisperLiveKit Demo" width="300">
| `None` |
| `--frame-threshold` | AlignAtt frame threshold (lower = faster, higher = more accurate) | `25` |
| `--beams` | Number of beams for beam search (1 = greedy decoding) | `1` |
@@ -371,7 +374,7 @@ docker compose up --build wlk-cpu
# Quick benchmark with the CLI
wlk bench
wlk bench --backend faster-whisper --model large-v3
wlk bench --json results.json
wlk bench --languages all --json results.json
# Install test dependencies for full suite
pip install -e ".[test]"
@@ -379,13 +382,11 @@ pip install -e ".[test]"
# Run unit tests (no model download required)
pytest tests/ -v
# Detailed multi-backend benchmark
python test_backend_offline.py --benchmark --no-realtime
python test_backend_offline.py --benchmark --no-realtime --json results.json
# Speed vs Accuracy scatter plot (all backends, compute-aware + unaware)
python scripts/create_long_samples.py # generate ~90s test samples (cached)
python scripts/run_scatter_benchmark.py # English (both modes)
python scripts/run_scatter_benchmark.py --lang fr # French
```
See [BENCHMARK.md](BENCHMARK.md) for a full comparison of backends, policies, WER, speed, and
timestamp accuracy on Apple Silicon.
## Use Cases
Capture discussions in real-time for meeting transcription, help hearing-impaired users follow conversations through accessibility tools, transcribe podcasts or videos automatically for content creation, transcribe support calls with speaker identification for customer service...

BIN
architecture.png
View File

Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 446 KiB

After

Width:  |  Height:  |  Size: 426 KiB

97
audio_tests/00_00_07_english_1_speaker.transcript.json
View File

@@ -1,97 +0,0 @@
[
{
"word": "This",
"start": 0.0,
"end": 0.24
},
{
"word": "is",
"start": 0.24,
"end": 0.56
},
{
"word": "a",
"start": 0.56,
"end": 0.76
},
{
"word": "transcription",
"start": 0.76,
"end": 1.32
},
{
"word": "test.",
"start": 1.32,
"end": 2.0
},
{
"word": "We",
"start": 2.4,
"end": 2.5
},
{
"word": "want",
"start": 2.5,
"end": 2.66
},
{
"word": "to",
"start": 2.66,
"end": 2.84
},
{
"word": "see",
"start": 2.84,
"end": 3.1
},
{
"word": "if",
"start": 3.1,
"end": 3.34
},
{
"word": "we",
"start": 3.34,
"end": 3.5
},
{
"word": "can",
"start": 3.5,
"end": 3.68
},
{
"word": "use",
"start": 3.68,
"end": 4.04
},
{
"word": "smaller",
"start": 4.04,
"end": 4.76
},
{
"word": "chunks.",
"start": 4.76,
"end": 5.16
},
{
"word": "What",
"start": 6.06,
"end": 6.32
},
{
"word": "do",
"start": 6.32,
"end": 6.44
},
{
"word": "you",
"start": 6.44,
"end": 6.58
},
{
"word": "think?",
"start": 6.58,
"end": 6.84
}
]

177
audio_tests/00_00_16_french_1_speaker.transcript.json
View File

@@ -1,177 +0,0 @@
[
{
"word": "Ok,",
"start": 2.02,
"end": 2.38
},
{
"word": "là",
"start": 2.52,
"end": 2.58
},
{
"word": "c",
"start": 2.58,
"end": 2.74
},
{
"word": "'est",
"start": 2.74,
"end": 2.76
},
{
"word": "un",
"start": 2.76,
"end": 2.86
},
{
"word": "test,",
"start": 2.86,
"end": 3.2
},
{
"word": "on",
"start": 3.34,
"end": 3.34
},
{
"word": "veut",
"start": 3.34,
"end": 3.48
},
{
"word": "voir",
"start": 3.48,
"end": 3.86
},
{
"word": "si",
"start": 3.86,
"end": 4.14
},
{
"word": "ça",
"start": 4.14,
"end": 4.26
},
{
"word": "arrive",
"start": 4.26,
"end": 4.36
},
{
"word": "à",
"start": 4.36,
"end": 4.5
},
{
"word": "capté",
"start": 4.5,
"end": 4.78
},
{
"word": "le",
"start": 4.78,
"end": 4.9
},
{
"word": "silence.",
"start": 4.9,
"end": 5.44
},
{
"word": "Là",
"start": 9.24,
"end": 9.6
},
{
"word": "il",
"start": 9.6,
"end": 9.78
},
{
"word": "est",
"start": 9.78,
"end": 9.84
},
{
"word": "une",
"start": 9.84,
"end": 9.96
},
{
"word": "telle",
"start": 9.96,
"end": 10.12
},
{
"word": "seconde",
"start": 10.12,
"end": 10.38
},
{
"word": "de",
"start": 10.38,
"end": 10.48
},
{
"word": "silence",
"start": 10.48,
"end": 10.78
},
{
"word": "et",
"start": 10.78,
"end": 11.06
},
{
"word": "je",
"start": 11.06,
"end": 11.16
},
{
"word": "vous",
"start": 11.16,
"end": 11.32
},
{
"word": "parle.",
"start": 11.32,
"end": 11.68
},
{
"word": "Et",
"start": 13.28,
"end": 13.64
},
{
"word": "voilà,",
"start": 13.64,
"end": 13.96
},
{
"word": "allez",
"start": 14.36,
"end": 14.62
},
{
"word": "on",
"start": 14.62,
"end": 14.78
},
{
"word": "va",
"start": 14.78,
"end": 14.88
},
{
"word": "tester",
"start": 14.88,
"end": 15.06
},
{
"word": "ça.",
"start": 15.06,
"end": 15.36
}
]

382
audio_tests/00_00_30_english_3_speakers.transcript.json
View File

@@ -1,382 +0,0 @@
[
{
"word": "Transcription",
"start": 0.0,
"end": 0.6
},
{
"word": "technology",
"start": 0.6,
"end": 1.24
},
{
"word": "has",
"start": 1.24,
"end": 1.5
},
{
"word": "improved",
"start": 1.5,
"end": 1.96
},
{
"word": "so",
"start": 1.96,
"end": 2.32
},
{
"word": "much",
"start": 2.32,
"end": 2.68
},
{
"word": "in",
"start": 2.68,
"end": 2.94
},
{
"word": "the",
"start": 2.94,
"end": 3.02
},
{
"word": "past",
"start": 3.02,
"end": 3.24
},
{
"word": "few",
"start": 3.24,
"end": 3.5
},
{
"word": "years.",
"start": 3.5,
"end": 3.96
},
{
"word": "Have",
"start": 4.56,
"end": 4.74
},
{
"word": "you",
"start": 4.74,
"end": 4.9
},
{
"word": "noticed",
"start": 4.9,
"end": 5.26
},
{
"word": "how",
"start": 5.26,
"end": 5.52
},
{
"word": "accurate",
"start": 5.52,
"end": 6.08
},
{
"word": "real",
"start": 6.08,
"end": 6.42
},
{
"word": "-time",
"start": 6.42,
"end": 6.74
},
{
"word": "speech",
"start": 6.74,
"end": 7.24
},
{
"word": "to",
"start": 7.24,
"end": 7.46
},
{
"word": "text",
"start": 7.46,
"end": 7.78
},
{
"word": "is",
"start": 7.78,
"end": 8.0
},
{
"word": "now?",
"start": 8.0,
"end": 8.3
},
{
"word": "Absolutely.",
"start": 8.7,
"end": 9.16
},
{
"word": "I",
"start": 10.04,
"end": 10.38
},
{
"word": "use",
"start": 10.38,
"end": 10.56
},
{
"word": "it",
"start": 10.56,
"end": 10.76
},
{
"word": "all",
"start": 10.76,
"end": 10.9
},
{
"word": "the",
"start": 10.9,
"end": 11.04
},
{
"word": "time",
"start": 11.04,
"end": 11.32
},
{
"word": "for",
"start": 11.32,
"end": 11.54
},
{
"word": "taking",
"start": 11.54,
"end": 11.86
},
{
"word": "notes",
"start": 11.86,
"end": 12.16
},
{
"word": "during",
"start": 12.16,
"end": 12.54
},
{
"word": "meetings.",
"start": 12.54,
"end": 12.94
},
{
"word": "It's",
"start": 13.6,
"end": 13.8
},
{
"word": "amazing",
"start": 13.8,
"end": 14.1
},
{
"word": "how",
"start": 14.1,
"end": 14.48
},
{
"word": "it",
"start": 14.48,
"end": 14.62
},
{
"word": "can",
"start": 14.62,
"end": 14.74
},
{
"word": "recognise",
"start": 14.74,
"end": 15.24
},
{
"word": "different",
"start": 15.24,
"end": 15.68
},
{
"word": "speakers",
"start": 15.68,
"end": 16.16
},
{
"word": "and",
"start": 16.16,
"end": 16.8
},
{
"word": "even",
"start": 16.8,
"end": 17.1
},
{
"word": "add",
"start": 17.1,
"end": 17.44
},
{
"word": "punctuation.",
"start": 17.44,
"end": 18.36
},
{
"word": "Yeah,",
"start": 18.88,
"end": 19.16
},
{
"word": "but",
"start": 19.36,
"end": 19.52
},
{
"word": "sometimes",
"start": 19.52,
"end": 20.16
},
{
"word": "noise",
"start": 20.16,
"end": 20.54
},
{
"word": "can",
"start": 20.54,
"end": 20.8
},
{
"word": "still",
"start": 20.8,
"end": 21.1
},
{
"word": "cause",
"start": 21.1,
"end": 21.44
},
{
"word": "mistakes.",
"start": 21.44,
"end": 21.94
},
{
"word": "Does",
"start": 22.68,
"end": 22.9
},
{
"word": "this",
"start": 22.9,
"end": 23.12
},
{
"word": "system",
"start": 23.12,
"end": 23.46
},
{
"word": "handle",
"start": 23.46,
"end": 23.88
},
{
"word": "that",
"start": 23.88,
"end": 24.12
},
{
"word": "well?",
"start": 24.12,
"end": 24.42
},
{
"word": "It",
"start": 24.42,
"end": 25.32
},
{
"word": "does",
"start": 25.32,
"end": 25.48
},
{
"word": "a",
"start": 25.48,
"end": 25.62
},
{
"word": "pretty",
"start": 25.62,
"end": 25.88
},
{
"word": "good",
"start": 25.88,
"end": 26.08
},
{
"word": "job",
"start": 26.08,
"end": 26.32
},
{
"word": "filtering",
"start": 26.32,
"end": 26.8
},
{
"word": "noise,",
"start": 26.8,
"end": 27.18
},
{
"word": "especially",
"start": 27.36,
"end": 28.0
},
{
"word": "with",
"start": 28.0,
"end": 28.28
},
{
"word": "models",
"start": 28.28,
"end": 28.62
},
{
"word": "that",
"start": 28.62,
"end": 28.94
},
{
"word": "use",
"start": 28.94,
"end": 29.22
},
{
"word": "voice",
"start": 29.22,
"end": 29.54
},
{
"word": "active.",
"start": 29.54,
"end": 29.9
}
]

58
audio_tests/generate_transcripts.py
View File

@@ -1,58 +0,0 @@
#!/usr/bin/env python3
"""Generate word-level timestamped transcripts using faster-whisper (offline).
Produces one JSON file per audio with: [{word, start, end}, ...]
"""
import json
import os
from faster_whisper import WhisperModel
AUDIO_DIR = os.path.dirname(os.path.abspath(__file__))
FILES = [
("00_00_07_english_1_speaker.wav", "en"),
("00_00_16_french_1_speaker.wav", "fr"),
("00_00_30_english_3_speakers.wav", "en"),
]
def main():
print("Loading faster-whisper model (base, cpu, float32)...")
model = WhisperModel("base", device="cpu", compute_type="float32")
for filename, lang in FILES:
audio_path = os.path.join(AUDIO_DIR, filename)
out_path = os.path.join(
AUDIO_DIR, filename.rsplit(".", 1)[0] + ".transcript.json"
)
print(f"\n{'='*60}")
print(f"Transcribing: {filename} (language={lang})")
print(f"{'='*60}")
segments, info = model.transcribe(
audio_path, word_timestamps=True, language=lang
)
words = []
for segment in segments:
if segment.words:
for w in segment.words:
words.append({
"word": w.word.strip(),
"start": round(w.start, 3),
"end": round(w.end, 3),
})
print(f" {w.start:6.2f} - {w.end:6.2f} {w.word.strip()}")
with open(out_path, "w", encoding="utf-8") as f:
json.dump(words, f, indent=2, ensure_ascii=False)
print(f"\n -> {len(words)} words written to {os.path.basename(out_path)}")
print("\nDone.")
if __name__ == "__main__":
main()

BIN
benchmark_chart.png
View File

Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 69 KiB

BIN
benchmark_scatter.png
View File

Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 95 KiB

BIN
benchmark_scatter_en_aware.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 100 KiB

BIN
benchmark_scatter_fr_aware.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 100 KiB

BIN
benchmarks/h100/acl6060_per_talk.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 63 KiB

BIN
benchmarks/h100/bars_wer_rtf_latency.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 130 KiB

124
benchmarks/h100/bench_voxtral_hf_batch.py Normal file
View File

@@ -0,0 +1,124 @@
#!/usr/bin/env python3
"""Standalone Voxtral benchmark — no whisperlivekit imports."""
import json, logging, re, time, wave, queue, threading
import numpy as np
import torch
logging.basicConfig(level=logging.WARNING)
for n in ["transformers","torch","httpx"]:
logging.getLogger(n).setLevel(logging.ERROR)
from jiwer import wer as compute_wer
from transformers import AutoProcessor, VoxtralRealtimeForConditionalGeneration, TextIteratorStreamer
def norm(t):
return re.sub(r' +', ' ', re.sub(r'[^a-z0-9 ]', ' ', t.lower())).strip()
def load_audio(path):
with wave.open(path, 'r') as wf:
return np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16).astype(np.float32) / 32768.0
# Load model
print("Loading Voxtral-Mini-4B...", flush=True)
MODEL_ID = "mistralai/Voxtral-Mini-4B-Realtime-2602"
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = VoxtralRealtimeForConditionalGeneration.from_pretrained(
MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda:0",
)
print(f"Loaded, GPU: {torch.cuda.memory_allocated()/1e9:.1f} GB", flush=True)
def transcribe_batch(audio_np):
"""Simple batch transcription (not streaming)."""
# Voxtral expects audio as input_features from processor
inputs = processor(
audio=audio_np, sampling_rate=16000, return_tensors="pt",
).to("cuda:0").to(torch.bfloat16)
t0 = time.perf_counter()
with torch.inference_mode():
generated = model.generate(**inputs, max_new_tokens=1024)
t1 = time.perf_counter()
text = processor.batch_decode(generated, skip_special_tokens=True)[0].strip()
return text, t1 - t0
# 1. LibriSpeech test-clean
print("\n=== Voxtral / LibriSpeech test-clean ===", flush=True)
clean = json.load(open("/home/cloud/benchmark_data/metadata.json"))
wers = []; ta = tp = 0
for i, s in enumerate(clean):
audio = load_audio(s['path'])
hyp, pt = transcribe_batch(audio)
w = compute_wer(norm(s['reference']), norm(hyp))
wers.append(w); ta += s['duration']; tp += pt
if i < 3 or i % 20 == 0:
print(f" [{i}] {s['duration']:.1f}s RTF={pt/s['duration']:.2f} WER={w:.1%} | {hyp[:60]}", flush=True)
clean_wer = np.mean(wers); clean_rtf = tp/ta
print(f" CLEAN: WER {clean_wer:.2%}, RTF {clean_rtf:.3f} ({len(clean)} samples, {ta:.0f}s)")
# 2. LibriSpeech test-other
print("\n=== Voxtral / LibriSpeech test-other ===", flush=True)
other = json.load(open("/home/cloud/benchmark_data/metadata_other.json"))
wers2 = []; ta2 = tp2 = 0
for i, s in enumerate(other):
audio = load_audio(s['path'])
hyp, pt = transcribe_batch(audio)
w = compute_wer(norm(s['reference']), norm(hyp))
wers2.append(w); ta2 += s['duration']; tp2 += pt
if i < 3 or i % 20 == 0:
print(f" [{i}] {s['duration']:.1f}s RTF={pt/s['duration']:.2f} WER={w:.1%}", flush=True)
other_wer = np.mean(wers2); other_rtf = tp2/ta2
print(f" OTHER: WER {other_wer:.2%}, RTF {other_rtf:.3f} ({len(other)} samples, {ta2:.0f}s)")
# 3. ACL6060
print("\n=== Voxtral / ACL6060 ===", flush=True)
acl_results = []
for talk in ["110", "117", "268", "367", "590"]:
audio = load_audio(f"/home/cloud/acl6060_audio/2022.acl-long.{talk}.wav")
dur = len(audio) / 16000
gw = []
with open(f"/home/cloud/iwslt26-sst/inputs/en/acl6060.ts/gold-jsonl/2022.acl-long.{talk}.jsonl") as f:
for line in f:
gw.append(json.loads(line)["text"].strip())
gold = " ".join(gw)
# For long audio, process in 30s chunks
all_hyp = []
t0 = time.perf_counter()
chunk_size = 30 * 16000
for start in range(0, len(audio), chunk_size):
chunk = audio[start:start + chunk_size]
if len(chunk) < 1600: # skip very short tail
continue
hyp, _ = transcribe_batch(chunk)
all_hyp.append(hyp)
t1 = time.perf_counter()
full_hyp = " ".join(all_hyp)
w = compute_wer(norm(gold), norm(full_hyp))
rtf = (t1 - t0) / dur
acl_results.append({"talk": talk, "wer": w, "rtf": rtf, "dur": dur})
print(f" Talk {talk}: {dur:.0f}s, WER {w:.2%}, RTF {rtf:.3f}", flush=True)
acl_wer = np.mean([r["wer"] for r in acl_results])
acl_rtf = np.mean([r["rtf"] for r in acl_results])
print(f" ACL6060 AVERAGE: WER {acl_wer:.2%}, RTF {acl_rtf:.3f}")
# Summary
print(f"\n{'='*60}")
print(f" VOXTRAL BENCHMARK SUMMARY (H100 80GB)")
print(f"{'='*60}")
print(f" {'Dataset':>25} {'WER':>7} {'RTF':>7}")
print(f" {'-'*42}")
print(f" {'LibriSpeech clean':>25} {clean_wer:>6.2%} {clean_rtf:>7.3f}")
print(f" {'LibriSpeech other':>25} {other_wer:>6.2%} {other_rtf:>7.3f}")
print(f" {'ACL6060 (5 talks)':>25} {acl_wer:>6.2%} {acl_rtf:>7.3f}")
results = {
"clean": {"avg_wer": round(float(clean_wer), 4), "rtf": round(float(clean_rtf), 3)},
"other": {"avg_wer": round(float(other_wer), 4), "rtf": round(float(other_rtf), 3)},
"acl6060": {"avg_wer": round(float(acl_wer), 4), "avg_rtf": round(float(acl_rtf), 3),
"talks": [{k: (round(float(v), 4) if isinstance(v, (float, np.floating)) else v) for k, v in r.items()} for r in acl_results]},
}
json.dump(results, open("/home/cloud/bench_voxtral_results.json", "w"), indent=2)
print(f"\nSaved to /home/cloud/bench_voxtral_results.json")

122
benchmarks/h100/bench_voxtral_vllm_realtime.py Normal file
View File

@@ -0,0 +1,122 @@
#!/usr/bin/env python3
"""Benchmark Voxtral via vLLM WebSocket /v1/realtime — proper streaming."""
import asyncio, json, base64, time, wave, re, os
import numpy as np
import websockets
import librosa
from jiwer import wer as compute_wer
MODEL = "mistralai/Voxtral-Mini-4B-Realtime-2602"
WS_URI = "ws://localhost:8000/v1/realtime"
def norm(t):
return re.sub(r' +', ' ', re.sub(r'[^a-z0-9 ]', ' ', t.lower())).strip()
async def transcribe(audio_path, max_tokens=4096):
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
pcm16 = (audio * 32767).astype(np.int16).tobytes()
dur = len(audio) / 16000
t0 = time.time()
transcript = ""
first_token_time = None
async with websockets.connect(WS_URI, max_size=2**24) as ws:
await ws.recv() # session.created
await ws.send(json.dumps({"type": "session.update", "model": MODEL}))
await ws.send(json.dumps({"type": "input_audio_buffer.commit"})) # signal ready
# Send audio in 4KB chunks
for i in range(0, len(pcm16), 4096):
await ws.send(json.dumps({
"type": "input_audio_buffer.append",
"audio": base64.b64encode(pcm16[i:i+4096]).decode(),
}))
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
while True:
try:
msg = json.loads(await asyncio.wait_for(ws.recv(), timeout=120))
if msg["type"] == "transcription.delta":
d = msg.get("delta", "")
if d.strip() and first_token_time is None:
first_token_time = time.time() - t0
transcript += d
elif msg["type"] == "transcription.done":
transcript = msg.get("text", transcript)
break
elif msg["type"] == "error":
break
except asyncio.TimeoutError:
break
elapsed = time.time() - t0
return transcript.strip(), dur, elapsed / dur, first_token_time or elapsed
async def main():
# Warmup
print("Warmup...", flush=True)
await transcribe("/home/cloud/benchmark_data/librispeech_clean_0000.wav")
# LibriSpeech clean (full 91 samples)
print("\n=== Voxtral vLLM Realtime / LibriSpeech clean ===", flush=True)
clean = json.load(open("/home/cloud/benchmark_data/metadata.json"))
wers = []; ta = tp = 0
for i, s in enumerate(clean):
hyp, dur, rtf, fwl = await transcribe(s['path'])
w = compute_wer(norm(s['reference']), norm(hyp)) if hyp else 1.0
wers.append(w); ta += dur; tp += dur * rtf
if i < 3 or i % 20 == 0:
print(f" [{i}] {dur:.1f}s RTF={rtf:.3f} FWL={fwl:.2f}s WER={w:.1%} | {hyp[:60]}", flush=True)
clean_wer = np.mean(wers); clean_rtf = tp / ta
print(f" CLEAN ({len(clean)}): WER {clean_wer:.2%}, RTF {clean_rtf:.3f}\n", flush=True)
# LibriSpeech other (full 133 samples)
print("=== Voxtral vLLM Realtime / LibriSpeech other ===", flush=True)
other = json.load(open("/home/cloud/benchmark_data/metadata_other.json"))
wers2 = []; ta2 = tp2 = 0
for i, s in enumerate(other):
hyp, dur, rtf, fwl = await transcribe(s['path'])
w = compute_wer(norm(s['reference']), norm(hyp)) if hyp else 1.0
wers2.append(w); ta2 += dur; tp2 += dur * rtf
if i < 3 or i % 20 == 0:
print(f" [{i}] {dur:.1f}s RTF={rtf:.3f} WER={w:.1%}", flush=True)
other_wer = np.mean(wers2); other_rtf = tp2 / ta2
print(f" OTHER ({len(other)}): WER {other_wer:.2%}, RTF {other_rtf:.3f}\n", flush=True)
# ACL6060 talks
print("=== Voxtral vLLM Realtime / ACL6060 ===", flush=True)
acl = []
for talk in ["110", "117", "268", "367", "590"]:
gw = []
with open(f"/home/cloud/iwslt26-sst/inputs/en/acl6060.ts/gold-jsonl/2022.acl-long.{talk}.jsonl") as f:
for line in f: gw.append(json.loads(line)["text"].strip())
gold = " ".join(gw)
hyp, dur, rtf, fwl = await transcribe(f"/home/cloud/acl6060_audio/2022.acl-long.{talk}.wav")
w = compute_wer(norm(gold), norm(hyp)) if hyp else 1.0
acl.append({"talk": talk, "wer": round(float(w),4), "rtf": round(float(rtf),3), "dur": round(dur,1)})
print(f" Talk {talk}: {dur:.0f}s, WER {w:.2%}, RTF {rtf:.3f}, FWL {fwl:.2f}s", flush=True)
acl_wer = np.mean([r["wer"] for r in acl])
acl_rtf = np.mean([r["rtf"] for r in acl])
print(f" ACL6060 AVERAGE: WER {acl_wer:.2%}, RTF {acl_rtf:.3f}\n", flush=True)
# Summary
print(f"{'='*55}")
print(f" VOXTRAL vLLM REALTIME BENCHMARK (H100)")
print(f"{'='*55}")
print(f" LS clean ({len(clean)}): WER {clean_wer:.2%}, RTF {clean_rtf:.3f}")
print(f" LS other ({len(other)}): WER {other_wer:.2%}, RTF {other_rtf:.3f}")
print(f" ACL6060 (5): WER {acl_wer:.2%}, RTF {acl_rtf:.3f}")
results = {
"clean": {"avg_wer": round(float(clean_wer),4), "rtf": round(float(clean_rtf),3), "n": len(clean)},
"other": {"avg_wer": round(float(other_wer),4), "rtf": round(float(other_rtf),3), "n": len(other)},
"acl6060": {"avg_wer": round(float(acl_wer),4), "avg_rtf": round(float(acl_rtf),3), "talks": acl},
}
json.dump(results, open("/home/cloud/bench_voxtral_realtime_results.json", "w"), indent=2)
print(f"\n Saved to /home/cloud/bench_voxtral_realtime_results.json")
asyncio.run(main())

270
benchmarks/h100/generate_figures.py Normal file
View File

@@ -0,0 +1,270 @@
#!/usr/bin/env python3
"""
Generate polished benchmark figures for WhisperLiveKit H100 results.
Reads data from results.json, outputs PNGs to this directory.
Run: python3 benchmarks/h100/generate_figures.py
"""
import json
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
DIR = os.path.dirname(os.path.abspath(__file__))
DATA = json.load(open(os.path.join(DIR, "results.json")))
# ── Style constants ──
COLORS = {
"whisper": "#d63031",
"qwen_b": "#6c5ce7",
"qwen_s": "#00b894",
"voxtral": "#fdcb6e",
"fw_m5": "#74b9ff",
"mlx_m5": "#55efc4",
"vox_m5": "#ffeaa7",
}
plt.rcParams.update({
"font.family": "sans-serif",
"font.size": 11,
"axes.spines.top": False,
"axes.spines.right": False,
})
def _save(fig, name):
path = os.path.join(DIR, name)
fig.savefig(path, dpi=180, bbox_inches="tight", facecolor="white")
plt.close(fig)
print(f" {name}")
# ──────────────────────────────────────────────────────────
# Figure 1: WER vs RTF scatter — H100 (LibriSpeech clean)
# ──────────────────────────────────────────────────────────
def fig_scatter_clean():
ls = DATA["librispeech_clean"]["systems"]
m5 = DATA["m5_reference"]["systems"]
fig, ax = plt.subplots(figsize=(9, 7.5))
ax.axhspan(0, 10, color="#f0fff0", alpha=0.5, zorder=0)
# M5 (ghost dots)
for k, v in m5.items():
ax.scatter(v["rtf"], v["wer"], s=50, c="silver", marker="o",
alpha=0.22, zorder=2, linewidths=0.4, edgecolors="gray")
# H100 systems — (name, data, color, marker, size, label_x_off, label_y_off)
pts = [
("Whisper large-v3", ls["whisper_large_v3_batch"], COLORS["whisper"], "h", 240, -8, -16),
("Qwen3-ASR 0.6B (batch)", ls["qwen3_0.6b_batch"], COLORS["qwen_b"], "h", 170, 8, 6),
("Qwen3-ASR 1.7B (batch)", ls["qwen3_1.7b_batch"], COLORS["qwen_b"], "h", 240, 8, -16),
("Voxtral 4B (vLLM)", ls["voxtral_4b_vllm_realtime"], COLORS["voxtral"], "D", 260, 8, 6),
("Qwen3 0.6B SimulStream+KV", ls["qwen3_0.6b_simulstream_kv"], COLORS["qwen_s"], "s", 220, 8, 6),
("Qwen3 1.7B SimulStream+KV", ls["qwen3_1.7b_simulstream_kv"], COLORS["qwen_s"], "s", 280, 8, -16),
]
for name, d, color, marker, sz, lx, ly in pts:
ax.scatter(d["rtf"], d["wer"], s=sz, c=color, marker=marker,
edgecolors="white", linewidths=1.5, zorder=5)
ax.annotate(name, (d["rtf"], d["wer"]), fontsize=8.5, fontweight="bold",
xytext=(lx, ly), textcoords="offset points",
arrowprops=dict(arrowstyle="-", color="#aaa", lw=0.5))
ax.set_xlabel("RTF (lower = faster)")
ax.set_ylabel("WER % (lower = better)")
ax.set_title("Speed vs Accuracy — LibriSpeech test-clean (H100 80 GB)",
fontsize=13, fontweight="bold", pad=12)
ax.set_xlim(-0.005, 0.20)
ax.set_ylim(-0.3, 10)
ax.grid(True, alpha=0.12)
legend = [
mpatches.Patch(color=COLORS["whisper"], label="Whisper large-v3"),
mpatches.Patch(color=COLORS["qwen_b"], label="Qwen3-ASR (batch)"),
mpatches.Patch(color=COLORS["qwen_s"], label="Qwen3 SimulStream+KV"),
mpatches.Patch(color=COLORS["voxtral"], label="Voxtral 4B (vLLM)"),
plt.Line2D([0],[0], marker="h", color="w", mfc="gray", ms=8, label="Batch"),
plt.Line2D([0],[0], marker="s", color="w", mfc="gray", ms=8, label="Streaming"),
]
ax.legend(handles=legend, fontsize=8.5, loc="upper right", framealpha=0.85, ncol=2)
_save(fig, "wer_vs_rtf_clean.png")
# ──────────────────────────────────────────────────────────
# Figure 2: ACL6060 conference talks — the realistic test
# ──────────────────────────────────────────────────────────
def fig_scatter_acl6060():
acl = DATA["acl6060"]["systems"]
fig, ax = plt.subplots(figsize=(10, 6.5))
ax.axhspan(0, 15, color="#f0fff0", alpha=0.4, zorder=0)
pts = [
("Voxtral 4B\n(vLLM Realtime)", acl["voxtral_4b_vllm_realtime"], COLORS["voxtral"], "D", 380),
("Qwen3 1.7B\nSimulStream+KV", acl["qwen3_1.7b_simulstream_kv"], COLORS["qwen_s"], "s", 380),
("Qwen3 0.6B\nSimulStream+KV", acl["qwen3_0.6b_simulstream_kv"], COLORS["qwen_s"], "s", 260),
("Whisper large-v3\n(batch)", acl["whisper_large_v3_batch"], COLORS["whisper"], "h", 320),
]
label_off = [(10, -12), (10, 6), (10, 6), (10, 6)]
for (name, d, color, marker, sz), (lx, ly) in zip(pts, label_off):
wer = d["avg_wer"]; rtf = d["avg_rtf"]
ax.scatter(rtf, wer, s=sz, c=color, marker=marker,
edgecolors="white", linewidths=1.5, zorder=5)
ax.annotate(name, (rtf, wer), fontsize=9.5, fontweight="bold",
xytext=(lx, ly), textcoords="offset points",
arrowprops=dict(arrowstyle="-", color="#aaa", lw=0.6))
# Cascade annotation
ax.annotate("Full STT+MT cascade\nRTF 0.15 (real-time)",
xy=(0.151, 1), xytext=(0.25, 4),
fontsize=9, fontstyle="italic", color="#1565c0",
arrowprops=dict(arrowstyle="->", color="#1565c0", lw=1.5),
bbox=dict(boxstyle="round,pad=0.3", fc="#e3f2fd", ec="#90caf9", alpha=0.9))
ax.set_xlabel("RTF (lower = faster)")
ax.set_ylabel("WER % (lower = better)")
ax.set_title("ACL6060 Conference Talks — 5 talks, 58 min (H100 80 GB)",
fontsize=13, fontweight="bold", pad=12)
ax.set_xlim(-0.005, 0.30)
ax.set_ylim(-1, 26)
ax.grid(True, alpha=0.12)
_save(fig, "wer_vs_rtf_acl6060.png")
# ──────────────────────────────────────────────────────────
# Figure 3: Bar chart — WER + RTF side-by-side
# ──────────────────────────────────────────────────────────
def fig_bars():
names = [
"Whisper\nlarge-v3", "Voxtral 4B\n(vLLM)", "Qwen3 0.6B\n(batch)",
"Qwen3 1.7B\n(batch)", "Qwen3 0.6B\nSimulStream", "Qwen3 1.7B\nSimulStream",
]
wer_c = [2.02, 2.71, 2.30, 2.46, 6.44, 8.09]
wer_o = [7.79, 9.26, 6.12, 5.34, 9.27, 9.56]
rtf_c = [0.071, 0.137, 0.065, 0.069, 0.109, 0.117]
fwl = [472, 137, 432, 457, 91, 94] # ms
cols = [COLORS["whisper"], COLORS["voxtral"], COLORS["qwen_b"],
COLORS["qwen_b"], COLORS["qwen_s"], COLORS["qwen_s"]]
cols_l = ["#ff7675", "#ffeaa7", "#a29bfe", "#a29bfe", "#55efc4", "#55efc4"]
x = np.arange(len(names))
fig, axes = plt.subplots(1, 3, figsize=(16, 6))
# WER
ax = axes[0]; w = 0.36
ax.bar(x - w/2, wer_c, w, color=cols, alpha=0.9, edgecolor="white", label="test-clean")
ax.bar(x + w/2, wer_o, w, color=cols_l, alpha=0.65, edgecolor="white", label="test-other")
ax.set_ylabel("WER %"); ax.set_title("Word Error Rate", fontweight="bold")
ax.set_xticks(x); ax.set_xticklabels(names, fontsize=7.5, rotation=25, ha="right")
ax.legend(fontsize=8); ax.grid(axis="y", alpha=0.15)
for i, v in enumerate(wer_c):
ax.text(i - w/2, v + 0.2, f"{v:.1f}", ha="center", fontsize=7, fontweight="bold")
# RTF
ax = axes[1]
ax.bar(x, rtf_c, 0.55, color=cols, alpha=0.9, edgecolor="white")
ax.set_ylabel("RTF (lower = faster)"); ax.set_title("Real-Time Factor (test-clean)", fontweight="bold")
ax.set_xticks(x); ax.set_xticklabels(names, fontsize=7.5, rotation=25, ha="right")
ax.grid(axis="y", alpha=0.15)
for i, v in enumerate(rtf_c):
ax.text(i, v + 0.003, f"{v:.3f}", ha="center", fontsize=8, fontweight="bold")
# First-word latency
ax = axes[2]
ax.bar(x, fwl, 0.55, color=cols, alpha=0.9, edgecolor="white")
ax.set_ylabel("ms"); ax.set_title("First Word Latency", fontweight="bold")
ax.set_xticks(x); ax.set_xticklabels(names, fontsize=7.5, rotation=25, ha="right")
ax.grid(axis="y", alpha=0.15)
for i, v in enumerate(fwl):
ax.text(i, v + 8, f"{v}", ha="center", fontsize=8, fontweight="bold")
fig.suptitle("LibriSpeech Benchmark — H100 80 GB", fontsize=14, fontweight="bold")
plt.tight_layout()
_save(fig, "bars_wer_rtf_latency.png")
# ──────────────────────────────────────────────────────────
# Figure 4: Clean vs Other robustness
# ──────────────────────────────────────────────────────────
def fig_robustness():
models = [
("Whisper large-v3", 2.02, 7.79, COLORS["whisper"], "h", 280),
("Qwen3 0.6B (batch)", 2.30, 6.12, COLORS["qwen_b"], "h", 180),
("Qwen3 1.7B (batch)", 2.46, 5.34, COLORS["qwen_b"], "h", 280),
("Voxtral 4B (vLLM)", 2.71, 9.26, COLORS["voxtral"], "D", 280),
("Qwen3 0.6B\nSimulStream", 6.44, 9.27, COLORS["qwen_s"], "s", 240),
("Qwen3 1.7B\nSimulStream", 8.09, 9.56, COLORS["qwen_s"], "s", 300),
]
# Manual label offsets — carefully placed to avoid overlap
offsets = [(-55, 10), (8, 10), (8, -18), (-55, -18), (-10, 12), (10, -18)]
fig, ax = plt.subplots(figsize=(8.5, 7))
ax.plot([0, 13], [0, 13], "--", color="#ccc", lw=1, zorder=1)
ax.fill_between([0, 13], [0, 13], [13, 13], color="#fff5f5", alpha=0.5, zorder=0)
ax.text(4, 11, "degrades more\non noisy audio", fontsize=9, color="#bbb", fontstyle="italic")
for (name, wc, wo, color, marker, sz), (lx, ly) in zip(models, offsets):
ax.scatter(wc, wo, s=sz, c=color, marker=marker,
edgecolors="white", linewidths=1.5, zorder=5)
ax.annotate(name, (wc, wo), fontsize=8.5, fontweight="bold",
xytext=(lx, ly), textcoords="offset points",
arrowprops=dict(arrowstyle="-", color="#aaa", lw=0.6))
deg = wo - wc
ax.annotate(f"+{deg:.1f}%", (wc, wo), fontsize=7, color="#999",
xytext=(-6, -13), textcoords="offset points")
ax.set_xlabel("WER % on test-clean")
ax.set_ylabel("WER % on test-other")
ax.set_title("Clean vs Noisy Robustness (H100 80 GB)", fontsize=13, fontweight="bold", pad=12)
ax.set_xlim(-0.3, 12); ax.set_ylim(-0.3, 12)
ax.set_aspect("equal"); ax.grid(True, alpha=0.12)
_save(fig, "robustness_clean_vs_other.png")
# ──────────────────────────────────────────────────────────
# Figure 5: ACL6060 per-talk breakdown (Qwen3 vs Voxtral)
# ──────────────────────────────────────────────────────────
def fig_per_talk():
q = DATA["acl6060"]["systems"]["qwen3_1.7b_simulstream_kv"]["per_talk"]
v = DATA["acl6060"]["systems"]["voxtral_4b_vllm_realtime"]["per_talk"]
talks = DATA["acl6060"]["talks"]
fig, ax = plt.subplots(figsize=(9, 5))
x = np.arange(len(talks)); w = 0.35
bars_v = ax.bar(x - w/2, [v[t] for t in talks], w, color=COLORS["voxtral"],
edgecolor="white", label="Voxtral 4B (vLLM)")
bars_q = ax.bar(x + w/2, [q[t] for t in talks], w, color=COLORS["qwen_s"],
edgecolor="white", label="Qwen3 1.7B SimulStream+KV")
for bar in bars_v:
ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
f"{bar.get_height():.1f}", ha="center", fontsize=8)
for bar in bars_q:
ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
f"{bar.get_height():.1f}", ha="center", fontsize=8)
ax.set_xlabel("ACL6060 Talk ID")
ax.set_ylabel("WER %")
ax.set_title("Per-Talk WER — ACL6060 Conference Talks (H100 80 GB)",
fontsize=13, fontweight="bold", pad=12)
ax.set_xticks(x); ax.set_xticklabels([f"Talk {t}" for t in talks])
ax.legend(fontsize=9); ax.grid(axis="y", alpha=0.15)
ax.set_ylim(0, 18)
_save(fig, "acl6060_per_talk.png")
if __name__ == "__main__":
print("Generating H100 benchmark figures...")
fig_scatter_clean()
fig_scatter_acl6060()
fig_bars()
fig_robustness()
fig_per_talk()
print("Done!")

56
benchmarks/h100/results.json Normal file
View File

@@ -0,0 +1,56 @@
{
"hardware": "NVIDIA H100 80GB HBM3, CUDA 12.4, Driver 550.163",
"date": "2026-03-15",
"librispeech_clean": {
"n_samples": 91,
"total_audio_s": 602,
"systems": {
"whisper_large_v3_batch": {"wer": 2.02, "rtf": 0.071, "first_word_latency_s": 0.472},
"qwen3_0.6b_batch": {"wer": 2.30, "rtf": 0.065, "first_word_latency_s": 0.432},
"qwen3_1.7b_batch": {"wer": 2.46, "rtf": 0.069, "first_word_latency_s": 0.457},
"voxtral_4b_vllm_realtime": {"wer": 2.71, "rtf": 0.137, "first_word_latency_s": 0.137},
"qwen3_0.6b_simulstream_kv": {"wer": 6.44, "rtf": 0.109, "first_word_latency_s": 0.091},
"qwen3_1.7b_simulstream_kv": {"wer": 8.09, "rtf": 0.117, "first_word_latency_s": 0.094}
}
},
"librispeech_other": {
"n_samples": 133,
"total_audio_s": 600,
"systems": {
"qwen3_1.7b_batch": {"wer": 5.34, "rtf": 0.088},
"qwen3_0.6b_batch": {"wer": 6.12, "rtf": 0.086},
"whisper_large_v3_batch": {"wer": 7.79, "rtf": 0.092},
"qwen3_0.6b_simulstream_kv": {"wer": 9.27, "rtf": 0.127},
"voxtral_4b_vllm_realtime": {"wer": 9.26, "rtf": 0.144},
"qwen3_1.7b_simulstream_kv": {"wer": 9.56, "rtf": 0.140}
}
},
"acl6060": {
"description": "5 ACL 2022 conference talks, 58 min total",
"talks": ["110", "117", "268", "367", "590"],
"systems": {
"voxtral_4b_vllm_realtime": {"avg_wer": 7.83, "avg_rtf": 0.203, "per_talk": {"110": 5.18, "117": 2.24, "268": 14.88, "367": 9.40, "590": 7.45}},
"qwen3_1.7b_simulstream_kv": {"avg_wer": 9.20, "avg_rtf": 0.074, "per_talk": {"110": 5.59, "117": 8.12, "268": 12.25, "367": 12.29, "590": 7.77}},
"qwen3_0.6b_simulstream_kv": {"avg_wer": 13.21, "avg_rtf": 0.098},
"whisper_large_v3_batch": {"avg_wer": 22.53, "avg_rtf": 0.125}
}
},
"m5_reference": {
"description": "MacBook M5 results (from WLK scatter benchmarks)",
"systems": {
"fw_la_base": {"wer": 17.0, "rtf": 0.82},
"fw_la_small": {"wer": 8.6, "rtf": 0.76},
"fw_ss_base": {"wer": 7.8, "rtf": 0.46},
"fw_ss_small": {"wer": 7.0, "rtf": 0.90},
"mlx_ss_base": {"wer": 7.7, "rtf": 0.34},
"mlx_ss_small": {"wer": 6.5, "rtf": 0.68},
"voxtral_mlx": {"wer": 7.0, "rtf": 0.26},
"qwen3_mlx_0.6b":{"wer": 5.5, "rtf": 0.55},
"qwen3_0.6b_batch":{"wer":24.0, "rtf": 1.42}
}
}
}

BIN
benchmarks/h100/robustness_clean_vs_other.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 101 KiB

BIN
benchmarks/h100/wer_vs_rtf_acl6060.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 95 KiB

BIN
benchmarks/h100/wer_vs_rtf_clean.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 110 KiB

3
pyproject.toml
View File

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "whisperlivekit"
version = "0.2.20"
description = "Real-time speech-to-text with speaker diarization using Whisper"
description = "Real-time speech-to-text models"
readme = "README.md"
authors = [{ name = "Quentin Fuxa" }]
license = { file = "LICENSE" }
@@ -144,6 +144,7 @@ packages = [
"whisperlivekit.local_agreement",
"whisperlivekit.voxtral_mlx",
"whisperlivekit.silero_vad_models",
"whisperlivekit.benchmark",
]
[tool.setuptools.package-data]

290
run_benchmark.py
View File

@@ -1,290 +0,0 @@
#!/usr/bin/env python3
"""
Comprehensive benchmark runner for WhisperLiveKit.
Tests all available backend+policy combinations across multiple audio files,
model sizes, and VAC on/off configurations. Outputs structured JSON that
is consumed by the report generator.
Usage:
python run_benchmark.py # full benchmark
python run_benchmark.py --quick # subset (tiny models, fewer combos)
python run_benchmark.py --json results.json # custom output path
"""
import argparse
import asyncio
import gc
import json
import logging
import platform
import subprocess
import sys
import time
from dataclasses import asdict
from pathlib import Path
logging.basicConfig(level=logging.WARNING, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("benchmark")
logger.setLevel(logging.INFO)
# Re-use harness functions
sys.path.insert(0, str(Path(__file__).parent))
from test_backend_offline import (
AUDIO_TESTS_DIR,
SAMPLE_RATE,
create_engine,
discover_audio_files,
download_sample_audio,
load_audio,
run_test,
)
CACHE_DIR = Path(__file__).parent / ".test_cache"
def get_system_info() -> dict:
"""Collect system metadata for the report."""
info = {
"platform": platform.platform(),
"machine": platform.machine(),
"processor": platform.processor(),
"python_version": platform.python_version(),
}
# macOS: get chip info
try:
chip = subprocess.check_output(
["sysctl", "-n", "machdep.cpu.brand_string"], text=True
).strip()
info["cpu"] = chip
except Exception:
info["cpu"] = platform.processor()
# RAM
try:
mem_bytes = int(
subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True).strip()
)
info["ram_gb"] = round(mem_bytes / (1024**3))
except Exception:
info["ram_gb"] = None
# Backend versions
versions = {}
try:
import faster_whisper
versions["faster-whisper"] = faster_whisper.__version__
except ImportError:
pass
try:
import mlx_whisper # noqa: F401
versions["mlx-whisper"] = "installed"
except ImportError:
pass
try:
import mlx.core as mx
versions["mlx"] = mx.__version__
except ImportError:
pass
try:
import transformers
versions["transformers"] = transformers.__version__
except ImportError:
pass
try:
import torch
versions["torch"] = torch.__version__
except ImportError:
pass
info["backend_versions"] = versions
return info
def detect_combos(quick: bool = False) -> list:
"""Build list of (backend, policy, model_size) combos to test."""
combos = []
# Model sizes to test
model_sizes = ["tiny", "base", "small"] if not quick else ["tiny", "base"]
# faster-whisper
try:
import faster_whisper # noqa: F401
for model in model_sizes:
combos.append({"backend": "faster-whisper", "policy": "localagreement", "model": model})
combos.append({"backend": "faster-whisper", "policy": "simulstreaming", "model": model})
except ImportError:
pass
# mlx-whisper
try:
import mlx_whisper # noqa: F401
for model in model_sizes:
combos.append({"backend": "mlx-whisper", "policy": "localagreement", "model": model})
combos.append({"backend": "mlx-whisper", "policy": "simulstreaming", "model": model})
except ImportError:
pass
# voxtral-mlx (single model, single policy)
try:
from whisperlivekit.voxtral_mlx import VoxtralMLXModel # noqa: F401
combos.append({"backend": "voxtral-mlx", "policy": "voxtral", "model": ""})
except ImportError:
pass
# voxtral HF (single model, single policy)
try:
from transformers import AutoModelForSpeechSeq2Seq # noqa: F401
combos.append({"backend": "voxtral", "policy": "voxtral", "model": ""})
except ImportError:
pass
return combos
def collect_audio_files() -> list:
"""Collect all benchmark audio files."""
files = []
# audio_tests/ directory
if AUDIO_TESTS_DIR.is_dir():
files.extend(discover_audio_files(str(AUDIO_TESTS_DIR)))
# JFK sample
jfk = CACHE_DIR / "jfk.wav"
if not jfk.exists():
jfk = download_sample_audio()
if jfk.exists():
files.append(jfk)
return files
async def run_single_combo(
combo: dict, audio_files: list, vac: bool, lan: str, max_duration: float,
) -> list:
"""Run one backend+policy+model combo across all audio files."""
backend = combo["backend"]
policy = combo["policy"]
model = combo["model"]
results = []
try:
engine = create_engine(
backend=backend,
model_size=model,
lan=lan,
vac=vac,
policy=policy,
)
# Quiet noisy loggers
for mod in (
"whisperlivekit.audio_processor",
"whisperlivekit.simul_whisper",
"whisperlivekit.tokens_alignment",
"whisperlivekit.simul_whisper.align_att_base",
"whisperlivekit.simul_whisper.simul_whisper",
):
logging.getLogger(mod).setLevel(logging.WARNING)
for audio_path in audio_files:
duration = len(load_audio(str(audio_path))) / SAMPLE_RATE
if duration > max_duration:
logger.info(f" Skipping {audio_path.name} ({duration:.0f}s > {max_duration:.0f}s)")
continue
file_lan = lan
if "french" in audio_path.name.lower() and lan == "en":
file_lan = "fr"
audio = load_audio(str(audio_path))
result = await run_test(
engine, audio, chunk_ms=100, realtime=False,
audio_file=audio_path.name, backend=backend,
policy=policy, lan=file_lan,
)
# Tag with extra metadata
result_dict = asdict(result)
result_dict["model_size"] = model
result_dict["vac"] = vac
results.append(result_dict)
except Exception as e:
logger.error(f" FAILED: {e}")
import traceback
traceback.print_exc()
return results
async def run_full_benchmark(combos, audio_files, max_duration=60.0):
"""Run all combos with VAC on and off."""
all_results = []
total = len(combos) * 2 # x2 for VAC on/off
idx = 0
for combo in combos:
for vac in [True, False]:
idx += 1
vac_str = "VAC=on" if vac else "VAC=off"
desc = f"{combo['backend']} / {combo['policy']}"
if combo["model"]:
desc += f" / {combo['model']}"
desc += f" / {vac_str}"
print(f"\n{'='*70}")
print(f"[{idx}/{total}] {desc}")
print(f"{'='*70}")
results = await run_single_combo(
combo, audio_files, vac=vac, lan="en", max_duration=max_duration,
)
all_results.extend(results)
# Free memory between combos
gc.collect()
return all_results
def main():
parser = argparse.ArgumentParser(description="Run comprehensive WhisperLiveKit benchmark")
parser.add_argument("--quick", action="store_true", help="Quick mode: fewer models and combos")
parser.add_argument("--json", default="benchmark_results.json", dest="json_output", help="Output JSON path")
parser.add_argument("--max-duration", type=float, default=60.0, help="Max audio duration in seconds")
args = parser.parse_args()
system_info = get_system_info()
combos = detect_combos(quick=args.quick)
audio_files = collect_audio_files()
print(f"System: {system_info.get('cpu', 'unknown')}, {system_info.get('ram_gb', '?')}GB RAM")
print(f"Backends: {list(system_info['backend_versions'].keys())}")
print(f"Combos to test: {len(combos)} x 2 (VAC on/off) = {len(combos)*2}")
print(f"Audio files: {[f.name for f in audio_files]}")
print()
t0 = time.time()
all_results = asyncio.run(
run_full_benchmark(combos, audio_files, max_duration=args.max_duration)
)
total_time = time.time() - t0
output = {
"system_info": system_info,
"benchmark_date": time.strftime("%Y-%m-%d %H:%M"),
"total_benchmark_time_s": round(total_time, 1),
"n_combos": len(combos) * 2,
"n_audio_files": len(audio_files),
"results": all_results,
}
Path(args.json_output).write_text(json.dumps(output, indent=2, ensure_ascii=False))
print(f"\nBenchmark complete in {total_time:.0f}s. Results: {args.json_output}")
if __name__ == "__main__":
main()

3346
scripts/alignment_heads_qwen3_asr_0.6B.json Normal file
View File

File diff suppressed because it is too large Load Diff

3445
scripts/alignment_heads_qwen3_asr_1.7B.json Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
scripts/alignment_heads_qwen3_asr_1.7B.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 83 KiB

3292
scripts/alignment_heads_qwen3_asr_1.7B_v2.json Normal file
View File

File diff suppressed because it is too large Load Diff

137
scripts/create_long_samples.py Normal file
View File

@@ -0,0 +1,137 @@
#!/usr/bin/env python3
"""Create long benchmark samples (5min+) by concatenating utterances from public datasets."""
import io
import json
import logging
import wave
from pathlib import Path
import numpy as np
logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)
CACHE = Path.home() / ".cache/whisperlivekit/benchmark_data"
CACHE.mkdir(parents=True, exist_ok=True)
SR = 16000
def save_wav(path, audio, sr=SR):
audio = np.clip(audio, -1, 1)
audio_int = (audio * 32767).astype(np.int16)
with wave.open(str(path), "w") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(sr)
wf.writeframes(audio_int.tobytes())
def decode_audio(audio_bytes):
import soundfile as sf
arr, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
return np.array(arr, dtype=np.float32), sr
def download_long_librispeech(config, lang_code, target_dur=300):
"""Concatenate LibriSpeech utterances into a ~5min sample."""
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info(f"Downloading LibriSpeech {config} for {lang_code} (~{target_dur}s)...")
ds = load_dataset("openslr/librispeech_asr", config, split="test", streaming=True)
ds = ds.cast_column("audio", Audio(decode=False))
chunks, texts = [], []
total = 0
for item in ds:
arr, sr = decode_audio(item["audio"]["bytes"])
chunks.append(arr)
texts.append(item["text"])
total += len(arr) / sr
if total >= target_dur:
break
if len(chunks) % 20 == 0:
logger.info(f" {total:.0f}s / {target_dur}s ({len(chunks)} utterances)")
# Insert small silences between utterances for natural transitions
silence = np.zeros(int(0.5 * sr), dtype=np.float32)
interleaved = []
for i, chunk in enumerate(chunks):
if i > 0:
interleaved.append(silence)
interleaved.append(chunk)
full = np.concatenate(interleaved)
total = len(full) / sr
ref = " ".join(texts)
name = f"{lang_code}_long_{config}"
path = CACHE / f"{name}.wav"
save_wav(path, full)
logger.info(f" -> {name}: {total:.1f}s ({len(texts)} utterances)")
return {"name": name, "path": str(path), "reference": ref,
"duration": round(total, 2), "language": lang_code.split("_")[0]}
def download_long_mls(config, lang_code, target_dur=300):
"""Concatenate MLS utterances into a ~5min sample."""
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info(f"Downloading MLS {config} for {lang_code} (~{target_dur}s)...")
ds = load_dataset("facebook/multilingual_librispeech", config, split="test", streaming=True)
ds = ds.cast_column("audio", Audio(decode=False))
chunks, texts = [], []
total = 0
for item in ds:
arr, sr = decode_audio(item["audio"]["bytes"])
chunks.append(arr)
texts.append(item.get("text", item.get("transcript", "")))
total += len(arr) / sr
if total >= target_dur:
break
if len(chunks) % 20 == 0:
logger.info(f" {total:.0f}s / {target_dur}s ({len(chunks)} utterances)")
silence = np.zeros(int(0.5 * sr), dtype=np.float32)
interleaved = []
for i, chunk in enumerate(chunks):
if i > 0:
interleaved.append(silence)
interleaved.append(chunk)
full = np.concatenate(interleaved)
total = len(full) / sr
ref = " ".join(texts)
name = f"{lang_code}_long"
path = CACHE / f"{name}.wav"
save_wav(path, full)
logger.info(f" -> {name}: {total:.1f}s ({len(texts)} utterances)")
return {"name": name, "path": str(path), "reference": ref,
"duration": round(total, 2), "language": lang_code}
def main():
samples = []
# English clean ~90s
samples.append(download_long_librispeech("clean", "en", target_dur=90))
# English noisy ~90s
samples.append(download_long_librispeech("other", "en_noisy", target_dur=90))
# French ~90s
samples.append(download_long_mls("french", "fr", target_dur=90))
# Save metadata
meta_path = CACHE / "long_samples.json"
meta_path.write_text(json.dumps(samples, indent=2))
logger.info(f"\nSaved metadata to {meta_path}")
total = sum(s["duration"] for s in samples)
logger.info(f"Total: {len(samples)} long samples, {total:.0f}s ({total/60:.1f}min)")
if __name__ == "__main__":
main()

703
scripts/detect_alignment_heads_qwen3.py Normal file
View File

@@ -0,0 +1,703 @@
#!/usr/bin/env python3
"""
Detect alignment heads in Qwen3-ASR for SimulStreaming-style inference.
Qwen3-ASR is a decoder-only multimodal model: audio is encoded by an audio
encoder and the resulting embeddings are injected into the text sequence
(replacing <|audio_pad|> placeholder tokens). The text decoder then attends
over the full sequence -- both audio-derived tokens and text tokens -- via
causal self-attention. There is **no** cross-attention.
For AlignAtt-style streaming, we need to find which (layer, head) pairs in
the text decoder's self-attention best track the monotonic alignment between
generated text tokens and their corresponding audio positions.
Algorithm
---------
For each audio sample with a known transcript:
1. Run Qwen3-ASR with output_attentions=True
2. Use the ForcedAligner to get ground-truth word->timestamp alignments
3. Convert timestamps to audio token positions in the input sequence
4. For each generated text token, check whether the argmax of each
attention head (over the audio-token region) points to the correct
audio position (as determined by the forced aligner)
5. Accumulate scores per (layer, head)
The heads whose attention argmax matches the ground-truth alignment most
often are the "alignment heads" usable for SimulStreaming.
Reference: Adapted from scripts/determine_alignment_heads.py (Whisper) and
iwslt26-sst/SimulMT_tests/heads/detect_translation_heads_qwen3.py
"""
import argparse
import io
import json
import logging
import re
import time
from difflib import SequenceMatcher
from typing import List, Optional, Tuple
import numpy as np
import soundfile as sf
import torch
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)
# ── Compatibility patches for qwen_asr 0.0.6 + transformers >= 5.3 ────
def _apply_transformers_compat_patches():
"""Apply all necessary patches to make qwen_asr work with transformers >= 5.3."""
# 1. check_model_inputs was removed
try:
import transformers.utils.generic as _g
if not hasattr(_g, "check_model_inputs"):
def check_model_inputs(*args, **kwargs):
def decorator(fn):
return fn
return decorator
_g.check_model_inputs = check_model_inputs
except ImportError:
pass
# 2. 'default' rope type was removed from ROPE_INIT_FUNCTIONS
try:
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
if "default" not in ROPE_INIT_FUNCTIONS:
def _compute_default_rope_parameters(config=None, device=None, seq_len=None, **kwargs):
if hasattr(config, "head_dim"):
head_dim = config.head_dim
else:
head_dim = config.hidden_size // config.num_attention_heads
partial = getattr(config, "partial_rotary_factor", 1.0)
dim = int(head_dim * partial)
base = config.rope_theta
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
return inv_freq, 1.0
ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters
except ImportError:
pass
# 3. pad_token_id missing on thinker config
try:
from qwen_asr.core.transformers_backend.configuration_qwen3_asr import (
Qwen3ASRThinkerConfig,
)
if not hasattr(Qwen3ASRThinkerConfig, "pad_token_id"):
Qwen3ASRThinkerConfig.pad_token_id = None
except ImportError:
pass
# 4. fix_mistral_regex is now handled internally by transformers 5.3;
# qwen_asr passes it explicitly, causing a duplicate-kwarg error.
try:
from transformers.models.auto import processing_auto
_orig_ap_from_pretrained = processing_auto.AutoProcessor.from_pretrained.__func__
@classmethod
def _patched_ap_from_pretrained(cls, *args, **kwargs):
kwargs.pop("fix_mistral_regex", None)
return _orig_ap_from_pretrained(cls, *args, **kwargs)
processing_auto.AutoProcessor.from_pretrained = _patched_ap_from_pretrained
except Exception:
pass
# 5. _finalize_model_loading calls initialize_weights which expects
# compute_default_rope_parameters on RotaryEmbedding modules.
try:
from qwen_asr.core.transformers_backend.modeling_qwen3_asr import (
Qwen3ASRThinkerTextRotaryEmbedding,
)
if not hasattr(Qwen3ASRThinkerTextRotaryEmbedding, "compute_default_rope_parameters"):
@staticmethod
def _compute_default_rope_parameters(config=None, device=None, seq_len=None, **kwargs):
if hasattr(config, "head_dim"):
head_dim = config.head_dim
else:
head_dim = config.hidden_size // config.num_attention_heads
partial = getattr(config, "partial_rotary_factor", 1.0)
dim = int(head_dim * partial)
base = config.rope_theta
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
return inv_freq, 1.0
Qwen3ASRThinkerTextRotaryEmbedding.compute_default_rope_parameters = _compute_default_rope_parameters
except ImportError:
pass
_apply_transformers_compat_patches()
# ── Constants ────────────────────────────────────────────────────────
SAMPLE_RATE = 16000
TS_THRESHOLD = 0.1 # Minimum Translation Score to qualify as alignment head
MIN_TEXT_SIMILARITY = 0.3 # Skip clips where generated text is too different from ground truth
def text_similarity(generated: str, reference: str) -> float:
"""Compute text similarity between generated and reference transcriptions.
Normalizes both strings (lowercase, remove punctuation, collapse whitespace)
then returns SequenceMatcher ratio.
"""
def normalize(s):
s = s.lower()
s = re.sub(r'[^\w\s]', '', s)
return re.sub(r'\s+', ' ', s).strip()
gen_norm = normalize(generated)
ref_norm = normalize(reference)
if not gen_norm or not ref_norm:
return 0.0
return SequenceMatcher(None, gen_norm, ref_norm).ratio()
def load_dataset_clips(name, config, split, limit):
"""Load audio clips from a HuggingFace dataset."""
from datasets import Audio as DatasetAudio
from datasets import load_dataset
ds = load_dataset(name, config, split=split)
ds = ds.cast_column("audio", DatasetAudio(decode=False))
clips = []
for idx, row in enumerate(ds):
if limit is not None and idx >= limit:
break
audio_field = row["audio"]
transcript = row["text"]
waveform_np, _ = sf.read(io.BytesIO(audio_field["bytes"]), dtype="float32")
if waveform_np.ndim > 1:
waveform_np = waveform_np.mean(axis=1)
clips.append((waveform_np, str(transcript)))
return clips
def get_device():
"""Select the best available device."""
if torch.backends.mps.is_available():
logger.info("Using MPS (Apple Silicon GPU)")
return torch.device("mps")
elif torch.cuda.is_available():
logger.info("Using CUDA (%s)", torch.cuda.get_device_name())
return torch.device("cuda")
else:
logger.info("Using CPU (will be slow)")
return torch.device("cpu")
def load_qwen3_asr(model_id: str, device: torch.device, dtype: torch.dtype):
"""Load Qwen3-ASR model, processor, and forced aligner."""
from qwen_asr.core.transformers_backend import (
Qwen3ASRConfig,
Qwen3ASRForConditionalGeneration,
Qwen3ASRProcessor,
)
from qwen_asr.inference.qwen3_forced_aligner import Qwen3ForcedAligner
from transformers import AutoConfig, AutoModel, AutoProcessor
AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
logger.info("Loading model: %s (dtype=%s, device=%s)", model_id, dtype, device)
model = AutoModel.from_pretrained(
model_id,
torch_dtype=dtype,
attn_implementation="eager",
device_map=str(device),
)
model.eval()
# Force eager attention on all sub-modules (attn_implementation="eager" doesn't
# propagate through nested model configs in qwen_asr's custom architecture)
for name, module in model.named_modules():
if hasattr(module, "config") and hasattr(module.config, "_attn_implementation"):
module.config._attn_implementation = "eager"
module.config._attn_implementation_internal = "eager"
try:
processor = AutoProcessor.from_pretrained(model_id, fix_mistral_regex=True)
except TypeError:
processor = AutoProcessor.from_pretrained(model_id)
logger.info("Loading forced aligner: Qwen/Qwen3-ForcedAligner-0.6B")
forced_aligner = Qwen3ForcedAligner.from_pretrained(
"Qwen/Qwen3-ForcedAligner-0.6B",
dtype=dtype,
device_map=str(device),
)
return model, processor, forced_aligner
def find_audio_token_range(input_ids: torch.Tensor, audio_token_id: int) -> Tuple[int, int]:
"""Find the start and end positions of audio tokens in the input sequence."""
mask = (input_ids == audio_token_id)
positions = mask.nonzero(as_tuple=True)[0]
if len(positions) == 0:
return 0, 0
return positions[0].item(), positions[-1].item() + 1
def timestamp_to_audio_token_position(
timestamp_sec: float,
audio_duration_sec: float,
audio_token_start: int,
audio_token_end: int,
) -> int:
"""Convert a timestamp in seconds to the corresponding audio token position.
Audio tokens span [audio_token_start, audio_token_end) in the input sequence.
We linearly interpolate within that range based on the timestamp fraction.
"""
n_audio_tokens = audio_token_end - audio_token_start
if n_audio_tokens <= 0 or audio_duration_sec <= 0:
return audio_token_start
fraction = min(timestamp_sec / audio_duration_sec, 1.0)
pos = audio_token_start + int(fraction * (n_audio_tokens - 1))
return max(audio_token_start, min(pos, audio_token_end - 1))
def run_detection(
model,
processor,
forced_aligner,
clips: List[Tuple[np.ndarray, str]],
language: Optional[str],
device: torch.device,
) -> Tuple[np.ndarray, int]:
"""Run alignment head detection on a set of audio clips.
Uses PyTorch forward hooks on each self_attn module to capture attention
weights that the decoder layer discards (``hidden_states, _ = self.self_attn(...)``).
With eager attention, ``self_attn`` always returns ``(attn_output, attn_weights)``
so the hook can read the weights from the return value.
Returns:
g: array of shape (total_heads,) with alignment hit counts
m: total number of alignment checks performed
"""
thinker = model.thinker
text_config = thinker.config.text_config
num_layers = text_config.num_hidden_layers
num_heads = text_config.num_attention_heads
total_heads = num_layers * num_heads
audio_token_id = thinker.config.audio_token_id
logger.info(
"Text decoder: %d layers x %d heads = %d total heads",
num_layers, num_heads, total_heads,
)
logger.info(
"KV heads: %d (GQA ratio: %d)",
text_config.num_key_value_heads,
num_heads // text_config.num_key_value_heads,
)
# Build prompt helper (same as Qwen3ASRModel._build_text_prompt)
from qwen_asr.inference.utils import normalize_language_name
def build_messages(audio_payload):
return [
{"role": "system", "content": ""},
{"role": "user", "content": [{"type": "audio", "audio": audio_payload}]},
]
def build_text_prompt(force_language=None):
msgs = build_messages("")
base = processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
if force_language:
base = base + f"language {force_language}<asr_text>"
return base
force_lang = None
if language:
force_lang = normalize_language_name(language)
# Stop token IDs
eos_ids = {151645, 151643} # <|im_end|>, <|endoftext|>
if processor.tokenizer.eos_token_id is not None:
eos_ids.add(processor.tokenizer.eos_token_id)
# Decoder layers: model.thinker.model.layers[i].self_attn
decoder_layers = thinker.model.layers
g = np.zeros(total_heads, dtype=np.int64)
m = 0
t0 = time.time()
for clip_idx, (waveform, transcript) in enumerate(clips):
if not transcript.strip():
continue
audio_duration = len(waveform) / SAMPLE_RATE
# 1. Get forced alignment timestamps
try:
align_results = forced_aligner.align(
audio=[(waveform, SAMPLE_RATE)],
text=[transcript],
language=[force_lang or "English"],
)
align_result = align_results[0]
except Exception as e:
logger.warning("Forced alignment failed for clip %d: %s", clip_idx, e)
continue
if not align_result.items:
continue
# Build word -> (start_time, end_time) mapping
word_timestamps = []
for item in align_result.items:
word_timestamps.append((item.text, item.start_time, item.end_time))
# 2. Prepare inputs
text_prompt = build_text_prompt(force_language=force_lang)
inputs = processor(
text=[text_prompt],
audio=[waveform],
return_tensors="pt",
padding=True,
)
inputs = inputs.to(model.device).to(model.dtype)
prompt_len = inputs.input_ids.shape[1]
# Find audio token range
audio_start, audio_end = find_audio_token_range(
inputs.input_ids[0], audio_token_id,
)
n_audio_tokens = audio_end - audio_start
if n_audio_tokens == 0:
logger.warning("No audio tokens found in clip %d", clip_idx)
continue
# 3. Register forward hooks on self_attn to capture attention weights.
# The decoder layer discards them: hidden_states, _ = self.self_attn(...)
# but eager_attention_forward always computes and returns attn_weights.
# We capture just the argmax over the audio region (memory-efficient).
# captured_argmax[layer_idx] = list of (num_heads,) tensors, one per decode step.
captured_argmax = {i: [] for i in range(num_layers)}
def _make_hook(store, a_start, a_end):
def hook_fn(module, args, output):
# output = (attn_output, attn_weights)
attn_weights = output[1]
if attn_weights is None:
return
# attn_weights shape: (batch, num_heads, q_len, kv_len)
# Only capture decode steps (q_len == 1), skip prefill
if attn_weights.shape[2] != 1:
return
kv_len = attn_weights.shape[-1]
if a_end > kv_len:
return
# Attention from the new token over audio region
audio_attn = attn_weights[0, :, 0, a_start:a_end] # (num_heads, n_audio)
store.append(audio_attn.argmax(dim=-1).cpu()) # (num_heads,)
return hook_fn
hooks = []
for layer_idx in range(num_layers):
h = decoder_layers[layer_idx].self_attn.register_forward_hook(
_make_hook(captured_argmax[layer_idx], audio_start, audio_end)
)
hooks.append(h)
# 4. Run generation
try:
with torch.inference_mode():
outputs = thinker.generate(
**inputs,
max_new_tokens=256,
do_sample=False,
)
except Exception as e:
for h in hooks:
h.remove()
logger.warning("Generation failed for clip %d: %s", clip_idx, e)
continue
finally:
for h in hooks:
h.remove()
# outputs is (batch, seq_len) tensor
all_generated = outputs[0, prompt_len:]
num_gen = len(all_generated)
for i, tid in enumerate(all_generated):
if tid.item() in eos_ids:
num_gen = i
break
generated_ids = all_generated[:num_gen]
if num_gen == 0:
del outputs, captured_argmax
continue
generated_text = processor.tokenizer.decode(generated_ids, skip_special_tokens=True)
# Filter out hallucinated clips (e.g. "!!!" patterns)
sim = text_similarity(generated_text, transcript)
if sim < MIN_TEXT_SIMILARITY:
logger.info(
"[%d/%d] SKIP (sim=%.2f) | %s...",
clip_idx + 1, len(clips), sim, generated_text[:60],
)
del outputs, captured_argmax
continue
# Verify hooks captured data
n_captured = len(captured_argmax[0])
if n_captured == 0:
logger.warning(
"No attention weights captured for clip %d (hooks may not have fired)", clip_idx
)
del outputs, captured_argmax
continue
# 5. Map generated tokens to word timestamps
gen_token_strings = [
processor.tokenizer.decode([tid.item()]) for tid in generated_ids
]
# Map each generated token index -> forced-aligner word index
accumulated_text = ""
word_idx = 0
token_to_word = {}
for tok_idx, tok_str in enumerate(gen_token_strings):
accumulated_text += tok_str
# Advance word index when accumulated text covers the current word
while (
word_idx < len(word_timestamps)
and len(accumulated_text.strip()) >= sum(
len(w[0]) + 1 for w in word_timestamps[:word_idx + 1]
)
):
word_idx += 1
actual_word_idx = min(word_idx, len(word_timestamps) - 1)
token_to_word[tok_idx] = actual_word_idx
# 6. Score each head using captured argmax data
for gen_step in range(num_gen):
word_idx = token_to_word.get(gen_step, None)
if word_idx is None or word_idx >= len(word_timestamps):
continue
_, word_start, word_end = word_timestamps[word_idx]
word_mid = (word_start + word_end) / 2.0
# Expected audio token position for this word
expected_pos = timestamp_to_audio_token_position(
word_mid, audio_duration, audio_start, audio_end,
)
# Tolerance: +/- a few audio tokens (proportional to word duration)
word_dur_tokens = max(1, int(
(word_end - word_start) / audio_duration * n_audio_tokens / 2
))
tolerance = max(3, word_dur_tokens)
m += 1
for layer_idx in range(num_layers):
if gen_step >= len(captured_argmax[layer_idx]):
continue
argmaxes = captured_argmax[layer_idx][gen_step].numpy() # (num_heads,)
for head_idx in range(num_heads):
attended_pos = argmaxes[head_idx] # relative to audio_start
attended_abs = audio_start + attended_pos
if abs(attended_abs - expected_pos) <= tolerance:
g[layer_idx * num_heads + head_idx] += 1
del outputs, captured_argmax
if device.type == "mps":
torch.mps.empty_cache()
elif device.type == "cuda":
torch.cuda.empty_cache()
elapsed = time.time() - t0
avg = elapsed / (clip_idx + 1)
eta = avg * (len(clips) - clip_idx - 1)
logger.info(
"[%d/%d] m=%d | %s... | %.1fs/clip | ETA: %.0fs",
clip_idx + 1, len(clips), m,
generated_text[:60], avg, eta,
)
return g, m
def main():
parser = argparse.ArgumentParser(
description="Detect alignment heads in Qwen3-ASR for SimulStreaming"
)
parser.add_argument(
"--model", type=str, default="Qwen/Qwen3-ASR-1.7B",
help="Qwen3-ASR model name or path",
)
parser.add_argument(
"--dataset", type=str, default="librispeech_asr",
help="HuggingFace dataset name",
)
parser.add_argument(
"--dataset-config", type=str, default="clean",
help="Dataset config/subset",
)
parser.add_argument(
"--dataset-split", type=str, default="validation",
help="Dataset split",
)
parser.add_argument(
"-n", "--num-samples", type=int, default=50,
help="Number of audio samples to process",
)
parser.add_argument(
"--language", type=str, default="English",
help="Language for forced alignment",
)
parser.add_argument(
"--dtype", type=str, default="bf16",
choices=["float32", "bf16", "float16"],
help="Model dtype",
)
parser.add_argument(
"-o", "--output", type=str, default="alignment_heads_qwen3_asr.json",
help="Output JSON file",
)
parser.add_argument(
"--heatmap", type=str, default="alignment_heads_qwen3_asr.png",
help="Output heatmap image",
)
parser.add_argument(
"--threshold", type=float, default=TS_THRESHOLD,
help="Minimum alignment score threshold",
)
args = parser.parse_args()
device = get_device()
dtype_map = {
"float32": torch.float32,
"bf16": torch.bfloat16,
"float16": torch.float16,
}
dtype = dtype_map[args.dtype]
# Load model
model, processor, forced_aligner = load_qwen3_asr(args.model, device, dtype)
# Load data
logger.info("Loading dataset: %s/%s [%s]", args.dataset, args.dataset_config, args.dataset_split)
clips = load_dataset_clips(
args.dataset, args.dataset_config, args.dataset_split, args.num_samples,
)
logger.info("Loaded %d clips", len(clips))
# Run detection
g, m = run_detection(model, processor, forced_aligner, clips, args.language, device)
# Compute alignment scores
thinker = model.thinker
text_config = thinker.config.text_config
num_layers = text_config.num_hidden_layers
num_heads = text_config.num_attention_heads
ts = g / max(m, 1)
ts_matrix = ts.reshape(num_layers, num_heads)
# Identify alignment heads
tah = []
for l in range(num_layers):
for h in range(num_heads):
score = ts_matrix[l, h]
if score > args.threshold:
tah.append({"layer": l, "head": h, "ts": round(float(score), 4)})
tah.sort(key=lambda x: x["ts"], reverse=True)
# Print results
print(f"\n{'=' * 60}")
print(f"ALIGNMENT HEADS (TS > {args.threshold}): {len(tah)} / {num_layers * num_heads}")
print(f"{'=' * 60}")
for entry in tah:
bar = "#" * int(entry["ts"] * 50)
print(f" L{entry['layer']:2d} H{entry['head']:2d} : TS={entry['ts']:.4f} {bar}")
n_active = sum(1 for s in ts if s > args.threshold)
n_low = sum(1 for s in ts if 0 < s <= args.threshold)
n_zero = sum(1 for s in ts if s == 0)
total_heads = num_layers * num_heads
print(f"\nDistribution:")
print(f" TS > {args.threshold} (alignment heads): {n_active} ({100 * n_active / total_heads:.1f}%)")
print(f" 0 < TS <= {args.threshold} (low activity): {n_low} ({100 * n_low / total_heads:.1f}%)")
print(f" TS = 0 (inactive): {n_zero} ({100 * n_zero / total_heads:.1f}%)")
print(f"\nTotal alignable tokens checked: m={m}")
# Save JSON
output = {
"model": args.model,
"language": args.language,
"num_layers": num_layers,
"num_heads": num_heads,
"num_kv_heads": text_config.num_key_value_heads,
"num_samples": len(clips),
"total_alignable_tokens": int(m),
"ts_threshold": args.threshold,
"ts_matrix": ts_matrix.tolist(),
"alignment_heads": tah,
# WhisperLiveKit-compatible format: list of [layer, head] pairs
"alignment_heads_compact": [[e["layer"], e["head"]] for e in tah],
}
with open(args.output, "w") as f:
json.dump(output, f, indent=2)
logger.info("Results saved to %s", args.output)
# Generate heatmap
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
fig, ax = plt.subplots(
figsize=(max(10, num_heads * 0.6), max(8, num_layers * 0.35)),
)
im = ax.imshow(
ts_matrix,
aspect="auto",
cmap="RdYlBu_r",
vmin=0,
vmax=max(0.4, ts_matrix.max()),
interpolation="nearest",
)
ax.set_xlabel("Head ID", fontsize=12)
ax.set_ylabel("Layer", fontsize=12)
ax.set_title(
f"Alignment Scores - {args.model}\n"
f"{len(tah)} alignment heads (TS > {args.threshold}), n={len(clips)}",
fontsize=13,
)
ax.set_xticks(range(num_heads))
ax.set_yticks(range(num_layers))
plt.colorbar(im, ax=ax, label="Alignment Score", shrink=0.8)
for entry in tah:
ax.add_patch(plt.Rectangle(
(entry["head"] - 0.5, entry["layer"] - 0.5),
1, 1, fill=False, edgecolor="red", linewidth=1.5,
))
plt.tight_layout()
plt.savefig(args.heatmap, dpi=150)
logger.info("Heatmap saved to %s", args.heatmap)
except Exception as e:
logger.warning("Could not generate heatmap: %s", e)
if __name__ == "__main__":
main()

216
scripts/generate_architecture.py Normal file
View File

@@ -0,0 +1,216 @@
#!/usr/bin/env python3
"""Generate the architecture.png diagram for WhisperLiveKit README."""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
# ── Colours ──
C_BG = "#1a1a2e"
C_PANEL = "#16213e"
C_PANEL2 = "#0f3460"
C_ACCENT = "#e94560"
C_GREEN = "#4ecca3"
C_ORANGE = "#f5a623"
C_BLUE = "#4a9eff"
C_PURPLE = "#b06af2"
C_PINK = "#ff6b9d"
C_YELLOW = "#f0e68c"
C_TEXT = "#e8e8e8"
C_TEXTDIM = "#a0a0b0"
C_BOX_BG = "#1e2d4a"
C_BOX_BG2 = "#2a1a3a"
C_BOX_BG3 = "#1a3a2a"
C_BORDER = "#3a4a6a"
fig, ax = plt.subplots(1, 1, figsize=(20, 12), facecolor=C_BG)
ax.set_xlim(0, 20)
ax.set_ylim(0, 12)
ax.set_aspect("equal")
ax.axis("off")
fig.subplots_adjust(left=0.01, right=0.99, top=0.97, bottom=0.01)
def box(x, y, w, h, label, color=C_BORDER, bg=C_BOX_BG, fontsize=8, bold=False,
text_color=C_TEXT, radius=0.15):
rect = FancyBboxPatch(
(x, y), w, h,
boxstyle=f"round,pad=0.05,rounding_size={radius}",
facecolor=bg, edgecolor=color, linewidth=1.2,
)
ax.add_patch(rect)
weight = "bold" if bold else "normal"
ax.text(x + w/2, y + h/2, label, ha="center", va="center",
fontsize=fontsize, color=text_color, fontweight=weight, family="monospace")
return rect
def arrow(x1, y1, x2, y2, color=C_TEXTDIM, style="->", lw=1.2):
ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
arrowprops=dict(arrowstyle=style, color=color, lw=lw))
def section_box(x, y, w, h, title, bg=C_PANEL, border=C_BORDER, title_color=C_ACCENT):
rect = FancyBboxPatch(
(x, y), w, h,
boxstyle="round,pad=0.05,rounding_size=0.2",
facecolor=bg, edgecolor=border, linewidth=1.5,
)
ax.add_patch(rect)
ax.text(x + 0.15, y + h - 0.25, title, ha="left", va="top",
fontsize=9, color=title_color, fontweight="bold", family="monospace")
# ═══════════════════════════════════════════════════════════════════
# Title
# ═══════════════════════════════════════════════════════════════════
ax.text(10, 11.7, "WhisperLiveKit Architecture", ha="center", va="center",
fontsize=16, color=C_TEXT, fontweight="bold", family="monospace")
ax.text(10, 11.35, "CLI commands: serve | listen | run | transcribe | bench | diagnose | models | pull | rm | check",
ha="center", va="center", fontsize=7, color=C_TEXTDIM, family="monospace")
# ═══════════════════════════════════════════════════════════════════
# Left: Client / Server
# ═══════════════════════════════════════════════════════════════════
section_box(0.1, 7.0, 3.5, 4.0, "FastAPI Server", border=C_GREEN)
box(0.3, 10.0, 1.5, 0.5, "Web UI\nHTML + JS", color=C_GREEN, fontsize=7)
box(2.0, 10.0, 1.4, 0.5, "Frontend\n(optional)", color=C_GREEN, fontsize=7)
box(0.3, 9.1, 3.1, 0.6, "WebSocket /asr • /v1/listen", color=C_GREEN, fontsize=7, bold=True)
box(0.3, 8.3, 3.1, 0.5, "REST /v1/audio/transcriptions", color=C_GREEN, fontsize=7)
box(0.3, 7.4, 3.1, 0.5, "Health • /v1/models", color=C_GREEN, fontsize=7)
# Clients
ax.text(0.2, 6.5, "Clients:", fontsize=7, color=C_TEXTDIM, family="monospace")
for i, client in enumerate(["Browser", "OpenAI SDK", "Deepgram SDK", "TestHarness"]):
box(0.3 + i * 0.9, 5.8, 0.8, 0.5, client, fontsize=5.5, bg="#1a2a1a", color="#3a6a3a")
# ═══════════════════════════════════════════════════════════════════
# Centre: Audio Processor (per-session pipeline)
# ═══════════════════════════════════════════════════════════════════
section_box(4.0, 5.5, 5.5, 5.5, "Audio Processor (per session)", border=C_BLUE)
box(4.3, 10.0, 2.0, 0.6, "FFmpeg\nDecoding", color=C_BLUE, bg="#1a2a4a", bold=True)
arrow(3.6, 9.4, 4.3, 10.2, color=C_GREEN)
box(6.6, 10.0, 2.6, 0.6, "Silero VAD\nspeech / silence", color=C_BLUE, bg="#1a2a4a")
arrow(6.3, 10.3, 6.6, 10.3, color=C_BLUE)
box(4.3, 8.8, 4.9, 0.8, "SessionASRProxy\nthread-safe per-session language override", color=C_BLUE, fontsize=7)
arrow(6.0, 10.0, 6.0, 9.6, color=C_BLUE)
box(4.3, 7.6, 2.3, 0.8, "DiffTracker\n(opt-in ?mode=diff)", color="#5a5a7a", fontsize=7)
box(6.9, 7.6, 2.3, 0.8, "Result Formatter\n→ FrontData.to_dict()", color=C_BLUE, fontsize=7)
# Streaming policies
ax.text(4.3, 7.1, "Streaming policies:", fontsize=7, color=C_ORANGE, fontweight="bold", family="monospace")
box(4.3, 6.2, 2.3, 0.7, "LocalAgreement\nHypothesisBuffer", color=C_ORANGE, bg="#2a2a1a", fontsize=7)
box(6.9, 6.2, 2.3, 0.7, "SimulStreaming\nAlignAtt (Whisper)", color=C_ORANGE, bg="#2a2a1a", fontsize=7)
# ═══════════════════════════════════════════════════════════════════
# Right: TranscriptionEngine (singleton)
# ═══════════════════════════════════════════════════════════════════
section_box(10.0, 0.3, 9.8, 10.7, "TranscriptionEngine (singleton — shared across sessions)",
border=C_ACCENT, bg="#1e1520")
ax.text(10.2, 10.5, "6 ASR Backends", fontsize=9, color=C_ACCENT, fontweight="bold", family="monospace")
# ── Whisper backends ──
section_box(10.2, 7.3, 4.5, 3.0, "Whisper Family (chunk-based)", border=C_PURPLE, bg=C_BOX_BG2)
box(10.4, 9.2, 1.3, 0.6, "Faster\nWhisper", color=C_PURPLE, bg="#2a1a3a", fontsize=7, bold=True)
box(11.9, 9.2, 1.3, 0.6, "MLX\nWhisper", color=C_PURPLE, bg="#2a1a3a", fontsize=7, bold=True)
box(13.4, 9.2, 1.1, 0.6, "OpenAI\nWhisper", color=C_PURPLE, bg="#2a1a3a", fontsize=7)
ax.text(10.4, 8.7, "PCM → Encoder → Decoder → Tokens", fontsize=6.5, color=C_TEXTDIM, family="monospace")
ax.text(10.4, 8.3, "Uses LocalAgreement or SimulStreaming (AlignAtt)", fontsize=6, color=C_PURPLE, family="monospace")
ax.text(10.4, 7.9, "Language detection • Buffer trimming", fontsize=6, color=C_TEXTDIM, family="monospace")
ax.text(10.4, 7.5, "CPU / CUDA / MLX", fontsize=6, color=C_TEXTDIM, family="monospace")
# ── Voxtral backends ──
section_box(10.2, 3.8, 4.5, 3.2, "Voxtral (native streaming)", border=C_PINK, bg="#2a1520")
box(10.4, 5.9, 1.8, 0.6, "Voxtral MLX\n(Apple Silicon)", color=C_PINK, bg="#2a1520", fontsize=7, bold=True)
box(12.5, 5.9, 2.0, 0.6, "Voxtral HF\n(CUDA/MPS/CPU)", color=C_PINK, bg="#2a1520", fontsize=7, bold=True)
ax.text(10.4, 5.4, "Incremental encoder → Autoregressive decoder", fontsize=6.5, color=C_TEXTDIM, family="monospace")
ax.text(10.4, 5.0, "Sliding KV cache • Token-by-token output", fontsize=6, color=C_PINK, family="monospace")
ax.text(10.4, 4.6, "No chunking needed — truly streams audio", fontsize=6, color=C_TEXTDIM, family="monospace")
ax.text(10.4, 4.2, "4B params • 15 languages • 6-bit quant (MLX)", fontsize=6, color=C_TEXTDIM, family="monospace")
# ── Qwen3 backend ──
section_box(15.0, 3.8, 4.6, 3.2, "Qwen3 ASR (batch + aligner)", border=C_GREEN, bg=C_BOX_BG3)
box(15.2, 5.9, 1.5, 0.6, "Qwen3 ASR\n1.7B / 0.6B", color=C_GREEN, bg="#1a3a2a", fontsize=7, bold=True)
box(16.9, 5.9, 1.5, 0.6, "Qwen3\nSimul", color=C_GREEN, bg="#1a3a2a", fontsize=7, bold=True)
box(18.6, 5.9, 1.0, 0.6, "Forced\nAligner", color=C_GREEN, bg="#1a3a2a", fontsize=6.5)
ax.text(15.2, 5.4, "Batch + SimulStreaming (AlignAtt)", fontsize=6.5, color=C_TEXTDIM, family="monospace")
ax.text(15.2, 5.0, "ForcedAligner provides word timestamps", fontsize=6, color=C_GREEN, family="monospace")
ax.text(15.2, 4.6, "LocalAgreement or border-distance policy", fontsize=6, color=C_TEXTDIM, family="monospace")
ax.text(15.2, 4.2, "29 languages • CUDA/MPS/CPU", fontsize=6, color=C_TEXTDIM, family="monospace")
# ── OpenAI API ──
box(15.2, 7.7, 4.2, 0.6, "OpenAI API (cloud)", color="#5a6a7a", fontsize=7)
ax.text(15.2, 7.4, "Remote transcription • API key required", fontsize=6, color=C_TEXTDIM, family="monospace")
# ── Shared components ──
section_box(10.2, 0.5, 9.4, 3.0, "Shared Components", border="#5a6a7a", bg="#151520")
box(10.4, 2.2, 2.5, 0.8, "Mel Spectrogram\ncached DFT + filterbank",
color="#5a6a7a", fontsize=7)
box(13.2, 2.2, 2.5, 0.8, "Diarization\nSortformer / pyannote",
color="#5a6a7a", fontsize=7)
box(16.0, 2.2, 3.4, 0.8, "Translation\nNLLB • CTranslate2",
color="#5a6a7a", fontsize=7)
box(10.4, 0.8, 4.0, 0.8, "WhisperLiveKitConfig\n(single source of truth)",
color=C_ACCENT, fontsize=7, bold=True)
box(14.8, 0.8, 2.3, 0.8, "TestHarness\npipeline testing",
color="#5a6a7a", fontsize=7)
box(17.3, 0.8, 2.3, 0.8, "Benchmark\n8 langs • 13 samples",
color=C_ORANGE, fontsize=7, bold=True)
# ═══════════════════════════════════════════════════════════════════
# Arrows: main data flow
# ═══════════════════════════════════════════════════════════════════
# Audio processor → TranscriptionEngine
arrow(9.5, 8.5, 10.2, 8.5, color=C_ACCENT, lw=2)
ax.text(9.6, 8.8, "PCM audio", fontsize=6, color=C_ACCENT, family="monospace")
# TranscriptionEngine → Audio processor (results)
arrow(10.2, 7.0, 9.5, 7.0, color=C_GREEN, lw=2)
ax.text(9.6, 7.3, "ASRTokens", fontsize=6, color=C_GREEN, family="monospace")
# Streaming policy connections
arrow(5.5, 6.2, 5.5, 5.5, color=C_ORANGE, style="->")
arrow(8.1, 6.2, 8.1, 5.5, color=C_ORANGE, style="->")
ax.text(4.3, 5.6, "Whisper + Qwen3", fontsize=5.5, color=C_ORANGE, family="monospace")
ax.text(6.9, 5.6, "Whisper + Qwen3-simul", fontsize=5.5, color=C_ORANGE, family="monospace")
# Voxtral note (no policy needed)
ax.text(10.2, 3.5, "Voxtral: own streaming processor (no external policy)", fontsize=6,
color=C_PINK, family="monospace", style="italic")
# ═══════════════════════════════════════════════════════════════════
# Legend
# ═══════════════════════════════════════════════════════════════════
legend_y = 5.0
ax.text(0.3, legend_y, "Streaming modes:", fontsize=7, color=C_TEXT, fontweight="bold", family="monospace")
for i, (label, color) in enumerate([
("Native streaming (Voxtral)", C_PINK),
("Chunk-based (Whisper)", C_PURPLE),
("Batch + aligner (Qwen3)", C_GREEN),
]):
ax.plot([0.3], [legend_y - 0.4 - i * 0.35], "s", color=color, markersize=6)
ax.text(0.6, legend_y - 0.4 - i * 0.35, label, fontsize=6.5, color=color,
va="center", family="monospace")
plt.savefig("architecture.png", dpi=200, facecolor=C_BG, bbox_inches="tight", pad_inches=0.1)
print("Saved architecture.png")

437
scripts/run_scatter_benchmark.py Normal file
View File

@@ -0,0 +1,437 @@
#!/usr/bin/env python3
"""Run benchmark across all backend x model x policy combos for scatter plot.
Tests each configuration on long audio samples in two modes:
- Compute-unaware (speed=0): all audio dumped instantly, measures pure model accuracy
- Compute-aware (speed=1.0): real-time simulation, slow models lose audio
Usage:
python scripts/run_scatter_benchmark.py
python scripts/run_scatter_benchmark.py --aware # only compute-aware
python scripts/run_scatter_benchmark.py --unaware # only compute-unaware
python scripts/run_scatter_benchmark.py --plot-only results.json
"""
import argparse
import asyncio
import gc
import json
import logging
import platform
import subprocess
import sys
import time
import warnings
warnings.filterwarnings("ignore")
logging.basicConfig(level=logging.WARNING)
for name in [
"whisperlivekit", "transformers", "torch", "httpx", "datasets",
"numexpr", "faster_whisper",
]:
logging.getLogger(name).setLevel(logging.ERROR)
LONG_SAMPLES_PATH = "~/.cache/whisperlivekit/benchmark_data/long_samples.json"
# ── All configurations to benchmark ──
COMBOS = [
# faster-whisper x LocalAgreement
{"backend": "faster-whisper", "model_size": "base", "policy": "localagreement",
"label": "fw LA base", "color": "#4a9eff", "marker": "o", "size": 100},
{"backend": "faster-whisper", "model_size": "small", "policy": "localagreement",
"label": "fw LA small", "color": "#4a9eff", "marker": "o", "size": 220},
# faster-whisper x SimulStreaming
{"backend": "faster-whisper", "model_size": "base", "policy": "simulstreaming",
"label": "fw SS base", "color": "#4a9eff", "marker": "s", "size": 100},
{"backend": "faster-whisper", "model_size": "small", "policy": "simulstreaming",
"label": "fw SS small", "color": "#4a9eff", "marker": "s", "size": 220},
# mlx-whisper x LocalAgreement
{"backend": "mlx-whisper", "model_size": "base", "policy": "localagreement",
"label": "mlx LA base", "color": "#4ecca3", "marker": "o", "size": 100},
{"backend": "mlx-whisper", "model_size": "small", "policy": "localagreement",
"label": "mlx LA small", "color": "#4ecca3", "marker": "o", "size": 220},
# mlx-whisper x SimulStreaming
{"backend": "mlx-whisper", "model_size": "base", "policy": "simulstreaming",
"label": "mlx SS base", "color": "#4ecca3", "marker": "s", "size": 100},
{"backend": "mlx-whisper", "model_size": "small", "policy": "simulstreaming",
"label": "mlx SS small", "color": "#4ecca3", "marker": "s", "size": 220},
# voxtral-mlx (4B, native streaming)
{"backend": "voxtral-mlx", "model_size": "", "policy": "",
"label": "voxtral mlx", "color": "#f5a623", "marker": "D", "size": 250},
]
def is_backend_available(backend):
try:
if backend == "faster-whisper":
import faster_whisper; return True # noqa
elif backend == "mlx-whisper":
import mlx_whisper; return True # noqa
elif backend == "whisper":
import whisper; return True # noqa
elif backend == "voxtral-mlx":
import mlx.core # noqa
from whisperlivekit.voxtral_mlx.loader import load_voxtral_model; return True # noqa
elif backend == "voxtral":
from transformers import VoxtralRealtimeForConditionalGeneration; return True # noqa
elif backend in ("qwen3", "qwen3-simul"):
from whisperlivekit.qwen3_asr import _patch_transformers_compat
_patch_transformers_compat()
from qwen_asr import Qwen3ASRModel; return True # noqa
except (ImportError, Exception):
pass
return False
def get_system_info():
info = {"platform": platform.platform(), "machine": platform.machine()}
try:
info["cpu"] = subprocess.check_output(
["sysctl", "-n", "machdep.cpu.brand_string"], text=True).strip()
except Exception:
info["cpu"] = platform.processor()
try:
mem = int(subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True).strip())
info["ram_gb"] = round(mem / (1024**3))
except Exception:
info["ram_gb"] = None
return info
async def run_combo_on_samples(combo, samples, lang="en", speed=0):
"""Run one config on all samples, return averaged result.
Args:
speed: 0 = compute-unaware (instant dump), 1.0 = compute-aware (real-time)
"""
from whisperlivekit.core import TranscriptionEngine
from whisperlivekit.metrics import compute_wer
from whisperlivekit.test_harness import TestHarness, _engine_cache
kwargs = {"lan": lang, "pcm_input": True}
if combo["backend"]:
kwargs["backend"] = combo["backend"]
if combo["model_size"]:
kwargs["model_size"] = combo["model_size"]
if combo.get("policy"):
kwargs["backend_policy"] = combo["policy"]
TranscriptionEngine.reset()
_engine_cache.clear()
gc.collect()
total_ref_words, total_errors = 0, 0
total_infer_time, total_audio_time = 0.0, 0.0
n_ok = 0
for sample in samples:
try:
async with TestHarness(**kwargs) as h:
await h.feed(sample["path"], speed=speed)
await h.drain(max(5.0, sample["duration"] * 0.5))
state = await h.finish(timeout=120)
metrics = h.metrics
hypothesis = state.committed_text or state.text
wer_result = compute_wer(sample["reference"], hypothesis)
total_ref_words += wer_result["ref_words"]
total_errors += (wer_result["substitutions"] +
wer_result["insertions"] +
wer_result["deletions"])
# Use actual inference time from metrics, not wall clock
if metrics and metrics.transcription_durations:
total_infer_time += sum(metrics.transcription_durations)
total_audio_time += sample["duration"]
n_ok += 1
except Exception as e:
print(f" [WARN: {sample['name']} failed: {e}]", end="")
if n_ok == 0:
return None
weighted_wer = total_errors / max(total_ref_words, 1)
# Real RTF = actual inference time / audio duration
real_rtf = total_infer_time / total_audio_time if total_audio_time > 0 else 0
return {
"label": combo["label"],
"backend": combo["backend"],
"model_size": combo.get("model_size", ""),
"policy": combo.get("policy", ""),
"color": combo["color"],
"marker": combo["marker"],
"size": combo["size"],
"rtf": round(real_rtf, 4),
"wer_pct": round(weighted_wer * 100, 1),
"n_samples": n_ok,
}
async def run_all(combos, samples, lang="en", speed=0):
mode_label = "compute-aware" if speed > 0 else "compute-unaware"
results = []
for i, combo in enumerate(combos):
if not is_backend_available(combo["backend"]):
print(f" [{i+1}/{len(combos)}] SKIP {combo['label']} (not installed)")
continue
print(f" [{i+1}/{len(combos)}] {combo['label']} ({mode_label})...", end="", flush=True)
result = await run_combo_on_samples(combo, samples, lang, speed=speed)
if result:
results.append(result)
print(f" RTF={result['rtf']:.2f}x WER={result['wer_pct']:.1f}% ({result['n_samples']} samples)")
else:
print(" FAILED (no results)")
return results
def get_long_samples_for_lang(lang="en"):
"""Load long benchmark samples from long_samples.json, filtered by language."""
import os
path = os.path.expanduser(LONG_SAMPLES_PATH)
if not os.path.exists(path):
print(f"ERROR: Long samples file not found: {path}")
print("Please generate it first (see benchmark_data/README).")
sys.exit(1)
with open(path) as f:
all_samples = json.load(f)
samples = [s for s in all_samples if s["language"] == lang]
return [{"name": s["name"], "path": s["path"], "reference": s["reference"],
"duration": s["duration"]} for s in samples]
LANG_NAMES = {
"en": "English", "fr": "French", "es": "Spanish", "de": "German",
"pt": "Portuguese", "it": "Italian", "nl": "Dutch", "pl": "Polish",
"zh": "Chinese", "ja": "Japanese", "ko": "Korean", "ru": "Russian",
}
def generate_scatter(results, system_info, output_path, n_samples, lang="en",
mode="unaware", sample_duration=0.0):
"""Generate scatter plot.
Args:
mode: "unaware" or "aware" -- shown in title
sample_duration: total audio duration in seconds -- shown in title
"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.lines import Line2D
fig, ax = plt.subplots(figsize=(12, 7), facecolor="white")
ax.set_facecolor("#fafafa")
# Show ALL points on chart (no outlier exclusion)
main = results
slow = []
# Axis limits: fit all data
if main:
xmax = max(r["rtf"] for r in main) * 1.15
ymax = max(r["wer_pct"] for r in main) * 1.15 + 1
else:
xmax, ymax = 0.5, 10
xmax = max(xmax, 1.15) # always show the real-time line
ymax = max(ymax, 8)
# Sweet spot zone: RTF < 1.0 (real-time) and WER < 12%
sweet_x = min(1.0, xmax * 0.85)
sweet_y = min(12, ymax * 0.45)
rect = plt.Rectangle((0, 0), sweet_x, sweet_y, alpha=0.07, color="#4ecca3",
zorder=0, linewidth=0)
ax.add_patch(rect)
ax.text(sweet_x - 0.005, sweet_y - 0.15, "sweet spot", ha="right", va="top",
fontsize=10, color="#2ecc71", fontstyle="italic", fontweight="bold", alpha=0.5)
# Real-time limit line
ax.axvline(x=1.0, color="#e94560", linestyle="--", linewidth=1.5, alpha=0.4, zorder=1)
ax.text(1.02, ymax * 0.97, "real-time\nlimit", fontsize=8, color="#e94560",
va="top", alpha=0.6)
# Manual label offsets keyed by label name — hand-tuned
OFFSETS = {
"fw LA base": (8, 8),
"fw LA small": (8, 8),
"fw SS base": (-55, -14),
"fw SS small": (8, 8),
"mlx LA base": (8, 10),
"mlx LA small": (8, 8),
"mlx SS base": (-55, 8),
"mlx SS small": (-55, -5),
"voxtral mlx": (10, -14),
"qwen3 0.6B": (10, 8),
"qwen3-mlx 0.6B": (10, -14),
"qwen3-mlx 1.7B": (10, 8),
"fw LA large-v3": (8, -5),
"fw SS large-v3": (8, 5),
}
# Plot main points
for r in main:
ax.scatter(r["rtf"], r["wer_pct"], c=r["color"], marker=r["marker"],
s=r["size"], edgecolors="white", linewidths=1.0, zorder=5, alpha=0.85)
ox, oy = OFFSETS.get(r["label"], (8, -4))
ax.annotate(r["label"], (r["rtf"], r["wer_pct"]),
textcoords="offset points", xytext=(ox, oy),
fontsize=8.5, color="#333333", fontweight="medium")
# Note slow backends outside main view
if slow:
lines = []
for r in slow:
lines.append(f"{r['label']}: RTF={r['rtf']:.1f}x, WER={r['wer_pct']:.1f}%")
note = "Beyond real-time:\n" + "\n".join(lines)
ax.text(xmax * 0.97, ymax * 0.97, note, ha="right", va="top",
fontsize=7.5, color="#777777", fontstyle="italic",
bbox=dict(boxstyle="round,pad=0.4", facecolor="#f8f8f8",
edgecolor="#dddddd", alpha=0.9))
# Axes
ax.set_xlim(left=-0.01, right=xmax)
ax.set_ylim(bottom=0, top=ymax)
ax.set_xlabel("RTF (lower = faster)", fontsize=13, fontweight="bold", labelpad=8)
ax.set_ylabel("WER % (lower = more accurate)", fontsize=13, fontweight="bold", labelpad=8)
ax.grid(True, alpha=0.15, linestyle="-", color="#cccccc")
ax.tick_params(labelsize=10)
# Title
cpu = system_info.get("cpu", "unknown").replace("Apple ", "")
lang_name = LANG_NAMES.get(lang, lang.upper())
mode_label = "compute-unaware" if mode == "unaware" else "compute-aware"
dur_str = f"{sample_duration / 60:.0f}min" if sample_duration >= 60 else f"{sample_duration:.0f}s"
ax.set_title(
f"Speed vs Accuracy ({mode_label}) — {n_samples} {lang_name} samples, {dur_str} ({cpu})",
fontsize=14, fontweight="bold", pad=12)
# Legend — backends
backend_handles = []
seen = set()
for r in results:
if r["backend"] not in seen:
seen.add(r["backend"])
backend_handles.append(mpatches.Patch(color=r["color"], label=r["backend"]))
# Legend — shapes
marker_map = {"o": "LocalAgreement", "s": "SimulStreaming", "D": "Native streaming",
"h": "Batch + aligner"}
active = set(r["marker"] for r in results)
shape_handles = [
Line2D([0], [0], marker=m, color="#888", label=lbl,
markerfacecolor="#888", markersize=8, linestyle="None")
for m, lbl in marker_map.items() if m in active
]
# sizes
shape_handles += [
Line2D([0], [0], marker="o", color="#888", label="base",
markerfacecolor="#888", markersize=5, linestyle="None"),
Line2D([0], [0], marker="o", color="#888", label="small / 4B",
markerfacecolor="#888", markersize=9, linestyle="None"),
]
leg1 = ax.legend(handles=backend_handles, loc="upper left", fontsize=9,
framealpha=0.95, edgecolor="#ddd", title="Backend", title_fontsize=9)
ax.add_artist(leg1)
ax.legend(handles=shape_handles, loc="lower right", fontsize=8,
framealpha=0.95, edgecolor="#ddd", ncol=2)
plt.tight_layout()
plt.savefig(output_path, dpi=150, bbox_inches="tight", pad_inches=0.15)
print(f"Saved {output_path}")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--plot-only", default=None)
parser.add_argument("--lang", default="en", help="Language code (en, fr, es, de, ...)")
parser.add_argument("--output", "-o", default=None,
help="Output path prefix (mode suffix added automatically)")
parser.add_argument("--json-output", default=None,
help="JSON output path prefix (mode suffix added automatically)")
parser.add_argument("--aware", action="store_true",
help="Run only compute-aware mode (speed=1.0)")
parser.add_argument("--unaware", action="store_true",
help="Run only compute-unaware mode (speed=0)")
args = parser.parse_args()
lang = args.lang
# Determine which modes to run
if args.aware and args.unaware:
modes = ["unaware", "aware"]
elif args.aware:
modes = ["aware"]
elif args.unaware:
modes = ["unaware"]
else:
# Default: run both
modes = ["unaware", "aware"]
if args.plot_only:
data = json.load(open(args.plot_only))
mode = data.get("mode", "unaware")
output_path = args.output or f"benchmark_scatter_{lang}_{mode}.png"
generate_scatter(data["results"], data["system_info"], output_path,
data["n_samples"], data.get("lang", "en"),
mode=mode,
sample_duration=data.get("total_audio_s", 0))
return
print(f"Loading long {lang} samples from {LONG_SAMPLES_PATH}...")
samples = get_long_samples_for_lang(lang)
if not samples:
print(f"ERROR: No long samples for language '{lang}'")
sys.exit(1)
print(f"Using {len(samples)} samples: {[s['name'] for s in samples]}")
total_dur = sum(s["duration"] for s in samples)
print(f"Total audio: {total_dur:.0f}s ({total_dur / 60:.1f}min)\n")
# Filter combos to backends that support this language
from whisperlivekit.benchmark.compat import backend_supports_language
combos = [c for c in COMBOS if backend_supports_language(c["backend"], lang)]
system_info = get_system_info()
for mode in modes:
speed = 1.0 if mode == "aware" else 0
mode_label = "compute-aware" if mode == "aware" else "compute-unaware"
print(f"\n{'='*60}")
print(f" Running {mode_label} (speed={speed})")
print(f"{'='*60}\n")
t0 = time.time()
results = asyncio.run(run_all(combos, samples, lang, speed=speed))
total = time.time() - t0
# Save JSON
json_path = args.json_output or f"/tmp/bench_scatter_{lang}"
json_file = f"{json_path}_{mode}.json"
output_data = {
"system_info": system_info,
"lang": lang,
"mode": mode,
"speed": speed,
"n_samples": len(samples),
"sample_names": [s["name"] for s in samples],
"total_audio_s": round(total_dur, 1),
"total_benchmark_time_s": round(total, 1),
"results": results,
}
with open(json_file, "w") as f:
json.dump(output_data, f, indent=2)
print(f"\nJSON: {json_file} ({total:.0f}s total)")
# Generate scatter plot
output_base = args.output or f"benchmark_scatter_{lang}"
output_path = f"{output_base}_{mode}.png"
generate_scatter(results, system_info, output_path, len(samples), lang,
mode=mode, sample_duration=total_dur)
if __name__ == "__main__":
main()

804
test_backend_offline.py
View File

@@ -1,804 +0,0 @@
#!/usr/bin/env python3
"""
Offline test harness and benchmark suite for WhisperLiveKit backends.
Simulates a client-server session by feeding audio files as PCM bytes through
the full AudioProcessor pipeline (the same path used by the WebSocket server),
without needing a browser or microphone.
Computes WER (Word Error Rate) and timestamp accuracy when ground truth
transcript files (.transcript.json) are available alongside audio files.
Usage:
# Test with a single audio file:
python test_backend_offline.py --backend faster-whisper --audio audio_tests/00_00_07_english_1_speaker.wav
# Test all files in audio_tests/:
python test_backend_offline.py --backend faster-whisper --no-realtime
# Override streaming policy:
python test_backend_offline.py --backend faster-whisper --policy simulstreaming --no-realtime
# Multi-backend benchmark (auto-detects all installed backends):
python test_backend_offline.py --benchmark --no-realtime
# Export results as JSON:
python test_backend_offline.py --benchmark --no-realtime --json results.json
# Insert silence for testing silence handling:
python test_backend_offline.py --backend faster-whisper --insert-silence 3.0 2.0
"""
import argparse
import asyncio
import json
import logging
import sys
import time
import urllib.request
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import List, Optional
import numpy as np
logging.basicConfig(
level=logging.WARNING,
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger("test_offline")
logger.setLevel(logging.INFO)
SAMPLE_RATE = 16000
JFK_WAV_URL = "https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav"
CACHE_DIR = Path(__file__).parent / ".test_cache"
AUDIO_TESTS_DIR = Path(__file__).parent / "audio_tests"
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
@dataclass
class WordTimestamp:
"""Word with its start/end time."""
word: str
start: float
end: float
@dataclass
class TestResult:
"""Structured result from a single test run."""
audio_file: str
audio_duration_s: float
backend: str
policy: str
language: str
chunk_ms: int
realtime_pacing: bool
# Timing
processing_time_s: float
rtf: float # real-time factor
# Transcription output
transcription: str
n_lines: int
n_responses: int
# WER metrics (None if no ground truth)
wer: Optional[float] = None
wer_details: Optional[dict] = None
# Timestamp accuracy (None if no ground truth)
timestamp_mae: Optional[float] = None
timestamp_max_delta: Optional[float] = None
timestamp_median_delta: Optional[float] = None
# Word-level timestamps
word_timestamps: List[WordTimestamp] = field(default_factory=list)
# Raw last response
last_response: Optional[dict] = None
def download_sample_audio() -> Path:
"""Download the jfk.wav sample if not cached."""
CACHE_DIR.mkdir(exist_ok=True)
path = CACHE_DIR / "jfk.wav"
if not path.exists():
logger.info(f"Downloading sample audio to {path} ...")
urllib.request.urlretrieve(JFK_WAV_URL, path)
logger.info("Done.")
return path
def load_audio(path: str) -> np.ndarray:
"""Load audio file as float32 mono 16kHz numpy array.
Supports WAV, FLAC (via soundfile) and MP3, OGG, M4A (via librosa).
"""
ext = Path(path).suffix.lower()
if ext in (".mp3", ".ogg", ".m4a"):
import librosa
audio, _ = librosa.load(path, sr=SAMPLE_RATE, mono=True)
return audio.astype(np.float32)
import soundfile as sf
audio, sr = sf.read(path, dtype="float32")
if audio.ndim > 1:
audio = audio.mean(axis=1)
if sr != SAMPLE_RATE:
import librosa
audio = librosa.resample(audio, orig_sr=sr, target_sr=SAMPLE_RATE)
return audio
def insert_silence(audio: np.ndarray, silence_sec: float, position_sec: float) -> np.ndarray:
"""Insert silence into audio at a given position.
Args:
audio: Float32 mono audio array at SAMPLE_RATE.
silence_sec: Duration of silence to insert in seconds.
position_sec: Position in seconds where silence starts.
Returns:
New audio array with silence inserted.
"""
pos_samples = int(position_sec * SAMPLE_RATE)
silence_samples = int(silence_sec * SAMPLE_RATE)
pos_samples = min(pos_samples, len(audio))
silence = np.zeros(silence_samples, dtype=np.float32)
return np.concatenate([audio[:pos_samples], silence, audio[pos_samples:]])
def float32_to_s16le_bytes(audio: np.ndarray) -> bytes:
"""Convert float32 audio to s16le PCM bytes (what the browser sends)."""
return (audio * 32768).clip(-32768, 32767).astype(np.int16).tobytes()
def create_engine(
backend: str, model_size: str, lan: str,
diarization: bool = False,
diarization_backend: str = "",
vac: bool = True,
policy: str = "",
):
"""Create a TranscriptionEngine with the given backend config."""
import gc
from whisperlivekit.core import TranscriptionEngine
# Reset singleton so we get a fresh instance
TranscriptionEngine._instance = None
TranscriptionEngine._initialized = False
gc.collect()
kwargs = dict(
backend=backend,
lan=lan,
pcm_input=True,
vac=vac,
transcription=True,
diarization=diarization,
)
if diarization_backend:
kwargs["diarization_backend"] = diarization_backend
if model_size:
kwargs["model_size"] = model_size
if policy:
kwargs["backend_policy"] = policy
return TranscriptionEngine(**kwargs)
def _extract_text_from_response(response_dict: dict) -> str:
"""Extract full transcription text from a FrontData dict."""
def _strip_or_empty(value: object) -> str:
return value.strip() if isinstance(value, str) else ""
segments = response_dict.get("lines", [])
full_text = " ".join(
text
for seg in segments
if isinstance(seg, dict)
for text in [_strip_or_empty(seg.get("text"))]
if text
)
buf = _strip_or_empty(response_dict.get("buffer_transcription"))
if buf:
full_text = f"{full_text} {buf}".strip() if full_text else buf
return full_text
async def run_test(
engine, audio: np.ndarray, chunk_ms: int, realtime: bool,
audio_file: str = "", backend: str = "", policy: str = "", lan: str = "",
) -> TestResult:
"""
Simulate a client session through the full AudioProcessor pipeline.
1. Create AudioProcessor (one per "client session")
2. Start async pipeline (transcription_processor, results_formatter, etc.)
3. Feed audio as PCM bytes in timed chunks
4. Collect and display FrontData responses
5. Signal EOF and cleanup
"""
from whisperlivekit.audio_processor import AudioProcessor
chunk_samples = int(SAMPLE_RATE * chunk_ms / 1000)
total_samples = len(audio)
audio_duration = total_samples / SAMPLE_RATE
logger.info(
f"Audio: {audio_duration:.2f}s | "
f"Chunk: {chunk_ms}ms ({chunk_samples} samples) | "
f"Steps: {total_samples // chunk_samples + 1} | "
f"Realtime: {realtime}"
)
# --- Server side: create processor and start pipeline ---
processor = AudioProcessor(transcription_engine=engine)
results_generator = await processor.create_tasks()
# Collect results in background (like handle_websocket_results)
all_responses = []
response_count = 0
last_printed_text = ""
async def collect_results():
nonlocal response_count, last_printed_text
async for response in results_generator:
all_responses.append(response)
response_count += 1
d = response.to_dict()
# Only print when transcription text actually changes
current_text = _extract_text_from_response(d)
if current_text and current_text != last_printed_text:
buf = d.get("buffer_transcription")
buf = buf.strip() if isinstance(buf, str) else ""
committed = current_text
if buf and committed.endswith(buf):
committed = committed[:-len(buf)].strip()
# Show committed text + buffer separately
display = committed
if buf:
display = f"{committed} \033[90m{buf}\033[0m" if committed else f"\033[90m{buf}\033[0m"
print(f" > {display}", flush=True)
last_printed_text = current_text
result_task = asyncio.create_task(collect_results())
# --- Client side: feed audio as PCM bytes ---
t_start = time.time()
for offset in range(0, total_samples, chunk_samples):
chunk = audio[offset : offset + chunk_samples]
pcm_bytes = float32_to_s16le_bytes(chunk)
await processor.process_audio(pcm_bytes)
if realtime:
await asyncio.sleep(chunk_ms / 1000)
feed_elapsed = time.time() - t_start
logger.info(f"Audio fed in {feed_elapsed:.2f}s. Signaling EOF...")
# Signal end of audio (like client disconnect / empty message)
await processor.process_audio(None)
# Wait for pipeline to drain completely
try:
await asyncio.wait_for(result_task, timeout=120.0)
except asyncio.TimeoutError:
logger.warning("Timed out waiting for results. Proceeding with cleanup.")
result_task.cancel()
try:
await result_task
except asyncio.CancelledError:
pass
# --- Capture word-level timestamps before cleanup ---
word_timestamps = []
try:
state = await processor.get_current_state()
for token in state.tokens:
if hasattr(token, 'start') and hasattr(token, 'text') and token.text:
word_timestamps.append(WordTimestamp(
word=token.text.strip(),
start=round(token.start, 3),
end=round(token.end, 3),
))
except Exception as e:
logger.warning(f"Could not capture word timestamps: {e}")
# Cleanup
await processor.cleanup()
total_elapsed = time.time() - t_start
# --- Build result ---
transcription = ""
n_lines = 0
last_response_dict = None
if all_responses:
last = all_responses[-1].to_dict()
last_response_dict = last
n_lines = len(last.get("lines", []))
transcription = _extract_text_from_response(last)
# --- Compute WER and timestamp accuracy against ground truth ---
from whisperlivekit.metrics import compute_timestamp_accuracy, compute_wer
wer_val = None
wer_details = None
ts_mae = None
ts_max_delta = None
ts_median_delta = None
gt_path = Path(audio_file).with_suffix(".transcript.json")
if not gt_path.exists():
gt_path = AUDIO_TESTS_DIR / gt_path
gt = None
if gt_path.exists():
with open(gt_path) as f:
gt = json.load(f)
# WER
gt_text = " ".join(w["word"] for w in gt)
wer_result = compute_wer(gt_text, transcription)
wer_val = round(wer_result["wer"], 4)
wer_details = wer_result
# Timestamp accuracy
if word_timestamps:
pred_dicts = [{"word": wt.word, "start": wt.start, "end": wt.end} for wt in word_timestamps]
ts_result = compute_timestamp_accuracy(pred_dicts, gt)
ts_mae = ts_result["mae_start"]
ts_max_delta = ts_result["max_delta_start"]
ts_median_delta = ts_result["median_delta_start"]
result = TestResult(
audio_file=audio_file,
audio_duration_s=round(audio_duration, 2),
backend=backend,
policy=policy,
language=lan,
chunk_ms=chunk_ms,
realtime_pacing=realtime,
processing_time_s=round(total_elapsed, 2),
rtf=round(total_elapsed / audio_duration, 2),
transcription=transcription,
n_lines=n_lines,
n_responses=response_count,
wer=wer_val,
wer_details=wer_details,
timestamp_mae=round(ts_mae, 3) if ts_mae is not None else None,
timestamp_max_delta=round(ts_max_delta, 3) if ts_max_delta is not None else None,
timestamp_median_delta=round(ts_median_delta, 3) if ts_median_delta is not None else None,
word_timestamps=word_timestamps,
last_response=last_response_dict,
)
# --- Print summary ---
print(f"\n{'=' * 60}")
print(f"RESULT: {audio_file}")
print(f"{'=' * 60}")
print(f"Transcription: {transcription}")
print(f"Lines: {n_lines} | Responses: {response_count}")
print(f"Audio: {audio_duration:.2f}s | Time: {total_elapsed:.2f}s | RTF: {result.rtf:.2f}x")
if wer_val is not None:
print(f"WER: {wer_val:.2%} (S={wer_details['substitutions']} I={wer_details['insertions']} D={wer_details['deletions']})")
# Print word timestamps if available
if word_timestamps:
print(f"\nWord timestamps ({len(word_timestamps)} words):")
for wt in word_timestamps:
print(f" [{wt.start:6.2f} - {wt.end:6.2f}] {wt.word}")
# Detailed comparison with ground truth
if gt:
print(f"\n vs Ground truth ({len(gt)} words):")
max_words = max(len(word_timestamps), len(gt))
for i in range(max_words):
pred = word_timestamps[i] if i < len(word_timestamps) else None
ref = gt[i] if i < len(gt) else None
p_str = f"[{pred.start:5.2f}-{pred.end:5.2f}] {pred.word:<15}" if pred else " " * 30
r_str = f"[{ref['start']:5.2f}-{ref['end']:5.2f}] {ref['word']:<15}" if ref else ""
delta = ""
if pred and ref:
d = pred.start - ref['start']
delta = f" Δstart={d:+.2f}"
print(f" {p_str} | {r_str}{delta}")
if ts_mae is not None:
print(f"\n Timestamp stats: MAE={ts_mae:.3f}s max|Δ|={ts_max_delta:.3f}s median|Δ|={ts_median_delta:.3f}s")
print(f"{'=' * 60}")
return result
def discover_audio_files(directory: str) -> List[Path]:
"""Find all supported audio files in directory."""
d = Path(directory)
files = sorted(
p for p in d.iterdir()
if p.is_file() and p.suffix.lower() in AUDIO_EXTENSIONS
)
return files
async def run_all_tests(
engine, audio_files: List[Path], chunk_ms: int, realtime: bool,
backend: str, policy: str, lan: str, max_duration: float = 60.0,
silence_insertions: Optional[List[List[float]]] = None,
) -> List[TestResult]:
"""Run tests on multiple audio files sequentially."""
results = []
for audio_path in audio_files:
# Detect language from filename if "french" in name
file_lan = lan
if "french" in audio_path.name.lower() and lan == "en":
file_lan = "fr"
logger.info("Auto-detected language 'fr' from filename")
audio = load_audio(str(audio_path))
# Insert silence segments (applied in reverse position order to keep offsets valid)
if silence_insertions:
for secs, at_sec in sorted(silence_insertions, key=lambda x: x[1], reverse=True):
logger.info(f"Inserting {secs:.1f}s silence at {at_sec:.1f}s")
audio = insert_silence(audio, secs, at_sec)
duration = len(audio) / SAMPLE_RATE
if duration > max_duration:
logger.info(f"Skipping {audio_path.name} ({duration:.0f}s > {max_duration:.0f}s max)")
continue
print(f"\n{'#' * 60}")
print(f"# Testing: {audio_path.name} ({duration:.1f}s)")
print(f"{'#' * 60}")
result = await run_test(
engine, audio, chunk_ms, realtime,
audio_file=audio_path.name, backend=backend, policy=policy, lan=file_lan,
)
results.append(result)
return results
def print_benchmark_summary(results: List[TestResult]):
"""Print a tabular summary of all test results."""
print(f"\n{'=' * 110}")
print("BENCHMARK SUMMARY")
print(f"{'=' * 110}")
print(
f"{'File':<40} {'Duration':>8} {'Time':>8} {'RTF':>6} "
f"{'WER':>7} {'MAE(s)':>7} {'Lines':>5}"
)
print(f"{'-' * 110}")
for r in results:
wer_str = f"{r.wer:.2%}" if r.wer is not None else " -"
mae_str = f"{r.timestamp_mae:.3f}" if r.timestamp_mae is not None else " -"
print(
f"{r.audio_file:<40} {r.audio_duration_s:>7.1f}s {r.processing_time_s:>7.1f}s "
f"{r.rtf:>5.2f}x {wer_str:>7} {mae_str:>7} {r.n_lines:>5}"
)
print(f"{'-' * 110}")
total_audio = sum(r.audio_duration_s for r in results)
total_time = sum(r.processing_time_s for r in results)
avg_rtf = total_time / total_audio if total_audio > 0 else 0
wer_vals = [r.wer for r in results if r.wer is not None]
avg_wer_str = f"{sum(wer_vals)/len(wer_vals):.2%}" if wer_vals else " -"
mae_vals = [r.timestamp_mae for r in results if r.timestamp_mae is not None]
avg_mae_str = f"{sum(mae_vals)/len(mae_vals):.3f}" if mae_vals else " -"
print(
f"{'TOTAL/AVG':<40} {total_audio:>7.1f}s {total_time:>7.1f}s "
f"{avg_rtf:>5.2f}x {avg_wer_str:>7} {avg_mae_str:>7}"
)
print(f"{'=' * 110}")
# Print transcription excerpts
print("\nTRANSCRIPTIONS:")
print(f"{'-' * 110}")
for r in results:
excerpt = r.transcription[:120] + "..." if len(r.transcription) > 120 else r.transcription
print(f" {r.audio_file}:")
print(f" {excerpt}")
print(f"{'=' * 110}")
def detect_available_backends() -> List[dict]:
"""Probe which backends can be imported and return (backend, policy) combos.
Returns list of dicts with keys: backend, policy, description.
"""
combos = []
# faster-whisper
try:
import faster_whisper # noqa: F401
combos.append({"backend": "faster-whisper", "policy": "localagreement", "description": "faster-whisper + LocalAgreement"})
combos.append({"backend": "faster-whisper", "policy": "simulstreaming", "description": "faster-whisper + SimulStreaming"})
except ImportError:
pass
# mlx-whisper (macOS only)
try:
import mlx_whisper # noqa: F401
combos.append({"backend": "mlx-whisper", "policy": "localagreement", "description": "mlx-whisper + LocalAgreement"})
combos.append({"backend": "mlx-whisper", "policy": "simulstreaming", "description": "mlx-whisper + SimulStreaming"})
except ImportError:
pass
# openai-whisper
try:
import whisper # noqa: F401
combos.append({"backend": "whisper", "policy": "localagreement", "description": "openai-whisper + LocalAgreement"})
combos.append({"backend": "whisper", "policy": "simulstreaming", "description": "openai-whisper + SimulStreaming"})
except ImportError:
pass
# voxtral-mlx
try:
from whisperlivekit.voxtral_mlx import VoxtralMLXModel # noqa: F401
combos.append({"backend": "voxtral-mlx", "policy": "voxtral", "description": "voxtral-mlx (MLX)"})
except ImportError:
pass
# voxtral (HuggingFace)
try:
from transformers import AutoModelForSpeechSeq2Seq # noqa: F401
combos.append({"backend": "voxtral", "policy": "voxtral", "description": "voxtral (HuggingFace)"})
except ImportError:
pass
return combos
def print_cross_backend_comparison(all_results: List[TestResult]):
"""Print a comparison table across backends and policies."""
print(f"\n{'=' * 110}")
print("CROSS-BACKEND BENCHMARK COMPARISON")
print(f"{'=' * 110}")
print(
f"{'Backend':<18} {'Policy':<16} {'File':<30} "
f"{'WER':>7} {'RTF':>6} {'MAE(s)':>7} {'MaxΔ(s)':>8}"
)
print(f"{'-' * 110}")
for r in all_results:
wer_str = f"{r.wer:.2%}" if r.wer is not None else " -"
rtf_str = f"{r.rtf:.2f}x"
mae_str = f"{r.timestamp_mae:.3f}" if r.timestamp_mae is not None else " -"
max_str = f"{r.timestamp_max_delta:.3f}" if r.timestamp_max_delta is not None else " -"
# Truncate filename for readability
fname = r.audio_file[:28] + ".." if len(r.audio_file) > 30 else r.audio_file
print(
f"{r.backend:<18} {r.policy:<16} {fname:<30} "
f"{wer_str:>7} {rtf_str:>6} {mae_str:>7} {max_str:>8}"
)
print(f"{'-' * 110}")
# Per-backend averages
from collections import defaultdict
by_combo = defaultdict(list)
for r in all_results:
by_combo[(r.backend, r.policy)].append(r)
print(f"\n{'Backend':<18} {'Policy':<16} {'Avg WER':>8} {'Avg RTF':>8} {'Avg MAE':>8} {'Files':>6}")
print(f"{'-' * 80}")
for (backend, policy), group in sorted(by_combo.items()):
wer_vals = [r.wer for r in group if r.wer is not None]
rtf_vals = [r.rtf for r in group]
mae_vals = [r.timestamp_mae for r in group if r.timestamp_mae is not None]
avg_wer = f"{sum(wer_vals)/len(wer_vals):.2%}" if wer_vals else " -"
avg_rtf = f"{sum(rtf_vals)/len(rtf_vals):.2f}x"
avg_mae = f"{sum(mae_vals)/len(mae_vals):.3f}" if mae_vals else " -"
print(
f"{backend:<18} {policy:<16} {avg_wer:>8} {avg_rtf:>8} {avg_mae:>8} {len(group):>6}"
)
print(f"{'=' * 110}")
def _quiet_loggers(verbose: bool):
"""Set internal module log levels to reduce noise."""
if verbose:
logging.getLogger().setLevel(logging.DEBUG)
else:
for mod in (
"whisperlivekit.audio_processor", "whisperlivekit.simul_whisper",
"whisperlivekit.tokens_alignment", "whisperlivekit.simul_whisper.align_att_base",
"whisperlivekit.simul_whisper.simul_whisper",
):
logging.getLogger(mod).setLevel(logging.WARNING)
async def run_benchmark(
audio_files: List[Path], chunk_ms: int, realtime: bool,
model_size: str, lan: str, max_duration: float, vac: bool,
verbose: bool,
) -> List[TestResult]:
"""Run benchmark across all available backend+policy combinations."""
combos = detect_available_backends()
if not combos:
logger.error("No backends available. Install at least one ASR backend.")
return []
logger.info(f"Detected {len(combos)} backend+policy combinations:")
for c in combos:
logger.info(f" - {c['description']}")
all_results = []
for i, combo in enumerate(combos, 1):
backend = combo["backend"]
policy = combo["policy"]
desc = combo["description"]
print(f"\n{'*' * 70}")
print(f"* BENCHMARK {i}/{len(combos)}: {desc}")
print(f"{'*' * 70}")
try:
engine = create_engine(
backend, model_size, lan, vac=vac, policy=policy,
)
_quiet_loggers(verbose)
results = await run_all_tests(
engine, audio_files, chunk_ms, realtime,
backend=backend, policy=policy, lan=lan,
max_duration=max_duration,
)
all_results.extend(results)
except Exception as e:
logger.error(f"Failed to run {desc}: {e}")
import traceback
traceback.print_exc()
return all_results
def main():
parser = argparse.ArgumentParser(
description="Offline backend test harness (AudioProcessor-level)"
)
parser.add_argument(
"--backend", default="faster-whisper",
help="Backend: voxtral, voxtral-mlx, auto, faster-whisper, mlx-whisper, whisper.",
)
parser.add_argument(
"--policy", default="",
help="Override backend policy: localagreement, simulstreaming, voxtral.",
)
parser.add_argument(
"--audio", default=None,
help="Path to a single audio file (WAV, MP3, FLAC, etc.).",
)
parser.add_argument(
"--audio-dir", default=None,
help="Directory of audio files to test. Defaults to audio_tests/ if neither --audio nor --audio-dir given.",
)
parser.add_argument(
"--chunk-ms", type=int, default=100,
help="Chunk size in milliseconds (simulates real-time interval).",
)
parser.add_argument(
"--model", default="", dest="model_size",
help="Model size or HF repo ID.",
)
parser.add_argument("--lan", default="en", help="Language code.")
parser.add_argument(
"--no-realtime", action="store_true",
help="Skip real-time pacing between chunks (faster but less realistic).",
)
parser.add_argument(
"--no-vac", action="store_true",
help="Disable Voice Activity Classification (send all audio without silence filtering).",
)
parser.add_argument(
"--diarization", action="store_true",
help="Enable speaker diarization.",
)
parser.add_argument(
"--diarization-backend",
default="",
choices=["diart", "sortformer"],
help="Diarization backend when --diarization is enabled.",
)
parser.add_argument(
"--benchmark", action="store_true",
help="Run benchmark across all detected backend+policy combinations.",
)
parser.add_argument(
"--json", default=None, dest="json_output",
help="Write structured JSON results to this file.",
)
parser.add_argument(
"--max-duration", type=float, default=60.0,
help="Skip audio files longer than this many seconds (default: 60).",
)
parser.add_argument(
"--insert-silence", nargs=2, type=float, metavar=("SECS", "AT_SEC"),
action="append", default=[],
help="Insert SECS of silence at AT_SEC position. Can be repeated. "
"E.g.: --insert-silence 3.0 2.0 --insert-silence 5.0 7.0",
)
parser.add_argument(
"-v", "--verbose", action="store_true",
help="Show debug-level logs from all components.",
)
args = parser.parse_args()
realtime = not args.no_realtime
vac = not args.no_vac
# Resolve audio file(s)
if args.audio:
audio_files = [Path(args.audio)]
elif args.audio_dir:
audio_files = discover_audio_files(args.audio_dir)
elif AUDIO_TESTS_DIR.is_dir():
audio_files = discover_audio_files(str(AUDIO_TESTS_DIR))
else:
# Fall back to jfk.wav download
audio_files = [download_sample_audio()]
if not audio_files:
logger.error("No audio files found.")
sys.exit(1)
logger.info(f"Audio files: {[f.name for f in audio_files]}")
if args.benchmark:
# --- Multi-backend benchmark mode ---
all_results = asyncio.run(
run_benchmark(
audio_files, args.chunk_ms, realtime,
args.model_size, args.lan, args.max_duration, vac,
args.verbose,
)
)
if all_results:
print_cross_backend_comparison(all_results)
results = all_results
else:
# --- Single-backend mode ---
policy = args.policy
logger.info(f"Creating {args.backend} engine...")
engine = create_engine(
args.backend, args.model_size, args.lan,
diarization=args.diarization,
diarization_backend=args.diarization_backend,
vac=vac,
policy=policy,
)
logger.info("Engine ready.")
_quiet_loggers(args.verbose)
results = asyncio.run(
run_all_tests(
engine, audio_files, args.chunk_ms, realtime,
args.backend, policy, args.lan,
max_duration=args.max_duration,
silence_insertions=args.insert_silence or None,
)
)
if len(results) > 1:
print_benchmark_summary(results)
# JSON output
if args.json_output and results:
json_results = []
for r in results:
d = asdict(r)
d.pop("last_response", None) # too verbose for summary
json_results.append(d)
Path(args.json_output).write_text(
json.dumps(json_results, indent=2, ensure_ascii=False)
)
logger.info(f"Results written to {args.json_output}")
if __name__ == "__main__":
main()

30
tests/test_pipeline.py
View File

@@ -43,8 +43,17 @@ except ImportError:
pass
try:
from whisperlivekit.qwen3_asr import _patch_transformers_compat
_patch_transformers_compat()
from qwen_asr import Qwen3ASRModel # noqa: F401
AVAILABLE_BACKENDS.append("qwen3")
AVAILABLE_BACKENDS.append("qwen3-simul")
except (ImportError, Exception):
pass
try:
import mlx_qwen3_asr # noqa: F401
AVAILABLE_BACKENDS.append("qwen3-mlx")
except ImportError:
pass
@@ -53,6 +62,12 @@ BACKEND_CONFIG = {
"voxtral-mlx": {"backend": "voxtral-mlx", "lan": "en"},
"voxtral-hf": {"backend": "voxtral", "lan": "en"},
"qwen3": {"backend": "qwen3", "lan": "en"},
"qwen3-simul": {
"backend": "qwen3-simul",
"lan": "en",
"custom_alignment_heads": "scripts/alignment_heads_qwen3_asr_1.7B.json",
},
"qwen3-mlx": {"backend": "qwen3-mlx", "lan": "en"},
}
# Voxtral backends flush all words at once with proportionally-distributed
@@ -62,7 +77,7 @@ BACKEND_CONFIG = {
VOXTRAL_BACKENDS = {"voxtral-mlx", "voxtral-hf"}
# Backends that use batch-flush and may have non-monotonic timestamps
BATCH_FLUSH_BACKENDS = {"voxtral-mlx", "voxtral-hf", "qwen3"}
BATCH_FLUSH_BACKENDS = {"voxtral-mlx", "voxtral-hf", "qwen3", "qwen3-simul", "qwen3-mlx"}
def backend_kwargs(backend: str) -> dict:
@@ -176,8 +191,11 @@ async def test_text_appears_progressively(backend, medium_sample):
)
if len(non_empty) >= 3:
mid = len(non_empty) // 2
assert len(non_empty[-1]) > len(non_empty[mid]), (
# Check that text grew at SOME point during streaming.
# Compare first vs last non-empty snapshot rather than mid vs last,
# because some streaming backends (e.g. qwen3-simul) produce all text
# during the feed phase and the latter half of snapshots are stable.
assert len(non_empty[-1]) > len(non_empty[0]), (
f"Text not growing during streaming for {backend}"
)
@@ -250,10 +268,12 @@ async def test_silence_flushes_all_words(backend, medium_sample):
# Key assertion: silence must have committed most words.
# Some backends (voxtral-hf) produce extra words from right-padding
# at finish(), and MPS inference may leave some words in the pipeline.
# At least 50% of final words must be committed at silence time.
# Generative backends (qwen3-simul) keep producing new text on each
# inference call, so finish() adds significantly more words.
if words_at_finish > 3:
min_pct = 0.20 if backend in BATCH_FLUSH_BACKENDS else 0.50
flushed_pct = words_at_silence / words_at_finish
assert flushed_pct >= 0.50, (
assert flushed_pct >= min_pct, (
f"[{backend}] Only {flushed_pct:.0%} of words flushed at silence. "
f"At silence: {words_at_silence}, at finish: {words_at_finish}. "
f"Buffer at silence: '{buffer_at_silence}'"

1
whisperlivekit/basic_server.py
View File

@@ -13,6 +13,7 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
logging.getLogger().setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
logging.getLogger("whisperlivekit.qwen3_asr").setLevel(logging.DEBUG)
config = parse_args()
transcription_engine = None

34
whisperlivekit/benchmark/__init__.py Normal file
View File

@@ -0,0 +1,34 @@
"""WhisperLiveKit benchmark suite.
Comprehensive benchmarking of ASR backends using public datasets,
run through the same pipeline as real-time streaming.
Usage:
wlk bench # benchmark current backend
wlk bench --backend whisper --json results.json
wlk bench --languages en,fr,es # multilingual
wlk bench --quick # fast subset
Programmatic:
from whisperlivekit.benchmark import BenchmarkRunner
import asyncio
runner = BenchmarkRunner(backend="whisper", model_size="base")
report = asyncio.run(runner.run())
print(report.summary_table())
"""
from whisperlivekit.benchmark.datasets import (
BENCHMARK_CATALOG,
get_benchmark_samples,
)
from whisperlivekit.benchmark.metrics import BenchmarkReport, SampleResult
from whisperlivekit.benchmark.runner import BenchmarkRunner
__all__ = [
"BENCHMARK_CATALOG",
"BenchmarkReport",
"BenchmarkRunner",
"SampleResult",
"get_benchmark_samples",
]

105
whisperlivekit/benchmark/compat.py Normal file
View File

@@ -0,0 +1,105 @@
"""Backend detection and language compatibility matrix."""
import logging
from typing import Dict, List, Optional, Set
logger = logging.getLogger(__name__)
# Language support per backend.
# None means all Whisper-supported languages.
# A set means only those languages are supported.
BACKEND_LANGUAGES: Dict[str, Optional[Set[str]]] = {
"whisper": None,
"faster-whisper": None,
"mlx-whisper": None,
"voxtral-mlx": None,
"voxtral": None,
"qwen3": {
"zh", "en", "yue", "ar", "de", "fr", "es", "pt", "id", "it",
"ko", "ru", "th", "vi", "ja", "tr", "hi", "ms", "nl", "sv",
"da", "fi", "pl", "cs", "fa", "el", "hu", "mk", "ro",
},
"qwen3-simul": {
"zh", "en", "yue", "ar", "de", "fr", "es", "pt", "id", "it",
"ko", "ru", "th", "vi", "ja", "tr", "hi", "ms", "nl", "sv",
"da", "fi", "pl", "cs", "fa", "el", "hu", "mk", "ro",
},
}
def backend_supports_language(backend: str, language: str) -> bool:
"""Check if a backend supports a given language code."""
langs = BACKEND_LANGUAGES.get(backend)
if langs is None:
return True
return language in langs
def detect_available_backends() -> List[str]:
"""Probe which ASR backends are importable."""
backends = []
try:
import whisper # noqa: F401
backends.append("whisper")
except ImportError:
pass
try:
import faster_whisper # noqa: F401
backends.append("faster-whisper")
except ImportError:
pass
try:
import mlx_whisper # noqa: F401
backends.append("mlx-whisper")
except ImportError:
pass
try:
import mlx.core # noqa: F401
from whisperlivekit.voxtral_mlx.loader import load_voxtral_model # noqa: F401
backends.append("voxtral-mlx")
except ImportError:
pass
try:
from transformers import VoxtralRealtimeForConditionalGeneration # noqa: F401
backends.append("voxtral")
except ImportError:
pass
try:
from whisperlivekit.qwen3_asr import _patch_transformers_compat
_patch_transformers_compat()
from qwen_asr import Qwen3ASRModel # noqa: F401
backends.append("qwen3")
backends.append("qwen3-simul")
except (ImportError, Exception):
pass
return backends
def resolve_backend(backend: str) -> str:
"""Resolve 'auto' to the best available backend."""
if backend != "auto":
return backend
available = detect_available_backends()
if not available:
raise RuntimeError(
"No ASR backend available. Install at least one: "
"pip install openai-whisper, faster-whisper, or mlx-whisper"
)
# Priority order
priority = [
"faster-whisper", "mlx-whisper", "voxtral-mlx", "voxtral",
"qwen3", "qwen3-simul", "whisper",
]
for p in priority:
if p in available:
return p
return available[0]

561
whisperlivekit/benchmark/datasets.py Normal file
View File

@@ -0,0 +1,561 @@
"""Benchmark audio datasets from public HuggingFace repositories.
Downloads curated samples across languages, noise conditions, and speaker
configurations. All datasets are public and freely accessible — no auth
tokens required.
Samples are cached in ~/.cache/whisperlivekit/benchmark_data/ and reused
across benchmark runs.
Datasets used:
- LibriSpeech test-clean (English, clean, single speaker)
- LibriSpeech test-other (English, noisy/hard, single speaker)
- Multilingual LibriSpeech (French, Spanish, German, Portuguese, Italian, Polish, Dutch)
- AMI (English, multi-speaker meeting)
"""
import json
import logging
import wave
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Set
import numpy as np
logger = logging.getLogger(__name__)
CACHE_DIR = Path.home() / ".cache" / "whisperlivekit" / "benchmark_data"
METADATA_FILE = "benchmark_metadata.json"
@dataclass
class BenchmarkSample:
"""A benchmark audio sample with metadata and ground truth."""
name: str
path: str
reference: str
duration: float
language: str
category: str # "clean", "noisy", "multilingual", "meeting"
sample_rate: int = 16000
n_speakers: int = 1
source: str = ""
tags: Set[str] = field(default_factory=set)
def to_dict(self) -> Dict:
return {
"name": self.name,
"file": Path(self.path).name,
"reference": self.reference,
"duration": self.duration,
"language": self.language,
"category": self.category,
"sample_rate": self.sample_rate,
"n_speakers": self.n_speakers,
"source": self.source,
"tags": list(self.tags),
}
# ---------------------------------------------------------------------------
# Dataset catalog — defines what to download
# ---------------------------------------------------------------------------
BENCHMARK_CATALOG = {
# English clean (LibriSpeech test-clean)
"en_clean_short": {
"dataset": "openslr/librispeech_asr",
"config": "clean",
"split": "test",
"language": "en",
"category": "clean",
"n_samples": 1,
"skip": 0,
"tags": {"short"},
},
"en_clean_medium": {
"dataset": "openslr/librispeech_asr",
"config": "clean",
"split": "test",
"language": "en",
"category": "clean",
"n_samples": 1,
"skip": 1,
"tags": {"medium"},
},
# English noisy (LibriSpeech test-other)
"en_noisy_1": {
"dataset": "openslr/librispeech_asr",
"config": "other",
"split": "test",
"language": "en",
"category": "noisy",
"n_samples": 1,
"skip": 0,
"tags": {"accented"},
},
"en_noisy_2": {
"dataset": "openslr/librispeech_asr",
"config": "other",
"split": "test",
"language": "en",
"category": "noisy",
"n_samples": 1,
"skip": 1,
"tags": {"accented"},
},
# French (Multilingual LibriSpeech)
"fr_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "french",
"split": "test",
"language": "fr",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
"fr_clean_2": {
"dataset": "facebook/multilingual_librispeech",
"config": "french",
"split": "test",
"language": "fr",
"category": "multilingual",
"n_samples": 1,
"skip": 1,
"tags": set(),
},
# Spanish (Multilingual LibriSpeech)
"es_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "spanish",
"split": "test",
"language": "es",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# German (Multilingual LibriSpeech)
"de_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "german",
"split": "test",
"language": "de",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# Portuguese (Multilingual LibriSpeech)
"pt_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "portuguese",
"split": "test",
"language": "pt",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# Italian (Multilingual LibriSpeech)
"it_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "italian",
"split": "test",
"language": "it",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# Polish (Multilingual LibriSpeech)
"pl_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "polish",
"split": "test",
"language": "pl",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# Dutch (Multilingual LibriSpeech)
"nl_clean_1": {
"dataset": "facebook/multilingual_librispeech",
"config": "dutch",
"split": "test",
"language": "nl",
"category": "multilingual",
"n_samples": 1,
"skip": 0,
"tags": set(),
},
# English multi-speaker meeting (AMI)
"en_meeting": {
"dataset": "edinburghcstr/ami",
"config": "ihm",
"split": "test",
"language": "en",
"category": "meeting",
"n_samples": 1,
"skip": 0,
"tags": {"multi_speaker", "long"},
"max_duration": 60.0,
},
}
# Quick mode: subset of samples for fast smoke tests
QUICK_SAMPLES = {"en_clean_short", "en_clean_medium", "en_noisy_1", "fr_clean_1"}
# ---------------------------------------------------------------------------
# Audio utilities
# ---------------------------------------------------------------------------
def _save_wav(path: Path, audio: np.ndarray, sample_rate: int = 16000) -> None:
if audio.ndim > 1:
audio = audio.mean(axis=-1)
if audio.dtype in (np.float32, np.float64):
audio = np.clip(audio, -1.0, 1.0)
audio = (audio * 32767).astype(np.int16)
elif audio.dtype != np.int16:
audio = audio.astype(np.int16)
path.parent.mkdir(parents=True, exist_ok=True)
with wave.open(str(path), "w") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(sample_rate)
wf.writeframes(audio.tobytes())
def _decode_audio(audio_bytes: bytes) -> tuple:
import io
import soundfile as sf
audio_array, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
return np.array(audio_array, dtype=np.float32), sr
def _ensure_datasets():
try:
import datasets # noqa: F401
except ImportError:
raise ImportError(
"The 'datasets' package is required for benchmark data. "
"Install with: pip install whisperlivekit[test]"
)
# ---------------------------------------------------------------------------
# Download functions per dataset type
# ---------------------------------------------------------------------------
def _download_librispeech(config: str, n_samples: int, skip: int,
category: str, language: str,
prefix: str) -> List[Dict]:
"""Download from openslr/librispeech_asr (clean or other)."""
_ensure_datasets()
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info("Downloading LibriSpeech %s samples...", config)
ds = load_dataset(
"openslr/librispeech_asr", config, split="test", streaming=True,
)
ds = ds.cast_column("audio", Audio(decode=False))
samples = []
for i, item in enumerate(ds):
if i < skip:
continue
if len(samples) >= n_samples:
break
audio_array, sr = _decode_audio(item["audio"]["bytes"])
duration = len(audio_array) / sr
text = item["text"]
wav_name = f"{prefix}_{i}.wav"
_save_wav(CACHE_DIR / wav_name, audio_array, sr)
samples.append({
"file": wav_name,
"reference": text,
"duration": round(duration, 2),
"sample_rate": sr,
"language": language,
"category": category,
"n_speakers": 1,
"source": f"openslr/librispeech_asr ({config})",
})
logger.info(" %.1fs - %s", duration, text[:60])
return samples
def _download_mls(config: str, n_samples: int, skip: int,
language: str, prefix: str) -> List[Dict]:
"""Download from facebook/multilingual_librispeech."""
_ensure_datasets()
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info("Downloading MLS %s samples...", config)
ds = load_dataset(
"facebook/multilingual_librispeech", config, split="test", streaming=True,
)
ds = ds.cast_column("audio", Audio(decode=False))
samples = []
for i, item in enumerate(ds):
if i < skip:
continue
if len(samples) >= n_samples:
break
audio_array, sr = _decode_audio(item["audio"]["bytes"])
duration = len(audio_array) / sr
text = item.get("text", item.get("transcript", ""))
wav_name = f"{prefix}_{i}.wav"
_save_wav(CACHE_DIR / wav_name, audio_array, sr)
samples.append({
"file": wav_name,
"reference": text,
"duration": round(duration, 2),
"sample_rate": sr,
"language": language,
"category": "multilingual",
"n_speakers": 1,
"source": f"facebook/multilingual_librispeech ({config})",
})
logger.info(" [%s] %.1fs - %s", language, duration, text[:60])
return samples
def _download_fleurs(config: str, n_samples: int, skip: int,
language: str, prefix: str) -> List[Dict]:
"""Download from google/fleurs."""
_ensure_datasets()
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info("Downloading FLEURS %s samples...", config)
ds = load_dataset(
"google/fleurs", config, split="test", streaming=True,
)
ds = ds.cast_column("audio", Audio(decode=False))
samples = []
for i, item in enumerate(ds):
if i < skip:
continue
if len(samples) >= n_samples:
break
audio_array, sr = _decode_audio(item["audio"]["bytes"])
duration = len(audio_array) / sr
text = item.get("transcription", item.get("raw_transcription", ""))
wav_name = f"{prefix}_{i}.wav"
_save_wav(CACHE_DIR / wav_name, audio_array, sr)
samples.append({
"file": wav_name,
"reference": text,
"duration": round(duration, 2),
"sample_rate": sr,
"language": language,
"category": "multilingual",
"n_speakers": 1,
"source": f"google/fleurs ({config})",
})
logger.info(" [%s] %.1fs - %s", language, duration, text[:60])
return samples
def _download_ami(max_duration: float = 60.0) -> List[Dict]:
"""Download one AMI meeting segment with multiple speakers."""
_ensure_datasets()
import datasets.config
datasets.config.TORCHCODEC_AVAILABLE = False
from datasets import Audio, load_dataset
logger.info("Downloading AMI meeting sample...")
ds = load_dataset("edinburghcstr/ami", "ihm", split="test", streaming=True)
ds = ds.cast_column("audio", Audio(decode=False))
meeting_id = None
audio_arrays = []
texts = []
sample_rate = None
for item in ds:
mid = item.get("meeting_id", "unknown")
if meeting_id is None:
meeting_id = mid
elif mid != meeting_id:
break
audio_array, sr = _decode_audio(item["audio"]["bytes"])
sample_rate = sr
texts.append(item.get("text", ""))
audio_arrays.append(audio_array)
total_dur = sum(len(a) / sr for a in audio_arrays)
if total_dur > max_duration:
break
if not audio_arrays:
return []
full_audio = np.concatenate(audio_arrays)
duration = len(full_audio) / sample_rate
reference = " ".join(t for t in texts if t)
wav_name = "ami_meeting.wav"
_save_wav(CACHE_DIR / wav_name, full_audio, sample_rate)
logger.info(" AMI meeting: %.1fs, %d utterances", duration, len(texts))
return [{
"file": wav_name,
"reference": reference,
"duration": round(duration, 2),
"sample_rate": sample_rate,
"language": "en",
"category": "meeting",
"n_speakers": 4,
"source": f"edinburghcstr/ami (ihm, meeting {meeting_id})",
}]
# ---------------------------------------------------------------------------
# Dispatcher — routes catalog entries to download functions
# ---------------------------------------------------------------------------
def _download_catalog_entry(name: str, spec: Dict) -> List[Dict]:
"""Download a single catalog entry and return metadata dicts."""
dataset = spec["dataset"]
config = spec.get("config", "")
n_samples = spec.get("n_samples", 1)
skip = spec.get("skip", 0)
language = spec["language"]
category = spec["category"]
if dataset == "openslr/librispeech_asr":
return _download_librispeech(
config=config, n_samples=n_samples, skip=skip,
category=category, language=language, prefix=name,
)
elif dataset == "facebook/multilingual_librispeech":
return _download_mls(
config=config, n_samples=n_samples, skip=skip,
language=language, prefix=name,
)
elif dataset == "google/fleurs":
return _download_fleurs(
config=config, n_samples=n_samples, skip=skip,
language=language, prefix=name,
)
elif dataset == "edinburghcstr/ami":
return _download_ami(max_duration=spec.get("max_duration", 60.0))
else:
logger.warning("Unknown dataset: %s", dataset)
return []
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def get_benchmark_samples(
languages: Optional[List[str]] = None,
categories: Optional[List[str]] = None,
quick: bool = False,
force: bool = False,
) -> List[BenchmarkSample]:
"""Download and return benchmark samples, filtered by language/category.
Args:
languages: List of language codes to include (None = all).
categories: List of categories to include (None = all).
quick: If True, only download a small subset for smoke tests.
force: Re-download even if cached.
Returns:
List of BenchmarkSample objects ready for benchmarking.
"""
CACHE_DIR.mkdir(parents=True, exist_ok=True)
meta_path = CACHE_DIR / METADATA_FILE
# Load cached metadata
cached = {}
if meta_path.exists() and not force:
cached = json.loads(meta_path.read_text())
# Determine which entries to download
entries = BENCHMARK_CATALOG
if quick:
entries = {k: v for k, v in entries.items() if k in QUICK_SAMPLES}
if languages:
lang_set = set(languages)
entries = {k: v for k, v in entries.items() if v["language"] in lang_set}
if categories:
cat_set = set(categories)
entries = {k: v for k, v in entries.items() if v["category"] in cat_set}
# Download missing entries
all_meta = cached.get("samples", {})
for name, spec in entries.items():
if name in all_meta and not force:
# Check file exists
file_path = CACHE_DIR / all_meta[name][0]["file"]
if file_path.exists():
continue
logger.info("Downloading benchmark sample: %s", name)
try:
downloaded = _download_catalog_entry(name, spec)
if downloaded:
all_meta[name] = downloaded
except Exception as e:
logger.warning("Failed to download %s: %s", name, e)
# Save metadata
meta_path.write_text(json.dumps({"samples": all_meta}, indent=2))
# Build BenchmarkSample objects
samples = []
for name, spec in entries.items():
if name not in all_meta:
continue
for meta in all_meta[name]:
file_path = CACHE_DIR / meta["file"]
if not file_path.exists():
continue
catalog_entry = BENCHMARK_CATALOG.get(name, {})
samples.append(BenchmarkSample(
name=name,
path=str(file_path),
reference=meta["reference"],
duration=meta["duration"],
language=meta["language"],
category=meta["category"],
sample_rate=meta.get("sample_rate", 16000),
n_speakers=meta.get("n_speakers", 1),
source=meta.get("source", ""),
tags=set(catalog_entry.get("tags", set())),
))
logger.info("Loaded %d benchmark samples", len(samples))
return samples

273
whisperlivekit/benchmark/metrics.py Normal file
View File

@@ -0,0 +1,273 @@
"""Benchmark result data structures and aggregation."""
import platform
import subprocess
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
@dataclass
class SampleResult:
"""Result from benchmarking one audio sample."""
sample_name: str
language: str
category: str
duration_s: float
# Quality
wer: float
wer_details: Dict[str, int]
# Speed
processing_time_s: float
rtf: float
# Latency (from SessionMetrics)
avg_latency_ms: float = 0.0
p95_latency_ms: float = 0.0
n_transcription_calls: int = 0
# Pipeline stats
n_lines: int = 0
n_tokens: int = 0
# Timing quality
timing_valid: bool = True
timing_monotonic: bool = True
# Memory
peak_memory_mb: Optional[float] = None
# Texts
hypothesis: str = ""
reference: str = ""
# Source
source: str = ""
tags: List[str] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
return {
"sample": self.sample_name,
"language": self.language,
"category": self.category,
"duration_s": round(self.duration_s, 2),
"wer": round(self.wer, 4),
"wer_details": self.wer_details,
"processing_time_s": round(self.processing_time_s, 2),
"rtf": round(self.rtf, 3),
"avg_latency_ms": round(self.avg_latency_ms, 1),
"p95_latency_ms": round(self.p95_latency_ms, 1),
"n_transcription_calls": self.n_transcription_calls,
"n_lines": self.n_lines,
"n_tokens": self.n_tokens,
"timing_valid": self.timing_valid,
"timing_monotonic": self.timing_monotonic,
"peak_memory_mb": round(self.peak_memory_mb, 1) if self.peak_memory_mb else None,
"hypothesis": self.hypothesis,
"reference": self.reference,
"source": self.source,
"tags": self.tags,
}
@dataclass
class BenchmarkReport:
"""Aggregated benchmark report with system info and per-sample results."""
backend: str
model_size: str
timestamp: str = field(default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%S"))
system_info: Dict[str, Any] = field(default_factory=dict)
results: List[SampleResult] = field(default_factory=list)
# --- Aggregate properties ---
@property
def n_samples(self) -> int:
return len(self.results)
@property
def total_audio_s(self) -> float:
return sum(r.duration_s for r in self.results)
@property
def total_processing_s(self) -> float:
return sum(r.processing_time_s for r in self.results)
@property
def avg_wer(self) -> float:
if not self.results:
return 0.0
return sum(r.wer for r in self.results) / len(self.results)
@property
def weighted_wer(self) -> float:
"""Micro-averaged WER: total errors / total reference words."""
total_errors = sum(
r.wer_details.get("substitutions", 0) +
r.wer_details.get("insertions", 0) +
r.wer_details.get("deletions", 0)
for r in self.results
)
total_ref = sum(r.wer_details.get("ref_words", 0) for r in self.results)
return total_errors / max(total_ref, 1)
@property
def avg_rtf(self) -> float:
if not self.results:
return 0.0
return sum(r.rtf for r in self.results) / len(self.results)
@property
def overall_rtf(self) -> float:
if self.total_audio_s <= 0:
return 0.0
return self.total_processing_s / self.total_audio_s
@property
def avg_latency_ms(self) -> float:
vals = [r.avg_latency_ms for r in self.results if r.avg_latency_ms > 0]
return sum(vals) / len(vals) if vals else 0.0
@property
def p95_latency_ms(self) -> float:
vals = [r.p95_latency_ms for r in self.results if r.p95_latency_ms > 0]
return sum(vals) / len(vals) if vals else 0.0
# --- Per-dimension breakdowns ---
def _group_by(self, key: str) -> Dict[str, List[SampleResult]]:
groups: Dict[str, List[SampleResult]] = {}
for r in self.results:
k = getattr(r, key, "unknown")
groups.setdefault(k, []).append(r)
return groups
def wer_by_language(self) -> Dict[str, float]:
return {
lang: sum(r.wer for r in group) / len(group)
for lang, group in sorted(self._group_by("language").items())
}
def rtf_by_language(self) -> Dict[str, float]:
return {
lang: sum(r.rtf for r in group) / len(group)
for lang, group in sorted(self._group_by("language").items())
}
def wer_by_category(self) -> Dict[str, float]:
return {
cat: sum(r.wer for r in group) / len(group)
for cat, group in sorted(self._group_by("category").items())
}
@property
def languages(self) -> List[str]:
return sorted(set(r.language for r in self.results))
@property
def categories(self) -> List[str]:
return sorted(set(r.category for r in self.results))
def to_dict(self) -> Dict[str, Any]:
return {
"benchmark_version": "1.0",
"timestamp": self.timestamp,
"system_info": self.system_info,
"config": {
"backend": self.backend,
"model_size": self.model_size,
},
"summary": {
"n_samples": self.n_samples,
"total_audio_s": round(self.total_audio_s, 1),
"total_processing_s": round(self.total_processing_s, 1),
"avg_wer": round(self.avg_wer, 4),
"weighted_wer": round(self.weighted_wer, 4),
"avg_rtf": round(self.avg_rtf, 3),
"overall_rtf": round(self.overall_rtf, 3),
"avg_latency_ms": round(self.avg_latency_ms, 1),
"p95_latency_ms": round(self.p95_latency_ms, 1),
"wer_by_language": {
k: round(v, 4) for k, v in self.wer_by_language().items()
},
"rtf_by_language": {
k: round(v, 3) for k, v in self.rtf_by_language().items()
},
"wer_by_category": {
k: round(v, 4) for k, v in self.wer_by_category().items()
},
},
"results": [r.to_dict() for r in self.results],
}
def get_system_info() -> Dict[str, Any]:
"""Collect system metadata for the benchmark report."""
info: Dict[str, Any] = {
"platform": platform.platform(),
"machine": platform.machine(),
"python_version": platform.python_version(),
}
# CPU info
try:
chip = subprocess.check_output(
["sysctl", "-n", "machdep.cpu.brand_string"], text=True,
).strip()
info["cpu"] = chip
except Exception:
info["cpu"] = platform.processor()
# RAM
try:
mem_bytes = int(
subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True).strip()
)
info["ram_gb"] = round(mem_bytes / (1024**3))
except Exception:
try:
import os
pages = os.sysconf("SC_PHYS_PAGES")
page_size = os.sysconf("SC_PAGE_SIZE")
info["ram_gb"] = round(pages * page_size / (1024**3))
except Exception:
info["ram_gb"] = None
# Accelerator
try:
import torch
if torch.cuda.is_available():
info["accelerator"] = torch.cuda.get_device_name(0)
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
info["accelerator"] = "Apple Silicon (MPS)"
else:
info["accelerator"] = "CPU"
except ImportError:
info["accelerator"] = "CPU"
# Backend versions
versions = {}
for pkg, name in [
("faster_whisper", "faster-whisper"),
("whisper", "openai-whisper"),
("mlx_whisper", "mlx-whisper"),
("transformers", "transformers"),
("torch", "torch"),
]:
try:
mod = __import__(pkg)
versions[name] = getattr(mod, "__version__", "installed")
except ImportError:
pass
try:
import mlx.core as mx
versions["mlx"] = mx.__version__
except ImportError:
pass
info["backend_versions"] = versions
return info

161
whisperlivekit/benchmark/report.py Normal file
View File

@@ -0,0 +1,161 @@
"""Benchmark report formatting — terminal tables and JSON export."""
import json
import sys
from pathlib import Path
from typing import TextIO
from whisperlivekit.benchmark.metrics import BenchmarkReport
# ANSI color codes
GREEN = "\033[32m"
YELLOW = "\033[33m"
RED = "\033[31m"
CYAN = "\033[36m"
BOLD = "\033[1m"
DIM = "\033[2m"
RESET = "\033[0m"
def _wer_color(wer: float) -> str:
if wer < 0.15:
return GREEN
elif wer < 0.30:
return YELLOW
return RED
def _rtf_color(rtf: float) -> str:
if rtf < 0.5:
return GREEN
elif rtf < 1.0:
return YELLOW
return RED
def _lat_color(ms: float) -> str:
if ms < 500:
return GREEN
elif ms < 1000:
return YELLOW
return RED
def print_report(report: BenchmarkReport, out: TextIO = sys.stderr) -> None:
"""Print a comprehensive benchmark report to the terminal."""
w = out.write
# Header
w(f"\n{BOLD} WhisperLiveKit Benchmark Report{RESET}\n")
w(f" {'─' * 72}\n")
si = report.system_info
w(f" Backend: {CYAN}{report.backend}{RESET}\n")
w(f" Model: {report.model_size}\n")
w(f" Accelerator: {si.get('accelerator', 'unknown')}\n")
w(f" CPU: {si.get('cpu', 'unknown')}\n")
w(f" RAM: {si.get('ram_gb', '?')} GB\n")
w(f" Timestamp: {report.timestamp}\n")
w(f" {'─' * 72}\n\n")
# Per-sample table
w(f" {BOLD}{'Sample':<20} {'Lang':>4} {'Dur':>5} {'WER':>7} "
f"{'RTF':>6} {'Lat(avg)':>8} {'Lat(p95)':>8} {'Calls':>5} {'Lines':>5}{RESET}\n")
w(f" {'─' * 72}\n")
for r in report.results:
wc = _wer_color(r.wer)
rc = _rtf_color(r.rtf)
lc = _lat_color(r.avg_latency_ms)
name = r.sample_name[:20]
w(f" {name:<20} {r.language:>4} {r.duration_s:>4.1f}s "
f"{wc}{r.wer * 100:>6.1f}%{RESET} "
f"{rc}{r.rtf:>5.2f}x{RESET} "
f"{lc}{r.avg_latency_ms:>7.0f}ms{RESET} "
f"{lc}{r.p95_latency_ms:>7.0f}ms{RESET} "
f"{r.n_transcription_calls:>5} {r.n_lines:>5}\n")
# Timing warnings
if not r.timing_valid:
w(f" {' ' * 20} {RED}⚠ invalid timestamps{RESET}\n")
if not r.timing_monotonic:
w(f" {' ' * 20} {YELLOW}⚠ non-monotonic timestamps{RESET}\n")
w(f" {'─' * 72}\n\n")
# Summary
w(f" {BOLD}Summary{RESET} ({report.n_samples} samples, "
f"{report.total_audio_s:.1f}s total audio)\n\n")
wc = _wer_color(report.avg_wer)
rc = _rtf_color(report.overall_rtf)
lc = _lat_color(report.avg_latency_ms)
w(f" Avg WER (macro): {wc}{report.avg_wer * 100:>6.1f}%{RESET}\n")
w(f" Weighted WER: {_wer_color(report.weighted_wer)}"
f"{report.weighted_wer * 100:>6.1f}%{RESET}\n")
w(f" Overall RTF: {rc}{report.overall_rtf:>6.3f}x{RESET} "
f"({report.total_processing_s:.1f}s for {report.total_audio_s:.1f}s audio)\n")
w(f" Avg latency: {lc}{report.avg_latency_ms:>6.0f}ms{RESET}\n")
w(f" P95 latency: {_lat_color(report.p95_latency_ms)}"
f"{report.p95_latency_ms:>6.0f}ms{RESET}\n")
# Per-language breakdown
wer_by_lang = report.wer_by_language()
rtf_by_lang = report.rtf_by_language()
if len(wer_by_lang) > 1:
w(f"\n {BOLD}By Language{RESET}\n")
w(f" {'─' * 40}\n")
w(f" {'Lang':>4} {'WER':>7} {'RTF':>6} {'Samples':>7}\n")
w(f" {'─' * 34}\n")
lang_groups = {}
for r in report.results:
lang_groups.setdefault(r.language, []).append(r)
for lang in sorted(lang_groups):
group = lang_groups[lang]
avg_wer = sum(r.wer for r in group) / len(group)
avg_rtf = sum(r.rtf for r in group) / len(group)
wc = _wer_color(avg_wer)
rc = _rtf_color(avg_rtf)
w(f" {lang:>4} {wc}{avg_wer * 100:>6.1f}%{RESET} "
f"{rc}{avg_rtf:>5.2f}x{RESET} {len(group):>7}\n")
# Per-category breakdown
wer_by_cat = report.wer_by_category()
if len(wer_by_cat) > 1:
w(f"\n {BOLD}By Category{RESET}\n")
w(f" {'─' * 40}\n")
w(f" {'Category':>12} {'WER':>7} {'Samples':>7}\n")
w(f" {'─' * 30}\n")
cat_groups = {}
for r in report.results:
cat_groups.setdefault(r.category, []).append(r)
for cat in sorted(cat_groups):
group = cat_groups[cat]
avg_wer = sum(r.wer for r in group) / len(group)
wc = _wer_color(avg_wer)
w(f" {cat:>12} {wc}{avg_wer * 100:>6.1f}%{RESET} {len(group):>7}\n")
w(f"\n {'─' * 72}\n\n")
def print_transcriptions(report: BenchmarkReport, out: TextIO = sys.stderr) -> None:
"""Print hypothesis vs reference for each sample."""
w = out.write
w(f"\n {BOLD}Transcriptions{RESET}\n")
w(f" {'─' * 72}\n")
for r in report.results:
wc = _wer_color(r.wer)
w(f"\n {BOLD}{r.sample_name}{RESET} ({r.language}, {r.category}) "
f"WER={wc}{r.wer * 100:.1f}%{RESET}\n")
ref = r.reference[:120] + "..." if len(r.reference) > 120 else r.reference
hyp = r.hypothesis[:120] + "..." if len(r.hypothesis) > 120 else r.hypothesis
w(f" {DIM}ref: {ref}{RESET}\n")
w(f" hyp: {hyp}\n")
w(f"\n {'─' * 72}\n\n")
def write_json(report: BenchmarkReport, path: str) -> None:
"""Export the full report as JSON."""
Path(path).write_text(json.dumps(report.to_dict(), indent=2, ensure_ascii=False))

181
whisperlivekit/benchmark/runner.py Normal file
View File

@@ -0,0 +1,181 @@
"""Benchmark runner — orchestrates runs through TestHarness."""
import logging
import resource
import time
from typing import Callable, List, Optional
from whisperlivekit.benchmark.compat import backend_supports_language, resolve_backend
from whisperlivekit.benchmark.datasets import BenchmarkSample, get_benchmark_samples
from whisperlivekit.benchmark.metrics import BenchmarkReport, SampleResult, get_system_info
logger = logging.getLogger(__name__)
class BenchmarkRunner:
"""Orchestrates benchmark runs through TestHarness.
Args:
backend: ASR backend name or "auto".
model_size: Model size (e.g. "base", "large-v3").
languages: Language codes to benchmark (None = all available).
categories: Categories to benchmark (None = all).
quick: Use a small subset for fast smoke tests.
speed: Feed speed (0 = instant, 1.0 = real-time).
on_progress: Callback(sample_name, i, total) for progress updates.
"""
def __init__(
self,
backend: str = "auto",
model_size: str = "base",
languages: Optional[List[str]] = None,
categories: Optional[List[str]] = None,
quick: bool = False,
speed: float = 0,
on_progress: Optional[Callable] = None,
):
self.backend = resolve_backend(backend)
self.model_size = model_size
self.languages = languages
self.categories = categories
self.quick = quick
self.speed = speed
self.on_progress = on_progress
async def run(self) -> BenchmarkReport:
"""Run the full benchmark suite and return a report."""
from whisperlivekit.metrics import compute_wer
from whisperlivekit.test_harness import TestHarness
# Get samples
samples = get_benchmark_samples(
languages=self.languages,
categories=self.categories,
quick=self.quick,
)
# Filter by backend language support
compatible = []
for s in samples:
if backend_supports_language(self.backend, s.language):
compatible.append(s)
else:
logger.info(
"Skipping %s (%s) — backend %s does not support %s",
s.name, s.language, self.backend, s.language,
)
samples = compatible
if not samples:
raise RuntimeError(
f"No benchmark samples available for backend={self.backend}, "
f"languages={self.languages}, categories={self.categories}"
)
# Build harness kwargs
harness_kwargs = {
"model_size": self.model_size,
"lan": "auto", # let the model auto-detect for multilingual
"pcm_input": True,
}
if self.backend not in ("auto",):
harness_kwargs["backend"] = self.backend
report = BenchmarkReport(
backend=self.backend,
model_size=self.model_size,
system_info=get_system_info(),
)
for i, sample in enumerate(samples):
if self.on_progress:
self.on_progress(sample.name, i, len(samples))
result = await self._run_sample(
sample, harness_kwargs, compute_wer,
)
report.results.append(result)
if self.on_progress:
self.on_progress("done", len(samples), len(samples))
return report
async def _run_sample(
self,
sample: BenchmarkSample,
harness_kwargs: dict,
compute_wer,
) -> SampleResult:
"""Benchmark a single sample through TestHarness."""
from whisperlivekit.test_harness import TestHarness
# Override language for the specific sample
kwargs = {**harness_kwargs, "lan": sample.language}
# Memory before
mem_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
t_start = time.perf_counter()
async with TestHarness(**kwargs) as h:
await h.feed(sample.path, speed=self.speed)
# Drain time scales with audio duration for slow backends
drain = max(5.0, sample.duration * 0.5)
await h.drain(drain)
state = await h.finish(timeout=120)
# Extract metrics from the pipeline
metrics = h.metrics
t_elapsed = time.perf_counter() - t_start
# Memory after
mem_after = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
# On macOS ru_maxrss is bytes, on Linux it's KB
import sys
divisor = 1024 * 1024 if sys.platform == "darwin" else 1024
mem_delta = (mem_after - mem_before) / divisor
# RTF
rtf = t_elapsed / sample.duration if sample.duration > 0 else 0
# WER
hypothesis = state.committed_text or state.text
wer_result = compute_wer(sample.reference, hypothesis)
# Latency from SessionMetrics
avg_lat = metrics.avg_latency_ms if metrics else 0
p95_lat = metrics.p95_latency_ms if metrics else 0
n_calls = metrics.n_transcription_calls if metrics else 0
n_tokens = metrics.n_tokens_produced if metrics else 0
return SampleResult(
sample_name=sample.name,
language=sample.language,
category=sample.category,
duration_s=sample.duration,
wer=wer_result["wer"],
wer_details={
"substitutions": wer_result["substitutions"],
"insertions": wer_result["insertions"],
"deletions": wer_result["deletions"],
"ref_words": wer_result["ref_words"],
"hyp_words": wer_result["hyp_words"],
},
processing_time_s=round(t_elapsed, 2),
rtf=round(rtf, 3),
avg_latency_ms=round(avg_lat, 1),
p95_latency_ms=round(p95_lat, 1),
n_transcription_calls=n_calls,
n_lines=len(state.speech_lines),
n_tokens=n_tokens,
timing_valid=state.timing_valid,
timing_monotonic=state.timing_monotonic,
peak_memory_mb=round(mem_delta, 1) if mem_delta > 0 else None,
hypothesis=hypothesis,
reference=sample.reference,
source=sample.source,
tags=list(sample.tags),
)

116
whisperlivekit/cascade_bridge.py Normal file
View File

@@ -0,0 +1,116 @@
"""
Bridge between WhisperLiveKit STT and IWSLT26 MT pipeline.
Converts streaming ASRToken output from SimulStreaming into the JSONL
format expected by the AlignAtt MT agent (iwslt26-sst).
Output format (one JSON per line):
{"text": "word or phrase", "emission_time": 1.234, "is_final": false, "speech_time": 1.0}
Where:
- text: the emitted word/phrase
- emission_time: wall-clock time when the word was emitted (for compute-aware eval)
- speech_time: timestamp in the audio (for compute-unaware eval)
- is_final: whether this is the last word of a segment/silence boundary
"""
import json
import time
from typing import List, TextIO
from whisperlivekit.timed_objects import ASRToken
class CascadeBridge:
"""Converts ASRToken stream to JSONL for the MT agent."""
def __init__(self, output_file: TextIO = None):
self.output_file = output_file
self.start_time = time.time()
self.entries: List[dict] = []
def emit_tokens(self, tokens: List[ASRToken], is_final: bool = False):
"""Emit a batch of tokens from the STT."""
wall_clock = time.time() - self.start_time
for i, token in enumerate(tokens):
entry = {
"text": token.text.strip(),
"emission_time": round(wall_clock, 3),
"speech_time": round(token.start, 3),
"is_final": is_final and (i == len(tokens) - 1),
}
self.entries.append(entry)
if self.output_file:
self.output_file.write(json.dumps(entry) + "\n")
self.output_file.flush()
def get_entries(self) -> List[dict]:
return self.entries
def get_text(self) -> str:
"""Get the full transcribed text."""
return " ".join(e["text"] for e in self.entries if e["text"])
def save(self, path: str):
"""Save all entries to a JSONL file."""
with open(path, "w") as f:
for entry in self.entries:
f.write(json.dumps(entry) + "\n")
def run_stt_to_jsonl(
audio_path: str,
output_path: str,
model_id: str = "Qwen/Qwen3-ASR-0.6B",
alignment_heads_path: str = None,
border_fraction: float = 0.20,
language: str = "en",
chunk_sec: float = 1.0,
):
"""Run STT on an audio file and save JSONL output for the MT agent.
This is the main entry point for the cascade: audio file → JSONL.
"""
import wave
import numpy as np
from whisperlivekit.qwen3_simul_kv import Qwen3SimulKVASR, Qwen3SimulKVOnlineProcessor
# Load audio
with wave.open(audio_path, 'r') as wf:
audio = np.frombuffer(
wf.readframes(wf.getnframes()), dtype=np.int16
).astype(np.float32) / 32768.0
# Initialize STT
asr = Qwen3SimulKVASR(
model_dir=model_id,
lan=language,
alignment_heads_path=alignment_heads_path,
border_fraction=border_fraction,
)
proc = Qwen3SimulKVOnlineProcessor(asr)
bridge = CascadeBridge()
# Stream audio in chunks
chunk_samples = int(chunk_sec * 16000)
offset = 0
stream_time = 0.0
while offset < len(audio):
chunk = audio[offset:offset + chunk_samples]
stream_time += len(chunk) / 16000
proc.insert_audio_chunk(chunk, stream_time)
words, _ = proc.process_iter(is_last=False)
if words:
bridge.emit_tokens(words, is_final=False)
offset += chunk_samples
# Final flush
final_words, _ = proc.finish()
if final_words:
bridge.emit_tokens(final_words, is_final=True)
# Save
bridge.save(output_path)
return bridge

462
whisperlivekit/cli.py
View File

@@ -57,6 +57,8 @@ BACKENDS = [
"install": "pip install faster-whisper",
"description": "CTranslate2-based Whisper (fast, CPU/CUDA)",
"policy": "localagreement",
"streaming": "chunk", # batch inference with LocalAgreement/SimulStreaming
"devices": ["cpu", "cuda"],
},
{
"id": "whisper",
@@ -65,6 +67,8 @@ BACKENDS = [
"install": "pip install openai-whisper",
"description": "Original OpenAI Whisper (PyTorch)",
"policy": "simulstreaming",
"streaming": "chunk",
"devices": ["cpu", "cuda"],
},
{
"id": "mlx-whisper",
@@ -74,21 +78,27 @@ BACKENDS = [
"description": "Apple Silicon native Whisper (MLX)",
"policy": "localagreement",
"platform": "darwin-arm64",
"streaming": "chunk",
"devices": ["mlx"],
},
{
"id": "voxtral-mlx",
"name": "Voxtral MLX",
"module": "mlx",
"install": "pip install whisperlivekit[voxtral-mlx]",
"description": "Mistral Voxtral Mini on Apple Silicon (MLX)",
"description": "Mistral Voxtral Mini on Apple Silicon (MLX, native streaming)",
"platform": "darwin-arm64",
"streaming": "native", # truly streaming (token-by-token)
"devices": ["mlx"],
},
{
"id": "voxtral",
"name": "Voxtral HF",
"module": "transformers",
"install": "pip install whisperlivekit[voxtral-hf]",
"description": "Mistral Voxtral Mini (HF Transformers, CUDA/CPU/MPS)",
"description": "Mistral Voxtral Mini (HF Transformers, native streaming)",
"streaming": "native",
"devices": ["cuda", "mps", "cpu"],
},
{
"id": "qwen3",
@@ -96,6 +106,18 @@ BACKENDS = [
"module": "qwen_asr",
"install": "pip install qwen-asr",
"description": "Qwen3-ASR with ForcedAligner timestamps",
"streaming": "chunk",
"devices": ["cuda", "mps", "cpu"],
},
{
"id": "qwen3-mlx",
"name": "Qwen3 MLX",
"module": "mlx_qwen3_asr",
"install": "pip install mlx-qwen3-asr",
"description": "Qwen3-ASR on Apple Silicon (MLX, native streaming)",
"platform": "darwin-arm64",
"streaming": "native",
"devices": ["mlx"],
},
{
"id": "openai-api",
@@ -103,6 +125,8 @@ BACKENDS = [
"module": "openai",
"install": "pip install openai",
"description": "Cloud-based transcription via OpenAI API",
"streaming": "cloud",
"devices": ["cloud"],
},
]
@@ -159,6 +183,31 @@ QWEN3_REPOS = {
}
QWEN3_ALIGNER_REPO = "Qwen/Qwen3-ForcedAligner-0.6B"
# Model catalog: metadata for display in `wlk models`
# params = approximate parameter count, disk = approximate download size
MODEL_CATALOG = [
# Whisper family (available across faster-whisper, mlx-whisper, whisper backends)
{"name": "tiny", "family": "whisper", "params": "39M", "disk": "75 MB", "languages": 99, "quality": "low", "speed": "fastest"},
{"name": "tiny.en", "family": "whisper", "params": "39M", "disk": "75 MB", "languages": 1, "quality": "low", "speed": "fastest"},
{"name": "base", "family": "whisper", "params": "74M", "disk": "142 MB", "languages": 99, "quality": "fair", "speed": "fast"},
{"name": "base.en", "family": "whisper", "params": "74M", "disk": "142 MB", "languages": 1, "quality": "fair", "speed": "fast"},
{"name": "small", "family": "whisper", "params": "244M", "disk": "466 MB", "languages": 99, "quality": "good", "speed": "medium"},
{"name": "small.en", "family": "whisper", "params": "244M", "disk": "466 MB", "languages": 1, "quality": "good", "speed": "medium"},
{"name": "medium", "family": "whisper", "params": "769M", "disk": "1.5 GB", "languages": 99, "quality": "great", "speed": "slow"},
{"name": "medium.en", "family": "whisper", "params": "769M", "disk": "1.5 GB", "languages": 1, "quality": "great", "speed": "slow"},
{"name": "large-v3", "family": "whisper", "params": "1.5B", "disk": "3.1 GB", "languages": 99, "quality": "best", "speed": "slowest"},
{"name": "large-v3-turbo", "family": "whisper", "params": "809M", "disk": "1.6 GB", "languages": 99, "quality": "great", "speed": "medium"},
# Voxtral (native streaming, single model)
{"name": "voxtral", "family": "voxtral", "params": "4B", "disk": "8.2 GB", "languages": 15, "quality": "great", "speed": "medium"},
{"name": "voxtral-mlx", "family": "voxtral", "params": "4B", "disk": "2.7 GB", "languages": 15, "quality": "great", "speed": "medium"},
# Qwen3 ASR
{"name": "qwen3:1.7b", "family": "qwen3", "params": "1.7B", "disk": "3.6 GB", "languages": 12, "quality": "good", "speed": "fast"},
{"name": "qwen3:0.6b", "family": "qwen3", "params": "0.6B", "disk": "1.4 GB", "languages": 12, "quality": "fair", "speed": "fastest"},
# Qwen3 MLX (native streaming on Apple Silicon)
{"name": "qwen3-mlx:1.7b", "family": "qwen3-mlx", "params": "1.7B", "disk": "1.8 GB", "languages": 12, "quality": "good", "speed": "fast"},
{"name": "qwen3-mlx:0.6b", "family": "qwen3-mlx", "params": "0.6B", "disk": "0.7 GB", "languages": 12, "quality": "fair", "speed": "fastest"},
]
def _check_platform(backend: dict) -> bool:
"""Check if backend is compatible with current platform."""
@@ -254,93 +303,131 @@ def print_banner(config, host: str, port: int, ssl: bool = False):
# `wlk models` subcommand
# ---------------------------------------------------------------------------
def cmd_models():
"""List available backends and their installation status."""
is_apple_silicon = platform.system() == "Darwin" and platform.machine() == "arm64"
def _model_is_downloaded(model_entry: dict, downloaded: dict) -> bool:
"""Check if a model catalog entry has been downloaded."""
name = model_entry["name"]
family = model_entry["family"]
print("\nAvailable backends:\n")
if family == "whisper":
# Check all whisper backends
repos = [
FASTER_WHISPER_REPOS.get(name),
MLX_WHISPER_REPOS.get(name),
f"openai/whisper-{name}",
]
return any(r in downloaded for r in repos if r)
elif name == "voxtral":
return VOXTRAL_HF_REPO in downloaded
elif name == "voxtral-mlx":
return VOXTRAL_MLX_REPO in downloaded
elif family == "qwen3":
size = name.split(":")[1] if ":" in name else "1.7b"
return QWEN3_REPOS.get(size, "") in downloaded
elif family == "qwen3-mlx":
size = name.split(":")[1] if ":" in name else "1.7b"
return QWEN3_REPOS.get(size, "") in downloaded
return False
def _best_backend_for_model(model_entry: dict) -> str:
"""Suggest the best available backend for a model."""
family = model_entry["family"]
is_apple = platform.system() == "Darwin" and platform.machine() == "arm64"
if family == "voxtral":
if "mlx" in model_entry["name"]:
return "voxtral-mlx"
return "voxtral"
elif family == "qwen3":
return "qwen3"
elif family == "qwen3-mlx":
return "qwen3-mlx"
elif family == "whisper":
if is_apple and _module_available("mlx_whisper"):
return "mlx-whisper"
if _module_available("faster_whisper"):
return "faster-whisper"
if _module_available("whisper"):
return "whisper"
# Suggest best installable
return "mlx-whisper" if is_apple else "faster-whisper"
return "auto"
def cmd_models():
"""List available models and backends (ollama-style)."""
is_apple_silicon = platform.system() == "Darwin" and platform.machine() == "arm64"
downloaded = _scan_downloaded_models()
# --- Installed backends ---
print("\n Backends:\n")
max_name = max(len(b["name"]) for b in BACKENDS)
for b in BACKENDS:
compatible = _check_platform(b)
installed = _is_installed(b)
streaming = b.get("streaming", "chunk")
stream_label = {"native": "streaming", "chunk": "chunked", "cloud": "cloud"}.get(streaming, streaming)
if installed:
status = "\033[32m installed\033[0m"
status = "\033[32m+\033[0m"
elif not compatible:
status = "\033[90m n/a (wrong platform)\033[0m"
status = "\033[90m-\033[0m"
else:
status = "\033[33m not installed\033[0m"
status = "\033[33m-\033[0m"
name_pad = b["name"].ljust(max_name)
print(f" {name_pad} [{status}] {b['description']}")
desc_short = b["description"]
print(f" {status} {name_pad} {desc_short} [{stream_label}]")
if not installed and compatible:
print(f" {''.ljust(max_name)} └─ {b['install']}")
print(f" {''.ljust(max_name)} \033[90m{b['install']}\033[0m")
# System info
# --- System info ---
print(f"\n Platform: {platform.system()} {platform.machine()}")
print(f" Python: {platform.python_version()}")
print(f" Accelerator: {_gpu_info()}")
print(f" ffmpeg: {'found' if _check_ffmpeg() else 'NOT FOUND (required)'}")
print(f" ffmpeg: {'found' if _check_ffmpeg() else '\033[31mNOT FOUND\033[0m (required)'}")
# --- Model catalog ---
print("\n Models:\n")
# Table header
hdr = f" {'NAME':<20} {'PARAMS':>7} {'SIZE':>8} {'QUALITY':<8} {'SPEED':<8} {'LANGS':>5} {'STATUS':<10}"
print(hdr)
print(f" {'─' * 20} {'─' * 7} {'─' * 8} {'─' * 8} {'─' * 8} {'─' * 5} {'─' * 10}")
for m in MODEL_CATALOG:
name = m["name"]
# Skip platform-incompatible models
if name == "voxtral-mlx" and not is_apple_silicon:
continue
if m["family"] == "qwen3-mlx" and not is_apple_silicon:
continue
is_dl = _model_is_downloaded(m, downloaded)
if is_dl:
status = "\033[32mpulled\033[0m "
else:
status = "\033[90mavailable\033[0m "
langs = str(m["languages"]) if m["languages"] < 99 else "99+"
print(
f" {name:<20} {m['params']:>7} {m['disk']:>8} "
f"{m['quality']:<8} {m['speed']:<8} {langs:>5} {status}"
)
# --- Quick start ---
print(f"\n Quick start:\n")
if is_apple_silicon:
print("\n Tip: On Apple Silicon, mlx-whisper and voxtral-mlx offer the best performance.")
# Scan for downloaded models
downloaded = _scan_downloaded_models()
print("\n Downloaded models:\n")
found_any = False
# Check Whisper-family models
all_repos = {
"faster-whisper": FASTER_WHISPER_REPOS,
"mlx-whisper": MLX_WHISPER_REPOS,
}
for backend_name, repos in all_repos.items():
for size, repo_id in repos.items():
if repo_id in downloaded:
found_any = True
print(f" \033[32m*\033[0m {backend_name}:{size} ({repo_id})")
# Check native whisper
for size in WHISPER_SIZES:
key = f"openai/whisper-{size}"
if key in downloaded:
found_any = True
print(f" \033[32m*\033[0m whisper:{size}")
# Check voxtral / qwen3
if VOXTRAL_HF_REPO in downloaded:
found_any = True
print(f" \033[32m*\033[0m voxtral ({VOXTRAL_HF_REPO})")
if VOXTRAL_MLX_REPO in downloaded:
found_any = True
print(f" \033[32m*\033[0m voxtral-mlx ({VOXTRAL_MLX_REPO})")
for qsize, repo_id in QWEN3_REPOS.items():
if repo_id in downloaded:
found_any = True
print(f" \033[32m*\033[0m qwen3:{qsize} ({repo_id})")
if QWEN3_ALIGNER_REPO in downloaded:
found_any = True
print(f" \033[32m*\033[0m qwen3-aligner ({QWEN3_ALIGNER_REPO})")
if not found_any:
print(" (none — models download automatically on first use, or use 'wlk pull')")
# Show pullable models
print("\n Available models (use 'wlk pull <name>'):\n")
print(" Whisper sizes: " + ", ".join(WHISPER_SIZES))
print(" Voxtral: voxtral, voxtral-mlx")
print(" Qwen3: qwen3:1.7b, qwen3:0.6b")
print()
print(" Examples:")
print(" wlk pull base # Download for best available backend")
print(" wlk pull faster-whisper:large-v3 # Specific backend + model")
print(" wlk pull voxtral # Voxtral HF model")
print(" wlk pull qwen3:1.7b # Qwen3-ASR 1.7B")
print(" wlk run voxtral-mlx # Best streaming on Apple Silicon")
print(" wlk run large-v3-turbo # Best quality/speed balance")
else:
print(" wlk run large-v3-turbo # Best quality/speed balance")
print(" wlk run voxtral # Native streaming (CUDA/CPU)")
print(" wlk pull base # Download smallest multilingual model")
print(" wlk transcribe audio.mp3 # Offline transcription")
print()
@@ -380,6 +467,18 @@ def _resolve_pull_target(spec: str):
targets.append(("voxtral-mlx", VOXTRAL_MLX_REPO, "Voxtral Mini (MLX)"))
return targets
# Handle qwen3-mlx (must check before generic qwen3)
if backend_part == "qwen3-mlx" or size_part.startswith("qwen3-mlx"):
qwen_size = size_part.split(":")[-1] if ":" in spec else "1.7b"
if qwen_size.startswith("qwen3"):
qwen_size = "1.7b" # default
repo = QWEN3_REPOS.get(qwen_size)
if not repo:
print(f" Unknown Qwen3 size: {qwen_size}. Available: {', '.join(QWEN3_REPOS.keys())}")
return []
targets.append(("qwen3-mlx", repo, f"Qwen3-ASR MLX {qwen_size}"))
return targets
# Handle qwen3
if backend_part == "qwen3" or size_part.startswith("qwen3"):
qwen_size = size_part.split(":")[-1] if ":" in spec else "1.7b"
@@ -436,7 +535,7 @@ def _resolve_pull_target(spec: str):
else:
print(f" Unknown model: {spec}")
print(f" Available sizes: {', '.join(WHISPER_SIZES)}")
print(" Other models: voxtral, voxtral-mlx, qwen3:1.7b, qwen3:0.6b")
print(" Other models: voxtral, voxtral-mlx, qwen3:1.7b, qwen3:0.6b, qwen3-mlx:1.7b, qwen3-mlx:0.6b")
return []
return targets
@@ -623,7 +722,11 @@ def _subtitle_timestamp(seconds: float, fmt: str) -> str:
# ---------------------------------------------------------------------------
def cmd_bench(args: list):
"""Benchmark the transcription pipeline on standard test audio.
"""Benchmark the transcription pipeline on public test audio.
Downloads samples from LibriSpeech, Multilingual LibriSpeech, FLEURS,
and AMI on first run. Supports multilingual benchmarking across all
available backends.
Usage: wlk bench [options]
"""
@@ -631,27 +734,48 @@ def cmd_bench(args: list):
parser = argparse.ArgumentParser(
prog="wlk bench",
description="Benchmark WhisperLiveKit on standard test audio.",
description="Benchmark WhisperLiveKit on public test audio.",
)
parser.add_argument("--backend", default="auto", help="ASR backend (default: auto)")
parser.add_argument("--model", default="base", dest="model_size", help="Model size (default: base)")
parser.add_argument("--language", "--lan", default="en", dest="lan", help="Language code (default: en)")
parser.add_argument("--samples", default="all", help="Sample name or 'all' (default: all)")
parser.add_argument("--json", default=None, dest="json_out", help="Export results to JSON file")
parser.add_argument("--verbose", "-v", action="store_true", help="Show detailed logs")
parser.add_argument("--backend", default="auto",
help="ASR backend (default: auto-detect)")
parser.add_argument("--model", default="base", dest="model_size",
help="Model size (default: base)")
parser.add_argument("--languages", "--lan", default=None,
help="Comma-separated language codes, or 'all' (default: en)")
parser.add_argument("--categories", default=None,
help="Comma-separated categories: clean,noisy,multilingual,meeting")
parser.add_argument("--quick", action="store_true",
help="Quick mode: small subset for smoke tests")
parser.add_argument("--json", default=None, dest="json_out",
help="Export full report to JSON file")
parser.add_argument("--transcriptions", action="store_true",
help="Show hypothesis vs reference for each sample")
parser.add_argument("--verbose", "-v", action="store_true",
help="Show detailed logs")
parsed = parser.parse_args(args)
# Parse languages
languages = None
if parsed.languages and parsed.languages != "all":
languages = [l.strip() for l in parsed.languages.split(",")]
elif parsed.languages is None:
languages = ["en"] # default to English only
categories = None
if parsed.categories:
categories = [c.strip() for c in parsed.categories.split(",")]
import asyncio
if not parsed.verbose:
asyncio.run(_run_bench_quiet(parsed))
else:
asyncio.run(_run_bench(parsed))
_suppress_logging()
asyncio.run(_run_bench_new(parsed, languages, categories))
async def _run_bench_quiet(parsed):
"""Run benchmark with suppressed logging."""
def _suppress_logging():
"""Suppress noisy logs during benchmark."""
import warnings
warnings.filterwarnings("ignore")
logging.root.setLevel(logging.ERROR)
@@ -659,130 +783,42 @@ async def _run_bench_quiet(parsed):
handler.setLevel(logging.ERROR)
for name in list(logging.Logger.manager.loggerDict.keys()):
logging.getLogger(name).setLevel(logging.ERROR)
await _run_bench(parsed)
async def _run_bench(parsed):
"""Run the benchmark."""
import json as json_module
import time
async def _run_bench_new(parsed, languages, categories):
"""Run the benchmark using the new benchmark module."""
from whisperlivekit.benchmark.report import print_report, print_transcriptions, write_json
from whisperlivekit.benchmark.runner import BenchmarkRunner
from whisperlivekit.metrics import compute_wer
from whisperlivekit.test_data import get_sample, get_samples
from whisperlivekit.test_harness import TestHarness
def on_progress(name, i, total):
if name == "done":
print(f"\r [{total}/{total}] Done.{' ' * 30}", file=sys.stderr)
else:
print(f"\r [{i + 1}/{total}] {name}...{' ' * 20}",
end="", file=sys.stderr, flush=True)
# Determine samples to run
if parsed.samples == "all":
print(" Downloading test samples (first run only)...", file=sys.stderr)
samples = get_samples()
# Filter to matching language
samples = [s for s in samples if s.language == parsed.lan]
if not samples:
# Fall back to all samples if none match the language
samples = get_samples()
else:
samples = [get_sample(parsed.samples)]
runner = BenchmarkRunner(
backend=parsed.backend,
model_size=parsed.model_size,
languages=languages,
categories=categories,
quick=parsed.quick,
on_progress=on_progress,
)
backend_label = parsed.backend
if backend_label == "auto":
backend_label = "auto-detect"
print(f"\n Downloading benchmark samples (cached after first run)...",
file=sys.stderr)
print(file=sys.stderr)
print(" WhisperLiveKit Benchmark", file=sys.stderr)
print(f" Backend: {backend_label} | Model: {parsed.model_size} | Language: {parsed.lan}", file=sys.stderr)
print(f" Samples: {len(samples)}", file=sys.stderr)
print(f" {'─' * 70}", file=sys.stderr)
report = await runner.run()
results = []
print_report(report)
kwargs = {
"model_size": parsed.model_size,
"lan": parsed.lan,
"pcm_input": True,
}
if parsed.backend != "auto":
kwargs["backend"] = parsed.backend
if parsed.transcriptions:
print_transcriptions(report)
for sample in samples:
print(f"\n {sample.name} ({sample.duration:.1f}s, {sample.language})", file=sys.stderr)
t_start = time.perf_counter()
async with TestHarness(**kwargs) as h:
await h.feed(sample.path, speed=0)
await h.drain(5.0)
state = await h.finish(timeout=120)
t_elapsed = time.perf_counter() - t_start
rtf = t_elapsed / sample.duration if sample.duration > 0 else 0
# Compute WER
hypothesis = state.committed_text or state.text
wer_result = compute_wer(sample.reference, hypothesis)
n_lines = len(state.speech_lines)
result_entry = {
"sample": sample.name,
"duration_s": round(sample.duration, 2),
"processing_time_s": round(t_elapsed, 2),
"rtf": round(rtf, 3),
"wer": round(wer_result["wer"], 4),
"wer_details": {
"substitutions": wer_result["substitutions"],
"insertions": wer_result["insertions"],
"deletions": wer_result["deletions"],
"ref_words": wer_result["ref_words"],
"hyp_words": wer_result["hyp_words"],
},
"n_lines": n_lines,
"transcription": hypothesis,
}
results.append(result_entry)
# Print per-sample result
wer_pct = wer_result["wer"] * 100
wer_color = "\033[32m" if wer_pct < 15 else "\033[33m" if wer_pct < 30 else "\033[31m"
rtf_color = "\033[32m" if rtf < 0.5 else "\033[33m" if rtf < 1.0 else "\033[31m"
print(f" WER: {wer_color}{wer_pct:5.1f}%\033[0m "
f"(S:{wer_result['substitutions']} I:{wer_result['insertions']} D:{wer_result['deletions']})",
file=sys.stderr)
print(f" RTF: {rtf_color}{rtf:.3f}x\033[0m "
f"({t_elapsed:.1f}s for {sample.duration:.1f}s audio)",
file=sys.stderr)
print(f" Lines: {n_lines}",
file=sys.stderr)
# Summary
if len(results) > 1:
avg_wer = sum(r["wer"] for r in results) / len(results)
avg_rtf = sum(r["rtf"] for r in results) / len(results)
total_audio = sum(r["duration_s"] for r in results)
total_proc = sum(r["processing_time_s"] for r in results)
print(f"\n {'─' * 70}", file=sys.stderr)
print(f" Summary ({len(results)} samples, {total_audio:.1f}s total audio)", file=sys.stderr)
wer_color = "\033[32m" if avg_wer * 100 < 15 else "\033[33m" if avg_wer * 100 < 30 else "\033[31m"
rtf_color = "\033[32m" if avg_rtf < 0.5 else "\033[33m" if avg_rtf < 1.0 else "\033[31m"
print(f" Avg WER: {wer_color}{avg_wer * 100:5.1f}%\033[0m", file=sys.stderr)
print(f" Avg RTF: {rtf_color}{avg_rtf:.3f}x\033[0m "
f"({total_proc:.1f}s for {total_audio:.1f}s audio)", file=sys.stderr)
print(file=sys.stderr)
# JSON export
if parsed.json_out:
export = {
"backend": parsed.backend,
"model_size": parsed.model_size,
"language": parsed.lan,
"accelerator": _gpu_info(),
"results": results,
}
with open(parsed.json_out, "w") as f:
json_module.dump(export, f, indent=2)
print(f" Results exported to: {parsed.json_out}", file=sys.stderr)
write_json(report, parsed.json_out)
print(f" Results exported to: {parsed.json_out}\n", file=sys.stderr)
# ---------------------------------------------------------------------------
@@ -982,6 +1018,9 @@ def _resolve_run_spec(spec: str):
if spec == "voxtral-mlx":
return "voxtral-mlx", None
if spec == "qwen3-mlx":
return "qwen3-mlx", None
if spec in WHISPER_SIZES:
return None, spec
@@ -1010,6 +1049,23 @@ def cmd_run(args: list):
if parsed.model:
backend_flag, model_flag = _resolve_run_spec(parsed.model)
# Show what we resolved
catalog_match = next(
(m for m in MODEL_CATALOG if m["name"] == parsed.model),
None,
)
if catalog_match:
print(
f"\n Model: {catalog_match['name']} "
f"({catalog_match['params']} params, {catalog_match['disk']})",
file=sys.stderr,
)
if backend_flag:
print(f" Backend: {backend_flag}", file=sys.stderr)
else:
best = _best_backend_for_model(catalog_match)
print(f" Backend: {best} (auto-detected)", file=sys.stderr)
# Auto-pull if needed
downloaded = _scan_downloaded_models()
targets = _resolve_pull_target(parsed.model)
@@ -1198,9 +1254,9 @@ def _probe_backend_state(processor) -> dict:
info["n_audio_tokens_fed"] = transcription._n_audio_tokens_fed
info["n_text_tokens_received"] = transcription._n_text_tokens_received
info["n_committed_words"] = transcription._n_committed_words
info["pending_audio_samples"] = len(transcription._pending_audio)
info["pending_audio_samples"] = transcription._pending_len
with transcription._text_lock:
info["accumulated_text"] = transcription._accumulated_text
info["accumulated_text"] = transcription._get_accumulated_text()
if transcription._generate_error:
info["generate_error"] = str(transcription._generate_error)
# Audio queue depth
@@ -1210,6 +1266,12 @@ def _probe_backend_state(processor) -> dict:
elif hasattr(transcription, "_mlx_processor"):
info["backend_type"] = "voxtral-mlx"
# Qwen3 MLX specifics
elif hasattr(transcription, "_session") and hasattr(transcription, "_state"):
info["backend_type"] = "qwen3-mlx"
info["samples_fed"] = getattr(transcription, "_samples_fed", 0)
info["committed_words"] = getattr(transcription, "_n_committed_words", 0)
# SimulStreaming specifics
elif hasattr(transcription, "prev_output"):
info["backend_type"] = "simulstreaming"

4
whisperlivekit/config.py
View File

@@ -72,6 +72,10 @@ class WhisperLiveKitConfig:
nllb_backend: str = "transformers"
nllb_size: str = "600M"
# vLLM Realtime backend
vllm_url: str = "ws://localhost:8000/v1/realtime"
vllm_model: str = ""
def __post_init__(self):
# .en model suffix forces English
if self.model_size and self.model_size.endswith(".en"):

45
whisperlivekit/core.py
View File

@@ -102,7 +102,16 @@ class TranscriptionEngine:
}
if config.transcription:
if config.backend == "voxtral-mlx":
if config.backend == "vllm-realtime":
from whisperlivekit.vllm_realtime import VLLMRealtimeASR
self.tokenizer = None
self.asr = VLLMRealtimeASR(
vllm_url=config.vllm_url,
model_name=config.vllm_model or "Qwen/Qwen3-ASR-1.7B",
lan=config.lan,
)
logger.info("Using vLLM Realtime streaming backend at %s", config.vllm_url)
elif config.backend == "voxtral-mlx":
from whisperlivekit.voxtral_mlx_asr import VoxtralMLXASR
self.tokenizer = None
self.asr = VoxtralMLXASR(**transcription_common_params)
@@ -112,6 +121,28 @@ class TranscriptionEngine:
self.tokenizer = None
self.asr = VoxtralHFStreamingASR(**transcription_common_params)
logger.info("Using Voxtral HF Transformers streaming backend")
elif config.backend == "qwen3-mlx":
from whisperlivekit.qwen3_mlx_asr import Qwen3MLXASR
self.tokenizer = None
self.asr = Qwen3MLXASR(**transcription_common_params)
logger.info("Using Qwen3 MLX native backend")
elif config.backend == "qwen3-simul-kv":
from whisperlivekit.qwen3_simul_kv import Qwen3SimulKVASR
self.tokenizer = None
self.asr = Qwen3SimulKVASR(
**transcription_common_params,
alignment_heads_path=config.custom_alignment_heads,
border_fraction=getattr(config, 'border_fraction', 0.25),
)
logger.info("Using Qwen3-ASR backend with SimulStreaming+KV policy")
elif config.backend == "qwen3-simul":
from whisperlivekit.qwen3_simul import Qwen3SimulStreamingASR
self.tokenizer = None
self.asr = Qwen3SimulStreamingASR(
**transcription_common_params,
alignment_heads_path=config.custom_alignment_heads,
)
logger.info("Using Qwen3-ASR backend with SimulStreaming policy")
elif config.backend == "qwen3":
from whisperlivekit.qwen3_asr import Qwen3ASR
self.asr = Qwen3ASR(**transcription_common_params)
@@ -210,6 +241,18 @@ def online_factory(args, asr, language=None):
asr = SessionASRProxy(asr, language)
backend = getattr(args, 'backend', None)
if backend == "vllm-realtime":
from whisperlivekit.vllm_realtime import VLLMRealtimeOnlineProcessor
return VLLMRealtimeOnlineProcessor(asr)
if backend == "qwen3-simul-kv":
from whisperlivekit.qwen3_simul_kv import Qwen3SimulKVOnlineProcessor
return Qwen3SimulKVOnlineProcessor(asr)
if backend == "qwen3-mlx":
from whisperlivekit.qwen3_mlx_asr import Qwen3MLXOnlineProcessor
return Qwen3MLXOnlineProcessor(asr)
if backend == "qwen3-simul":
from whisperlivekit.qwen3_simul import Qwen3SimulStreamingOnlineProcessor
return Qwen3SimulStreamingOnlineProcessor(asr)
if backend == "voxtral-mlx":
from whisperlivekit.voxtral_mlx_asr import VoxtralMLXOnlineProcessor
return VoxtralMLXOnlineProcessor(asr)

20
whisperlivekit/parse_args.py
View File

@@ -147,8 +147,8 @@ def parse_args():
"--backend",
type=str,
default="auto",
choices=["auto", "mlx-whisper", "faster-whisper", "whisper", "openai-api", "voxtral", "voxtral-mlx", "qwen3"],
help="Select the ASR backend implementation (auto: prefer MLX on macOS, otherwise Faster-Whisper, else Whisper). Use 'voxtral' for HF Transformers Voxtral (CUDA/CPU/MPS). Use 'voxtral-mlx' for native MLX Voxtral on Apple Silicon. Use 'qwen3' for Qwen3-ASR.",
choices=["auto", "mlx-whisper", "faster-whisper", "whisper", "openai-api", "voxtral", "voxtral-mlx", "qwen3", "qwen3-mlx", "qwen3-simul", "vllm-realtime"],
help="Select the ASR backend implementation. Use 'qwen3' for Qwen3-ASR with LocalAgreement. Use 'qwen3-mlx' for Qwen3-ASR on Apple Silicon (MLX). Use 'qwen3-simul' for Qwen3-ASR with SimulStreaming (requires alignment heads). Use 'vllm-realtime' for vLLM Realtime WebSocket.",
)
parser.add_argument(
"--no-vac",
@@ -196,6 +196,22 @@ def parse_args():
default=False,
help="If set, raw PCM (s16le) data is expected as input and FFmpeg will be bypassed. Frontend will use AudioWorklet instead of MediaRecorder."
)
# vLLM Realtime backend arguments
parser.add_argument(
"--vllm-url",
type=str,
default="ws://localhost:8000/v1/realtime",
dest="vllm_url",
help="URL of the vLLM realtime WebSocket endpoint.",
)
parser.add_argument(
"--vllm-model",
type=str,
default="",
dest="vllm_model",
help="Model name to use with vLLM (e.g. Qwen/Qwen3-ASR-1.7B).",
)
# SimulStreaming-specific arguments
simulstreaming_group = parser.add_argument_group('SimulStreaming arguments (only used with --backend simulstreaming)')

98
whisperlivekit/qwen3_asr.py
View File

@@ -1,4 +1,5 @@
import logging
import re
import sys
from typing import List, Optional
@@ -11,12 +12,10 @@ logger = logging.getLogger(__name__)
def _patch_transformers_compat():
"""Patch transformers for qwen_asr compatibility.
"""Patch transformers for qwen_asr 0.0.6 + transformers >= 5.3 compatibility."""
import torch
qwen_asr imports ``check_model_inputs`` from ``transformers.utils.generic``,
but this decorator hasn't been released yet in any public transformers
version. We inject a no-op stub so the import succeeds.
"""
# 1. check_model_inputs was removed
try:
import transformers.utils.generic as _g
if not hasattr(_g, "check_model_inputs"):
@@ -28,6 +27,63 @@ def _patch_transformers_compat():
except ImportError:
pass
# 2. 'default' rope type was removed from ROPE_INIT_FUNCTIONS
try:
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
if "default" not in ROPE_INIT_FUNCTIONS:
def _compute_default_rope_parameters(config=None, device=None, seq_len=None, **kwargs):
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
partial = getattr(config, "partial_rotary_factor", 1.0)
dim = int(head_dim * partial)
base = config.rope_theta
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
return inv_freq, 1.0
ROPE_INIT_FUNCTIONS["default"] = _compute_default_rope_parameters
except ImportError:
pass
# 3. pad_token_id missing on thinker config
try:
from qwen_asr.core.transformers_backend.configuration_qwen3_asr import (
Qwen3ASRThinkerConfig,
)
if not hasattr(Qwen3ASRThinkerConfig, "pad_token_id"):
Qwen3ASRThinkerConfig.pad_token_id = None
except ImportError:
pass
# 4. fix_mistral_regex kwarg not accepted by newer transformers
try:
from transformers.models.auto import processing_auto
_orig_ap_from_pretrained = processing_auto.AutoProcessor.from_pretrained.__func__
@classmethod
def _patched_ap_from_pretrained(cls, *args, **kwargs):
kwargs.pop("fix_mistral_regex", None)
return _orig_ap_from_pretrained(cls, *args, **kwargs)
processing_auto.AutoProcessor.from_pretrained = _patched_ap_from_pretrained
except Exception:
pass
# 5. compute_default_rope_parameters missing on RotaryEmbedding
try:
from qwen_asr.core.transformers_backend.modeling_qwen3_asr import (
Qwen3ASRThinkerTextRotaryEmbedding,
)
if not hasattr(Qwen3ASRThinkerTextRotaryEmbedding, "compute_default_rope_parameters"):
@staticmethod
def _rope_params(config=None, device=None, seq_len=None, **kwargs):
head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
partial = getattr(config, "partial_rotary_factor", 1.0)
dim = int(head_dim * partial)
base = config.rope_theta
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))
return inv_freq, 1.0
Qwen3ASRThinkerTextRotaryEmbedding.compute_default_rope_parameters = _rope_params
except ImportError:
pass
_patch_transformers_compat()
@@ -62,6 +118,9 @@ QWEN3_MODEL_MAPPING = {
}
_PUNCTUATION_ENDS = set(".!?。！？；;")
# Qwen3 raw output starts with "language <Name>" metadata before <asr_text> tag.
# When the tag is missing (silence/noise), this metadata leaks as transcription text.
_GARBAGE_RE = re.compile(r"^language\s+\S+$", re.IGNORECASE)
class Qwen3ASR(ASRBase):
@@ -88,8 +147,12 @@ class Qwen3ASR(ASRBase):
else:
model_id = "Qwen/Qwen3-ASR-1.7B"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
device = "cuda:0" if torch.cuda.is_available() else "cpu"
if torch.cuda.is_available():
dtype, device = torch.bfloat16, "cuda:0"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
dtype, device = torch.float32, "mps"
else:
dtype, device = torch.float32, "cpu"
logger.info(f"Loading Qwen3-ASR: {model_id} ({dtype}, {device})")
model = Qwen3ASRModel.from_pretrained(
@@ -126,17 +189,32 @@ class Qwen3ASR(ASRBase):
result = results[0]
# Stash audio length for timestamp estimation fallback
result._audio_duration = len(audio) / 16000
logger.info(
"Qwen3 result: language=%r text=%r ts=%s",
result.language, result.text[:80] if result.text else "",
bool(result.time_stamps),
)
return result
@staticmethod
def _detected_language(result) -> Optional[str]:
"""Extract Whisper-style language code from Qwen3 result."""
lang = getattr(result, 'language', None)
if lang:
return QWEN3_TO_WHISPER_LANGUAGE.get(lang, lang.lower())
return None
if not lang or lang.lower() == "none":
return None
# merge_languages may return comma-separated; take the first
first = lang.split(",")[0].strip()
if not first or first.lower() == "none":
return None
return QWEN3_TO_WHISPER_LANGUAGE.get(first, first.lower())
def ts_words(self, result) -> List[ASRToken]:
# Filter garbage model output (e.g. "language None" for silence/noise)
text = (result.text or "").strip()
if not text or _GARBAGE_RE.match(text):
if text:
logger.info("Filtered garbage Qwen3 output: %r", text)
return []
detected = self._detected_language(result)
if result.time_stamps:
tokens = []

392
whisperlivekit/qwen3_mlx_asr.py Normal file
View File

@@ -0,0 +1,392 @@
"""
MLX-accelerated Qwen3-ASR backend for WhisperLiveKit.
Provides ``Qwen3MLXASR`` (model holder) and ``Qwen3MLXOnlineProcessor``
(batch-based processor) that plug into WhisperLiveKit's audio processing
pipeline via ``insert_audio_chunk`` / ``process_iter`` / ``get_buffer`` etc.
Uses the ``mlx-qwen3-asr`` package for fast Qwen3 inference on Apple Silicon.
The batch ``session.transcribe()`` API is called on the full accumulated audio
buffer, and LocalAgreement-style diffing (HypothesisBuffer) commits stable
words across consecutive inferences.
"""
import logging
import sys
import time
from typing import List, Tuple
import numpy as np
from whisperlivekit.timed_objects import ASRToken, Transcript
logger = logging.getLogger(__name__)
# Whisper language codes -> Qwen3 canonical language names
# (duplicated from qwen3_asr.py to avoid importing torch at module level)
WHISPER_TO_QWEN3_LANGUAGE = {
"zh": "Chinese", "en": "English", "yue": "Cantonese",
"ar": "Arabic", "de": "German", "fr": "French", "es": "Spanish",
"pt": "Portuguese", "id": "Indonesian", "it": "Italian",
"ko": "Korean", "ru": "Russian", "th": "Thai", "vi": "Vietnamese",
"ja": "Japanese", "tr": "Turkish", "hi": "Hindi", "ms": "Malay",
"nl": "Dutch", "sv": "Swedish", "da": "Danish", "fi": "Finnish",
"pl": "Polish", "cs": "Czech", "fa": "Persian",
"el": "Greek", "hu": "Hungarian", "mk": "Macedonian", "ro": "Romanian",
}
# Model size aliases -> HuggingFace model IDs
QWEN3_MLX_MODEL_MAPPING = {
"base": "Qwen/Qwen3-ASR-0.6B",
"tiny": "Qwen/Qwen3-ASR-0.6B",
"small": "Qwen/Qwen3-ASR-0.6B",
"large": "Qwen/Qwen3-ASR-1.7B",
"medium": "Qwen/Qwen3-ASR-1.7B",
"large-v3": "Qwen/Qwen3-ASR-1.7B",
"qwen3-asr-1.7b": "Qwen/Qwen3-ASR-1.7B",
"qwen3-asr-0.6b": "Qwen/Qwen3-ASR-0.6B",
"qwen3-1.7b": "Qwen/Qwen3-ASR-1.7B",
"qwen3-0.6b": "Qwen/Qwen3-ASR-0.6B",
"1.7b": "Qwen/Qwen3-ASR-1.7B",
"0.6b": "Qwen/Qwen3-ASR-0.6B",
}
# ---------------------------------------------------------------------------
# Model holder
# ---------------------------------------------------------------------------
class Qwen3MLXASR:
"""Lightweight model holder -- loads the mlx-qwen3-asr model once and
keeps it alive for the lifetime of the server."""
sep = ""
SAMPLING_RATE = 16_000
def __init__(self, logfile=sys.stderr, **kwargs):
import mlx.core as mx
import mlx_qwen3_asr
self.logfile = logfile
self.transcribe_kargs = {}
lan = kwargs.get("lan", "auto")
self.original_language = None if lan == "auto" else lan
# Resolve model ID from size aliases or explicit path
model_path = kwargs.get("model_dir") or kwargs.get("model_path")
if not model_path:
model_size = kwargs.get("model_size", "")
if model_size and ("/" in model_size or model_size.startswith(".")):
model_path = model_size
else:
model_path = QWEN3_MLX_MODEL_MAPPING.get(
(model_size or "base").lower(), "Qwen/Qwen3-ASR-0.6B"
)
t0 = time.time()
logger.info("Loading Qwen3 MLX model '%s' ...", model_path)
self.session = mlx_qwen3_asr.Session(model_path, dtype=mx.float16)
logger.info("Qwen3 MLX model loaded in %.2fs", time.time() - t0)
self.backend_choice = "qwen3-mlx"
self.tokenizer = None
def transcribe(self, audio):
pass # all work happens in the online processor
# ---------------------------------------------------------------------------
# Online processor
# ---------------------------------------------------------------------------
class Qwen3MLXOnlineProcessor:
"""Batch-based processor that accumulates audio and periodically calls
``session.transcribe()`` on the full buffer.
Uses LocalAgreement-style diffing (HypothesisBuffer) to commit stable
words across consecutive inferences, exactly like the PyTorch Qwen3
backend with ``OnlineASRProcessor``.
Lifecycle (called by ``AudioProcessor.transcription_processor``):
insert_audio_chunk(pcm, time) -> process_iter() -> get_buffer()
... repeat ...
start_silence() / end_silence()
finish()
"""
SAMPLING_RATE = 16_000
def __init__(self, asr: Qwen3MLXASR, logfile=sys.stderr):
self.asr = asr
self.logfile = logfile
self.end = 0.0
self._session = asr.session
lan = asr.original_language
self._language = WHISPER_TO_QWEN3_LANGUAGE.get(lan, "English") if lan else None
# Audio accumulation
self.audio_buffer = np.array([], dtype=np.float32)
self._buffer_time_offset: float = 0.0 # absolute time of audio_buffer[0]
# Throttle: minimum new audio (in samples) before re-running inference
self._min_new_samples: int = int(1.0 * self.SAMPLING_RATE) # 1 second
self._samples_since_last_inference: int = 0
# Buffer trimming — keep buffer short for fast re-transcription.
# The model produces ~0.2x RTF, so 15s buffer = ~3s per call.
self._max_buffer_sec: float = 15.0
self._trim_sec: float = 10.0 # keep this many seconds after trimming
# HypothesisBuffer for LocalAgreement diffing
self._committed: List[ASRToken] = []
self._prev_tokens: List[ASRToken] = [] # previous hypothesis (buffer role)
self._last_committed_time: float = 0.0
# Global time tracking
self._global_time_offset: float = 0.0 # extra offset from silences
# -- audio ingestion --
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float):
self.end = audio_stream_end_time
self.audio_buffer = np.append(self.audio_buffer, audio)
self._samples_since_last_inference += len(audio)
# -- batch transcription --
def _transcribe_buffer(self) -> List[ASRToken]:
"""Run batch transcription on the full audio buffer and return tokens."""
if len(self.audio_buffer) < 400: # too short for meaningful transcription
return []
t0 = time.time()
try:
result = self._session.transcribe(
self.audio_buffer,
language=self._language,
return_timestamps=True,
)
except Exception as e:
logger.warning("[qwen3-mlx] transcribe error: %s", e, exc_info=True)
return []
dur = time.time() - t0
audio_dur = len(self.audio_buffer) / self.SAMPLING_RATE
logger.debug(
"[qwen3-mlx] transcribed %.1fs audio in %.2fs (%.2fx RTF)",
audio_dur, dur, dur / max(audio_dur, 0.01),
)
text = (result.text or "").strip()
if not text:
return []
# Build tokens from segments (word-level timestamps)
tokens: List[ASRToken] = []
if result.segments:
for i, seg in enumerate(result.segments):
word = seg["text"]
start = self._buffer_time_offset + seg["start"]
end = self._buffer_time_offset + seg["end"]
label = word if i == 0 else " " + word
tokens.append(ASRToken(start=start, end=end, text=label))
else:
# Fallback: estimate timestamps from word count
words = text.split()
step = audio_dur / max(len(words), 1)
for i, w in enumerate(words):
t_start = self._buffer_time_offset + i * step
t_end = self._buffer_time_offset + (i + 1) * step
label = w if i == 0 else " " + w
tokens.append(ASRToken(start=t_start, end=t_end, text=label))
return tokens
def _local_agreement(self, new_tokens: List[ASRToken]) -> List[ASRToken]:
"""LocalAgreement diffing: commit the longest common prefix between
the previous hypothesis (``self._prev_tokens``) and the new tokens.
Before comparing, strips tokens that correspond to already-committed
audio (i.e., tokens whose start time is before ``_last_committed_time``).
Also deduplicates boundary tokens (ngram matching) to avoid re-committing
the tail of the previous committed output.
Returns the newly committed tokens.
"""
# Step 1: Only keep tokens that are roughly "new" (after last committed time)
fresh_tokens = [
t for t in new_tokens
if t.start > self._last_committed_time - 0.1
]
# Step 2: Remove duplicates at the boundary with committed tokens
# (like HypothesisBuffer.insert's ngram dedup)
if fresh_tokens and self._committed:
max_ngram = min(len(self._committed), len(fresh_tokens), 5)
for n in range(1, max_ngram + 1):
committed_ngram = " ".join(
t.text.strip() for t in self._committed[-n:]
)
fresh_ngram = " ".join(
t.text.strip() for t in fresh_tokens[:n]
)
if committed_ngram == fresh_ngram:
fresh_tokens = fresh_tokens[n:]
break
# Step 3: LocalAgreement -- longest common prefix between prev and fresh
committed: List[ASRToken] = []
prev = self._prev_tokens
i = 0
j = 0
while i < len(fresh_tokens) and j < len(prev):
if fresh_tokens[i].text.strip() == prev[j].text.strip():
# Agreement: commit this token (use the new token's timestamps)
committed.append(fresh_tokens[i])
i += 1
j += 1
else:
break
# The remaining fresh tokens become the new "previous hypothesis"
self._prev_tokens = fresh_tokens[i:] if i < len(fresh_tokens) else []
return committed
def _trim_buffer_if_needed(self):
"""Trim the audio buffer if it exceeds max_buffer_sec.
Keeps the last ``_trim_sec`` seconds of audio. Also adjusts
committed token tracking and buffer_time_offset.
"""
buffer_dur = len(self.audio_buffer) / self.SAMPLING_RATE
if buffer_dur <= self._max_buffer_sec:
return
keep_sec = self._trim_sec
keep_samples = int(keep_sec * self.SAMPLING_RATE)
cut_samples = len(self.audio_buffer) - keep_samples
if cut_samples <= 0:
return
cut_sec = cut_samples / self.SAMPLING_RATE
self.audio_buffer = self.audio_buffer[cut_samples:]
self._buffer_time_offset += cut_sec
# Remove committed tokens that are before the new buffer start
self._committed = [
t for t in self._committed if t.end > self._buffer_time_offset
]
logger.debug(
"[qwen3-mlx] trimmed buffer: cut %.1fs, new offset %.1f, buffer %.1fs",
cut_sec, self._buffer_time_offset, len(self.audio_buffer) / self.SAMPLING_RATE,
)
# -- interface methods --
def process_iter(self, is_last=False) -> Tuple[List[ASRToken], float]:
"""Process the current audio buffer.
Throttles inference to at least 1s of new audio between calls.
Returns (newly_committed_tokens, audio_processed_upto_time).
"""
try:
# Throttle: skip if not enough new audio since last inference
if (not is_last
and self._samples_since_last_inference < self._min_new_samples):
return [], self.end
self._samples_since_last_inference = 0
# Trim buffer if too long
self._trim_buffer_if_needed()
# Run batch transcription
new_tokens = self._transcribe_buffer()
# LocalAgreement diffing
committed = self._local_agreement(new_tokens)
if committed:
self._committed.extend(committed)
self._last_committed_time = committed[-1].end
return committed, self.end
except Exception as e:
logger.warning("[qwen3-mlx] process_iter error: %s", e, exc_info=True)
return [], self.end
def get_buffer(self) -> Transcript:
"""Return the unconfirmed text (the tail of the last hypothesis
that was not committed by LocalAgreement)."""
if not self._prev_tokens:
return Transcript(start=None, end=None, text="")
text = "".join(t.text for t in self._prev_tokens)
start = self._prev_tokens[0].start
end = self._prev_tokens[-1].end
return Transcript(start=start, end=end, text=text)
def _flush_all(self) -> List[ASRToken]:
"""Force a final transcription and commit all remaining words."""
# Run one last transcription on the full buffer
self._samples_since_last_inference = self._min_new_samples # bypass throttle
new_tokens = self._transcribe_buffer()
# Commit everything: first the agreed prefix, then the remainder
committed = self._local_agreement(new_tokens)
# Also commit any remaining buffer tokens
remaining = self._prev_tokens
self._prev_tokens = []
all_new = committed + remaining
if all_new:
self._committed.extend(all_new)
self._last_committed_time = all_new[-1].end
return all_new
def _reset_for_new_utterance(self):
"""Reset buffers for a new utterance, preserving time continuity."""
new_offset = self._buffer_time_offset + len(self.audio_buffer) / self.SAMPLING_RATE
saved_end = self.end
self.audio_buffer = np.array([], dtype=np.float32)
self._buffer_time_offset = new_offset
self._samples_since_last_inference = 0
self._committed = []
self._prev_tokens = []
self.end = saved_end
def start_silence(self) -> Tuple[List[ASRToken], float]:
"""Flush pending words when silence starts.
Unlike other backends, does NOT reset the audio buffer — the model
produces better results re-transcribing the full accumulated audio.
Buffer trimming at 30s handles memory naturally.
"""
words = self._flush_all()
logger.info("[qwen3-mlx] start_silence: flushed %d words", len(words))
return words, self.end
def end_silence(self, silence_duration: float, offset: float):
self._global_time_offset += silence_duration
self.end += silence_duration
def new_speaker(self, change_speaker):
self.start_silence()
def warmup(self, audio, init_prompt=""):
pass
def finish(self) -> Tuple[List[ASRToken], float]:
words = self._flush_all()
logger.info("[qwen3-mlx] finish: flushed %d words", len(words))
return words, self.end

1190
whisperlivekit/qwen3_simul.py Normal file
View File

File diff suppressed because it is too large Load Diff

791
whisperlivekit/qwen3_simul_kv.py Normal file
View File

@@ -0,0 +1,791 @@
"""
Qwen3-ASR SimulStreaming with KV cache reuse.
This is an optimized version of qwen3_simul.py that reuses the KV cache
across inference calls, avoiding redundant prefill of prompt + old audio.
Architecture:
1. First call: full prefill (prompt + audio tokens), greedy decode with
alignment-head stopping, save KV cache + generated tokens
2. Subsequent calls: invalidate KV for old audio suffix, prefill only
new audio tokens, continue decoding from saved state
3. Audio encoder caching: reuse embeddings for stable attention windows
This gives ~3-5x speedup over the original generate()-based approach.
"""
import json
import logging
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional, Tuple
import numpy as np
import torch
from transformers import DynamicCache
from whisperlivekit.timed_objects import ASRToken, ChangeSpeaker, Transcript
logger = logging.getLogger(__name__)
SAMPLE_RATE = 16000
@dataclass
class Qwen3SimulKVConfig:
"""Configuration for Qwen3 SimulStreaming with KV cache."""
model_id: str = "Qwen/Qwen3-ASR-1.7B"
alignment_heads_path: Optional[str] = None
language: str = "auto"
border_fraction: float = 0.20
rewind_fraction: float = 0.12
audio_min_len: float = 0.5
audio_max_len: float = 30.0
max_context_tokens: int = 20
init_prompt: Optional[str] = None
max_alignment_heads: int = 10
min_new_seconds: float = 2.0 # minimum new audio before running inference
@dataclass
class _AudioEmbedCache:
"""Cache for audio encoder outputs."""
encoded_samples: int = 0
embeddings: Optional[torch.Tensor] = None
encoded_mel_frames: int = 0
stable_tokens: int = 0
def reset(self):
self.encoded_samples = 0
self.embeddings = None
self.encoded_mel_frames = 0
self.stable_tokens = 0
@dataclass
class Qwen3SimulKVState:
"""Per-session mutable state with KV cache."""
# Audio
audio_buffer: np.ndarray = field(
default_factory=lambda: np.array([], dtype=np.float32)
)
cumulative_time_offset: float = 0.0
global_time_offset: float = 0.0
speaker: int = -1
# KV cache state
kv_cache: Optional[DynamicCache] = None
kv_seq_len: int = 0 # sequence length when KV was saved
prompt_token_count: int = 0 # tokens before audio (system prompt etc)
audio_token_count: int = 0 # audio tokens in the cached KV
generated_token_ids: List[int] = field(default_factory=list)
# Alignment tracking
last_attend_frame: int = -15
committed_text: str = ""
committed_word_count: int = 0
committed_token_ids: List[int] = field(default_factory=list)
# Tracking
first_timestamp: Optional[float] = None
detected_language: Optional[str] = None
last_infer_samples: int = 0
# Audio embedding cache
audio_cache: _AudioEmbedCache = field(default_factory=_AudioEmbedCache)
def reset_kv(self):
"""Reset KV cache (e.g., when audio is trimmed from front)."""
self.kv_cache = None
self.kv_seq_len = 0
self.prompt_token_count = 0
self.audio_token_count = 0
self.generated_token_ids = []
# Reset alignment tracking — old frame references are invalid
# after audio is trimmed from the front
self.last_attend_frame = -15
class Qwen3SimulKVASR:
"""
Shared backend for Qwen3-ASR SimulStreaming with KV cache reuse.
"""
sep = ""
def __init__(
self,
model_size: str = None,
model_dir: str = None,
lan: str = "auto",
alignment_heads_path: Optional[str] = None,
border_fraction: float = 0.15,
min_chunk_size: float = 0.1,
warmup_file: Optional[str] = None,
model_cache_dir: Optional[str] = None,
model_path: Optional[str] = None,
lora_path: Optional[str] = None,
direct_english_translation: bool = False,
**kwargs,
):
self.transcribe_kargs = {}
self.original_language = None if lan == "auto" else lan
self.warmup_file = warmup_file
self.cfg = Qwen3SimulKVConfig(
language=lan,
alignment_heads_path=alignment_heads_path,
border_fraction=border_fraction,
)
self._load_model(model_size, model_dir, model_cache_dir, model_path)
self.alignment_heads = self._load_alignment_heads(alignment_heads_path)
# Pre-compute heads by layer for efficient hook installation
self.heads_by_layer = {}
for layer_idx, head_idx in self.alignment_heads:
self.heads_by_layer.setdefault(layer_idx, []).append(head_idx)
if warmup_file:
from whisperlivekit.warmup import load_file
audio = load_file(warmup_file)
if audio is not None:
self._warmup(audio)
def _load_model(self, model_size, model_dir, model_cache_dir, model_path):
from whisperlivekit.qwen3_asr import QWEN3_MODEL_MAPPING, _patch_transformers_compat
_patch_transformers_compat()
from qwen_asr.core.transformers_backend import (
Qwen3ASRConfig, Qwen3ASRForConditionalGeneration, Qwen3ASRProcessor,
)
from transformers import AutoConfig, AutoModel, AutoProcessor
AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
AutoModel.register(Qwen3ASRConfig, Qwen3ASRForConditionalGeneration)
AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)
if model_dir:
model_id = model_dir
elif model_path:
model_id = model_path
elif model_size:
model_id = QWEN3_MODEL_MAPPING.get(model_size.lower(), model_size)
else:
model_id = "Qwen/Qwen3-ASR-1.7B"
if torch.cuda.is_available():
dtype, device = torch.bfloat16, "cuda:0"
else:
dtype, device = torch.float32, "cpu"
logger.info("Loading Qwen3-ASR for SimulStreaming+KV: %s", model_id)
self.model = AutoModel.from_pretrained(model_id, dtype=dtype, device_map=device)
self.model.eval()
self.processor = AutoProcessor.from_pretrained(model_id, fix_mistral_regex=True)
thinker = self.model.thinker
text_config = thinker.config.text_config
self.num_layers = text_config.num_hidden_layers
self.num_heads = text_config.num_attention_heads
self.num_kv_heads = text_config.num_key_value_heads
self.audio_token_id = thinker.config.audio_token_id
self.device = next(self.model.parameters()).device
self.dtype = next(self.model.parameters()).dtype
self.asr_text_token_id = self.processor.tokenizer.convert_tokens_to_ids("<asr_text>")
# EOS tokens
self.eos_ids = {151645, 151643}
if self.processor.tokenizer.eos_token_id is not None:
self.eos_ids.add(self.processor.tokenizer.eos_token_id)
logger.info(
"Qwen3-ASR loaded: %d layers x %d heads, device=%s",
self.num_layers, self.num_heads, self.device,
)
def _load_alignment_heads(self, path):
max_heads = self.cfg.max_alignment_heads
if path and Path(path).exists():
with open(path) as f:
data = json.load(f)
all_heads = [tuple(h) for h in data["alignment_heads_compact"]]
heads = all_heads[:max_heads]
logger.info("Loaded top %d alignment heads from %s", len(heads), path)
return heads
default_heads = []
start_layer = self.num_layers * 3 // 4
for layer in range(start_layer, self.num_layers):
for head in range(self.num_heads):
default_heads.append((layer, head))
logger.warning("No alignment heads file. Using %d default heads.", len(default_heads))
return default_heads[:max_heads]
def _warmup(self, audio):
try:
audio = audio[:SAMPLE_RATE * 2]
msgs = [{"role": "system", "content": ""}, {"role": "user", "content": [{"type": "audio", "audio": ""}]}]
text_prompt = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
inputs = self.processor(text=[text_prompt], audio=[audio], return_tensors="pt", padding=True)
inputs = inputs.to(self.device).to(self.dtype)
with torch.inference_mode():
self.model.thinker.generate(**inputs, max_new_tokens=5, do_sample=False)
logger.info("Warmup complete")
except Exception as e:
logger.warning("Warmup failed: %s", e)
def transcribe(self, audio):
pass
class Qwen3SimulKVOnlineProcessor:
"""
Per-session online processor with KV cache reuse.
Key optimization: instead of calling generate() each time (which does
full prefill), we maintain a DynamicCache and do incremental prefill
+ manual greedy decoding with alignment head hooks.
"""
SAMPLING_RATE = 16000
MIN_DURATION_REAL_SILENCE = 5
def __init__(self, asr: Qwen3SimulKVASR, logfile=sys.stderr):
self.asr = asr
self.logfile = logfile
self.end = 0.0
self.buffer: List[ASRToken] = []
self.state = Qwen3SimulKVState()
self._build_prompt_template()
def _build_prompt_template(self):
from whisperlivekit.qwen3_asr import WHISPER_TO_QWEN3_LANGUAGE
msgs = [
{"role": "system", "content": ""},
{"role": "user", "content": [{"type": "audio", "audio": ""}]},
]
self._base_prompt = self.asr.processor.apply_chat_template(
msgs, add_generation_prompt=True, tokenize=False,
)
lan = self.asr.cfg.language
if lan and lan != "auto":
lang_name = WHISPER_TO_QWEN3_LANGUAGE.get(lan, lan)
self._base_prompt += f"language {lang_name}<asr_text>"
@property
def speaker(self):
return self.state.speaker
@speaker.setter
def speaker(self, value):
self.state.speaker = value
@property
def global_time_offset(self):
return self.state.global_time_offset
@global_time_offset.setter
def global_time_offset(self, value):
self.state.global_time_offset = value
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float):
self.end = audio_stream_end_time
self.state.audio_buffer = np.append(self.state.audio_buffer, audio)
max_samples = int(self.asr.cfg.audio_max_len * self.SAMPLING_RATE)
if len(self.state.audio_buffer) > max_samples:
trim = len(self.state.audio_buffer) - max_samples
self.state.audio_buffer = self.state.audio_buffer[trim:]
self.state.cumulative_time_offset += trim / self.SAMPLING_RATE
self.state.last_infer_samples = max(0, self.state.last_infer_samples - trim)
self.state.audio_cache.reset()
self.state.reset_kv() # Must invalidate KV when audio is trimmed
def start_silence(self) -> Tuple[List[ASRToken], float]:
all_tokens = []
for _ in range(5):
tokens, _ = self.process_iter(is_last=True)
if not tokens:
break
all_tokens.extend(tokens)
return all_tokens, self.end
def end_silence(self, silence_duration: float, offset: float):
self.end += silence_duration
long_silence = silence_duration >= self.MIN_DURATION_REAL_SILENCE
if not long_silence:
gap_len = int(self.SAMPLING_RATE * silence_duration)
if gap_len > 0:
self.state.audio_buffer = np.append(
self.state.audio_buffer, np.zeros(gap_len, dtype=np.float32),
)
else:
self.state = Qwen3SimulKVState()
self.state.global_time_offset = silence_duration + offset
def new_speaker(self, change_speaker: ChangeSpeaker):
self.process_iter(is_last=True)
self.state = Qwen3SimulKVState()
self.state.speaker = change_speaker.speaker
self.state.global_time_offset = change_speaker.start
def get_buffer(self) -> Transcript:
return Transcript.from_tokens(tokens=self.buffer, sep='')
def _encode_audio(self) -> Tuple[torch.Tensor, int]:
"""Encode full audio buffer, with caching for stable windows."""
asr = self.asr
state = self.state
from qwen_asr.core.transformers_backend.processing_qwen3_asr import (
_get_feat_extract_output_lengths,
)
feat_out = asr.processor.feature_extractor(
[state.audio_buffer], sampling_rate=16000,
padding=True, truncation=False,
return_attention_mask=True, return_tensors="pt",
)
input_features = feat_out["input_features"].to(asr.device).to(asr.dtype)
feature_attention_mask = feat_out["attention_mask"].to(asr.device)
total_mel_frames = feature_attention_mask.sum().item()
total_audio_tokens = _get_feat_extract_output_lengths(
torch.tensor(total_mel_frames),
).item()
cache = state.audio_cache
audio_cfg = asr.model.thinker.audio_tower.config
n_window_infer = getattr(audio_cfg, "n_window_infer", 400)
n_complete_windows = total_mel_frames // n_window_infer
if n_complete_windows <= 0 or cache.embeddings is None:
# Full encode
audio_embeds = asr.model.thinker.get_audio_features(
input_features, feature_attention_mask=feature_attention_mask,
)
if audio_embeds.dim() == 3:
audio_embeds = audio_embeds[0]
stable_mel = n_complete_windows * n_window_infer if n_complete_windows > 0 else 0
stable_tokens = _get_feat_extract_output_lengths(
torch.tensor(stable_mel),
).item() if stable_mel > 0 else 0
else:
stable_mel = n_complete_windows * n_window_infer
stable_tokens = _get_feat_extract_output_lengths(
torch.tensor(stable_mel),
).item()
if cache.stable_tokens > 0 and cache.stable_tokens <= stable_tokens:
cached_prefix = cache.embeddings[:stable_tokens] if cache.embeddings.dim() == 2 else cache.embeddings[0, :stable_tokens]
tail_features = input_features[:, :, stable_mel:]
tail_mel_frames = total_mel_frames - stable_mel
if tail_mel_frames > 0:
tail_mask = torch.ones(
(1, tail_features.shape[2]),
dtype=feature_attention_mask.dtype,
device=feature_attention_mask.device,
)
tail_embeds = asr.model.thinker.get_audio_features(
tail_features, feature_attention_mask=tail_mask,
)
if tail_embeds.dim() == 3:
tail_embeds = tail_embeds[0]
audio_embeds = torch.cat([cached_prefix, tail_embeds], dim=0)
else:
audio_embeds = cached_prefix
else:
audio_embeds = asr.model.thinker.get_audio_features(
input_features, feature_attention_mask=feature_attention_mask,
)
if audio_embeds.dim() == 3:
audio_embeds = audio_embeds[0]
# Update cache
cache.embeddings = audio_embeds if audio_embeds.dim() == 2 else audio_embeds[0]
cache.encoded_samples = len(state.audio_buffer)
cache.encoded_mel_frames = total_mel_frames
stable_mel_final = n_complete_windows * n_window_infer if n_complete_windows > 0 else 0
cache.stable_tokens = _get_feat_extract_output_lengths(
torch.tensor(stable_mel_final),
).item() if stable_mel_final > 0 else 0
return audio_embeds, total_audio_tokens
def _build_full_inputs(self, audio_embeds: torch.Tensor) -> dict:
"""Build full input embeddings from prompt + audio embeddings + context."""
asr = self.asr
state = self.state
thinker = asr.model.thinker
from qwen_asr.core.transformers_backend.processing_qwen3_asr import (
_get_feat_extract_output_lengths,
)
n_audio_tokens = audio_embeds.shape[0]
prompt_with_placeholders = asr.processor.replace_multimodal_special_tokens(
[self._base_prompt], iter([n_audio_tokens]),
)[0]
text_ids = asr.processor.tokenizer(
[prompt_with_placeholders], return_tensors="pt", padding=True,
)
input_ids = text_ids["input_ids"].to(asr.device)
attention_mask = text_ids.get("attention_mask")
if attention_mask is not None:
attention_mask = attention_mask.to(asr.device)
# Append committed context tokens
if state.committed_token_ids:
ctx = state.committed_token_ids[-asr.cfg.max_context_tokens:]
ctx_ids = torch.tensor([ctx], dtype=input_ids.dtype, device=input_ids.device)
input_ids = torch.cat([input_ids, ctx_ids], dim=1)
if attention_mask is not None:
ctx_mask = torch.ones_like(ctx_ids)
attention_mask = torch.cat([attention_mask, ctx_mask], dim=1)
# Build inputs_embeds
inputs_embeds = thinker.get_input_embeddings()(input_ids)
audio_mask = (input_ids == asr.audio_token_id)
n_placeholders = audio_mask.sum().item()
if n_placeholders != n_audio_tokens:
logger.warning("Audio token mismatch: %d vs %d", n_placeholders, n_audio_tokens)
return None
audio_embeds_cast = audio_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
expand_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds)
inputs_embeds = inputs_embeds.masked_scatter(expand_mask, audio_embeds_cast)
# Find audio token range
audio_positions = audio_mask[0].nonzero(as_tuple=True)[0]
audio_start = audio_positions[0].item()
audio_end = audio_positions[-1].item() + 1
return {
"input_ids": input_ids,
"inputs_embeds": inputs_embeds,
"attention_mask": attention_mask,
"audio_start": audio_start,
"audio_end": audio_end,
"n_audio_tokens": n_audio_tokens,
}
@torch.inference_mode()
def process_iter(self, is_last=False) -> Tuple[List[ASRToken], float]:
audio_duration = len(self.state.audio_buffer) / self.SAMPLING_RATE
if audio_duration < self.asr.cfg.audio_min_len:
return [], self.end
new_samples = len(self.state.audio_buffer) - self.state.last_infer_samples
min_new_seconds = self.asr.cfg.min_new_seconds
if not is_last and new_samples < int(min_new_seconds * self.SAMPLING_RATE):
return [], self.end
self.state.last_infer_samples = len(self.state.audio_buffer)
try:
timestamped_words = self._infer(is_last)
except Exception as e:
logger.exception("Inference error: %s", e)
self.state.reset_kv()
return [], self.end
if not timestamped_words:
return [], self.end
self.buffer = []
return timestamped_words, self.end
def _infer(self, is_last: bool) -> List[ASRToken]:
"""Run inference with KV cache reuse and alignment-head stopping."""
asr = self.asr
state = self.state
thinker = asr.model.thinker
# Step 1: Encode audio (with caching)
audio_embeds, n_audio_tokens_total = self._encode_audio()
# Step 2: Build full inputs
full_inputs = self._build_full_inputs(audio_embeds)
if full_inputs is None:
state.reset_kv()
return []
input_ids = full_inputs["input_ids"]
inputs_embeds = full_inputs["inputs_embeds"]
attention_mask = full_inputs["attention_mask"]
audio_start = full_inputs["audio_start"]
audio_end = full_inputs["audio_end"]
n_audio_tokens = full_inputs["n_audio_tokens"]
audio_duration = len(state.audio_buffer) / self.SAMPLING_RATE
# Step 3: Full prefill (we always re-prefill since audio tokens change)
# Future optimization: partial prefill when only tail audio changes
out = thinker(
input_ids=input_ids,
inputs_embeds=inputs_embeds,
attention_mask=attention_mask,
use_cache=True,
)
kv_cache = out.past_key_values
prompt_len = input_ids.shape[1]
# Step 4: Greedy decode with alignment head stopping
border_threshold = max(2, int(n_audio_tokens * asr.cfg.border_fraction))
rewind_threshold = max(2, int(n_audio_tokens * asr.cfg.rewind_fraction))
last_attend_frame = state.last_attend_frame
# Install hooks for alignment head attention extraction
decoder_layers = thinker.model.layers
num_kv_heads = asr.num_kv_heads
num_heads = asr.num_heads
gqa_ratio = num_heads // num_kv_heads
from qwen_asr.core.transformers_backend.modeling_qwen3_asr import apply_rotary_pos_emb
per_step_frames: List[List[int]] = []
current_step_frames: List[int] = []
hooks = []
def _make_attn_hook(layer_idx):
head_indices = asr.heads_by_layer[layer_idx]
def hook_fn(module, args, kwargs, output):
hidden_states = kwargs.get('hidden_states')
if hidden_states is None:
hidden_states = args[0] if args else None
if hidden_states is None or hidden_states.shape[1] != 1:
return
position_embeddings = kwargs.get('position_embeddings')
if position_embeddings is None and len(args) > 1:
position_embeddings = args[1]
past_kv = kwargs.get('past_key_values')
if position_embeddings is None or past_kv is None:
return
hidden_shape = (*hidden_states.shape[:-1], -1, module.head_dim)
q = module.q_norm(module.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
cos, sin = position_embeddings
q, _ = apply_rotary_pos_emb(q, q, cos, sin)
cache_layer = past_kv.layers[module.layer_idx]
k = cache_layer.keys
if k is None or audio_end > k.shape[2]:
return
for h_idx in head_indices:
if h_idx >= q.shape[1]:
continue
kv_h_idx = h_idx // gqa_ratio
q_h = q[0, h_idx, 0]
k_audio = k[0, kv_h_idx, audio_start:audio_end]
scores = torch.matmul(k_audio, q_h)
frame = scores.argmax().item()
current_step_frames.append(frame)
return hook_fn
for layer_idx in asr.heads_by_layer:
if layer_idx < len(decoder_layers):
h = decoder_layers[layer_idx].self_attn.register_forward_hook(
_make_attn_hook(layer_idx), with_kwargs=True,
)
hooks.append(h)
try:
# Greedy decoding with alignment-based stopping
next_token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
generated_ids = []
border_stop_step = None
tokens_per_sec = 6
if is_last:
max_tokens = min(int(audio_duration * tokens_per_sec) + 10, 120)
else:
new_audio_secs = (len(state.audio_buffer) - state.last_infer_samples) / self.SAMPLING_RATE
max_tokens = min(int(max(new_audio_secs, 1.0) * tokens_per_sec) + 5, 40)
for step in range(max_tokens):
tid = next_token.item()
if tid in asr.eos_ids:
break
generated_ids.append(tid)
# Collect alignment frames for this step
if current_step_frames:
per_step_frames.append(current_step_frames)
current_step_frames = []
# Check stopping criteria (after 3 tokens)
if not is_last and len(per_step_frames) >= 3:
latest = per_step_frames[-1]
if latest:
frames_sorted = sorted(latest)
attended = frames_sorted[len(frames_sorted) // 2]
if last_attend_frame - attended > rewind_threshold:
border_stop_step = max(0, len(per_step_frames) - 2)
break
last_attend_frame = attended
if (n_audio_tokens - attended) <= border_threshold:
border_stop_step = len(per_step_frames) - 1
break
# Next token
out = thinker(
input_ids=next_token,
past_key_values=kv_cache,
use_cache=True,
)
kv_cache = out.past_key_values
next_token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
# Flush remaining frames
if current_step_frames:
per_step_frames.append(current_step_frames)
finally:
for h in hooks:
h.remove()
state.last_attend_frame = last_attend_frame
if not generated_ids:
return []
# Strip metadata prefix (<asr_text> token)
all_generated = torch.tensor(generated_ids, device=asr.device)
num_gen = len(generated_ids)
asr_text_id = asr.asr_text_token_id
metadata_offset = 0
for i in range(min(num_gen, 10)):
if generated_ids[i] == asr_text_id:
if state.detected_language is None and i > 0:
from whisperlivekit.qwen3_asr import QWEN3_TO_WHISPER_LANGUAGE
prefix_text = asr.processor.tokenizer.decode(
generated_ids[:i], skip_special_tokens=True,
).strip()
parts = prefix_text.split()
if len(parts) >= 2:
lang_name = parts[-1]
if lang_name.lower() != "none":
state.detected_language = QWEN3_TO_WHISPER_LANGUAGE.get(
lang_name, lang_name.lower(),
)
metadata_offset = i + 1
break
if metadata_offset > 0:
generated_ids = generated_ids[metadata_offset:]
num_gen -= metadata_offset
per_step_frames = per_step_frames[metadata_offset:]
if num_gen <= 0:
return []
# Determine emit count
if border_stop_step is not None:
emit_up_to = min(border_stop_step, num_gen)
else:
emit_up_to = num_gen
emitted_ids = generated_ids[:emit_up_to]
if not emitted_ids:
return []
# Build timestamped words
words = self._build_timestamped_words(
emitted_ids, per_step_frames, emit_up_to,
n_audio_tokens, audio_duration,
)
state.committed_word_count += len(words)
# Include metadata in committed tokens for context
all_emitted = generated_ids[:emit_up_to]
if metadata_offset > 0:
all_emitted = generated_ids[:emit_up_to] # already stripped
state.committed_token_ids.extend(all_emitted)
return words
def _build_timestamped_words(
self,
generated_ids: list,
step_frames: List[List[int]],
emit_up_to: int,
n_audio_tokens: int,
audio_duration: float,
) -> List[ASRToken]:
asr = self.asr
state = self.state
per_token_frame = []
for step in range(emit_up_to):
if step < len(step_frames) and step_frames[step]:
frames = sorted(step_frames[step])
per_token_frame.append(frames[len(frames) // 2])
else:
per_token_frame.append(None)
tokenizer = asr.processor.tokenizer
full_text = tokenizer.decode(generated_ids[:emit_up_to], skip_special_tokens=True)
text_words = full_text.split()
all_frames = [f for f in per_token_frame if f is not None]
words = []
for wi, word in enumerate(text_words):
if all_frames:
frac = wi / max(len(text_words), 1)
frame_idx = min(int(frac * len(all_frames)), len(all_frames) - 1)
frame = all_frames[frame_idx]
else:
frame = None
words.append((word, frame))
tokens = []
for i, (text, frame) in enumerate(words):
text = text.strip()
if not text:
continue
if frame is not None and n_audio_tokens > 0:
timestamp = (
frame / n_audio_tokens * audio_duration
+ state.cumulative_time_offset
)
else:
timestamp = (
(i / max(len(words), 1)) * audio_duration
+ state.cumulative_time_offset
)
is_very_first_word = (i == 0 and state.committed_word_count == 0)
display_text = text if is_very_first_word else " " + text
token = ASRToken(
start=round(timestamp, 2),
end=round(timestamp + 0.1, 2),
text=display_text,
speaker=state.speaker,
detected_language=state.detected_language,
).with_offset(state.global_time_offset)
tokens.append(token)
return tokens
def warmup(self, audio: np.ndarray, init_prompt: str = ""):
try:
self.state.audio_buffer = audio[:SAMPLE_RATE]
self.process_iter(is_last=True)
self.state = Qwen3SimulKVState()
except Exception as e:
logger.warning("Warmup failed: %s", e)
self.state = Qwen3SimulKVState()
def finish(self) -> Tuple[List[ASRToken], float]:
all_tokens = []
for _ in range(5):
tokens, _ = self.process_iter(is_last=True)
if not tokens:
break
all_tokens.extend(tokens)
return all_tokens, self.end

416
whisperlivekit/vllm_realtime.py Normal file
View File

@@ -0,0 +1,416 @@
"""
vLLM Realtime WebSocket streaming backend for WhisperLiveKit.
Connects to a vLLM server's ``/v1/realtime`` WebSocket endpoint to stream
audio and receive transcription deltas. Uses ``websockets.sync.client``
for simplicity since ``process_iter`` runs inside ``asyncio.to_thread``.
Provides ``VLLMRealtimeASR`` (lightweight model holder) and
``VLLMRealtimeOnlineProcessor`` (streaming processor) that plug into
WhisperLiveKit's audio processing pipeline.
"""
import base64
import json
import logging
import threading
import time
from typing import List, Optional, Tuple
import numpy as np
from whisperlivekit.timed_objects import ASRToken, Transcript
logger = logging.getLogger(__name__)
class VLLMRealtimeASR:
"""Lightweight model holder — stores connection info for the vLLM server."""
sep = " "
SAMPLING_RATE = 16000
backend_choice = "vllm-realtime"
def __init__(self, vllm_url="ws://localhost:8000/v1/realtime",
model_name="Qwen/Qwen3-ASR-1.7B", lan="auto", **kwargs):
self.vllm_url = vllm_url
self.model_name = model_name
self.original_language = None if lan == "auto" else lan
self.tokenizer = None
def transcribe(self, audio):
pass
class VLLMRealtimeOnlineProcessor:
"""
Online processor that streams audio to a vLLM Realtime WebSocket.
Uses a background thread for WebSocket receiving and
``websockets.sync.client`` for the sync WebSocket connection.
"""
SAMPLING_RATE = 16000
# Minimum audio samples before connecting (0.5s of audio)
_MIN_CONNECT_SAMPLES = SAMPLING_RATE // 2
def __init__(self, asr: VLLMRealtimeASR):
self.asr = asr
self.end = 0.0
self.buffer = []
self.audio_buffer = np.array([], dtype=np.float32)
self._reset_state()
logger.info(
"[vllm-realtime] Initialized. url=%s model=%s",
asr.vllm_url, asr.model_name,
)
def _reset_state(self):
self._pending_audio = np.zeros(0, dtype=np.float32)
self._ws = None
self._recv_thread: Optional[threading.Thread] = None
self._connected = False
self._done = False
self._recv_error: Optional[Exception] = None
# Text accumulation and word extraction
self._accumulated_text = ""
self._n_committed_words = 0
self._total_audio_duration = 0.0
self._global_time_offset = 0.0
# Lock for text state accessed from both recv thread and main thread
self._text_lock = threading.Lock()
# ── Interface methods ──
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float):
self.end = audio_stream_end_time
self._pending_audio = np.append(self._pending_audio, audio)
self.audio_buffer = self._pending_audio
def process_iter(self, is_last=False) -> Tuple[List[ASRToken], float]:
try:
return self._process_iter_inner(is_last)
except Exception as e:
logger.warning("[vllm-realtime] process_iter exception: %s", e, exc_info=True)
return [], self.end
def get_buffer(self) -> Transcript:
"""Return all uncommitted text as buffer."""
self._drain_deltas()
with self._text_lock:
text = self._accumulated_text
if not text:
return Transcript(start=None, end=None, text="")
words = text.split()
uncommitted = words[self._n_committed_words:]
if uncommitted:
return Transcript(start=self.end, end=self.end, text=" ".join(uncommitted))
return Transcript(start=None, end=None, text="")
def start_silence(self) -> Tuple[List[ASRToken], float]:
"""Flush all pending words when silence starts.
Sends commit(final=true) to signal end of utterance, waits for
transcription.done, flushes all words, then prepares for reconnection
on the next utterance.
"""
if not self._connected or self._done:
words = self._flush_all_pending_words()
logger.info("[vllm-realtime] start_silence (not connected): flushed %d words", len(words))
return words, self.end
# Send any remaining buffered audio
self._send_pending_audio()
# Signal end of stream
self._send_commit(final=True)
# Wait for transcription.done
self._wait_for_done(timeout=10.0)
# Flush all remaining words
words = self._flush_all_pending_words()
# Close and reset for next utterance
self._close_ws()
old_offset = self._global_time_offset + self._total_audio_duration
self._reset_state()
self._global_time_offset = old_offset
logger.info("[vllm-realtime] start_silence: flushed %d words", len(words))
return words, self.end
def end_silence(self, silence_duration: float, offset: float):
self._global_time_offset += silence_duration
self.end += silence_duration
def new_speaker(self, change_speaker):
self.start_silence()
def warmup(self, audio, init_prompt=""):
pass
def finish(self) -> Tuple[List[ASRToken], float]:
"""Close connection and flush all remaining words."""
if self._connected and not self._done:
# Send remaining audio
self._send_pending_audio()
# Signal final commit
self._send_commit(final=True)
# Wait for transcription.done
self._wait_for_done(timeout=30.0)
# Flush all words
words = self._flush_all_pending_words()
# Close WebSocket
self._close_ws()
logger.info("[vllm-realtime] finish: flushed %d words", len(words))
return words, self.end
# ── WebSocket connection management ──
def _connect(self):
"""Connect to the vLLM realtime WebSocket and start the receive thread."""
from websockets.sync.client import connect
url = self.asr.vllm_url
logger.info("[vllm-realtime] Connecting to %s", url)
self._ws = connect(url)
# Send session.update to select model
self._ws.send(json.dumps({
"type": "session.update",
"model": self.asr.model_name,
}))
# Send initial commit(final=false) to start generation
self._send_commit(final=False)
# Start receive thread
self._recv_thread = threading.Thread(target=self._recv_loop, daemon=True)
self._recv_thread.start()
self._connected = True
logger.info("[vllm-realtime] Connected and started receive thread")
def _close_ws(self):
"""Close the WebSocket connection and join the receive thread."""
if self._ws is not None:
try:
self._ws.close()
except Exception:
pass
self._ws = None
if self._recv_thread is not None:
self._recv_thread.join(timeout=5.0)
self._recv_thread = None
def _recv_loop(self):
"""Background thread: receive messages from the vLLM WebSocket."""
try:
while not self._done and self._ws is not None:
try:
raw = self._ws.recv(timeout=0.1)
except TimeoutError:
continue
except Exception:
break
try:
msg = json.loads(raw)
except (json.JSONDecodeError, TypeError):
continue
msg_type = msg.get("type", "")
if msg_type == "transcription.delta":
delta = msg.get("delta", "")
if delta:
with self._text_lock:
self._accumulated_text += delta
elif msg_type == "transcription.done":
done_text = msg.get("text", "")
if done_text:
with self._text_lock:
# Replace accumulated text with final text
self._accumulated_text = done_text
self._done = True
break
except Exception as e:
logger.error("[vllm-realtime] recv_loop error: %s", e, exc_info=True)
self._recv_error = e
self._done = True
# ── Protocol messages ──
def _send_commit(self, final: bool):
"""Send input_audio_buffer.commit message."""
if self._ws is None:
return
try:
self._ws.send(json.dumps({
"type": "input_audio_buffer.commit",
"final": final,
}))
except Exception as e:
logger.warning("[vllm-realtime] Failed to send commit: %s", e)
def _send_audio(self, audio: np.ndarray):
"""Send audio as a base64-encoded PCM16 append message."""
if self._ws is None:
return
# Convert float32 [-1, 1] to int16 PCM
pcm16 = (audio * 32767).astype(np.int16)
audio_bytes = pcm16.tobytes()
audio_b64 = base64.b64encode(audio_bytes).decode("ascii")
try:
self._ws.send(json.dumps({
"type": "input_audio_buffer.append",
"audio": audio_b64,
}))
except Exception as e:
logger.warning("[vllm-realtime] Failed to send audio: %s", e)
def _send_pending_audio(self):
"""Send all pending audio to the vLLM server."""
if len(self._pending_audio) == 0:
return
# Track total audio duration for timestamp estimation
self._total_audio_duration += len(self._pending_audio) / self.SAMPLING_RATE
# Send in chunks of 0.5s to avoid overwhelming the WebSocket
chunk_samples = self.SAMPLING_RATE // 2
while len(self._pending_audio) >= chunk_samples:
chunk = self._pending_audio[:chunk_samples]
self._send_audio(chunk)
self._pending_audio = self._pending_audio[chunk_samples:]
# Send remaining audio if any
if len(self._pending_audio) > 0:
self._send_audio(self._pending_audio)
self._pending_audio = np.zeros(0, dtype=np.float32)
self.audio_buffer = self._pending_audio
# ── Receive helpers ──
def _drain_deltas(self):
"""No-op since the recv thread accumulates text directly."""
pass
def _wait_for_done(self, timeout: float = 10.0):
"""Wait for transcription.done message from the server."""
deadline = time.time() + timeout
while not self._done and time.time() < deadline:
time.sleep(0.05)
if not self._done:
logger.warning("[vllm-realtime] Timed out waiting for transcription.done")
# ── Word extraction (same approach as VoxtralHF) ──
def _time_for_word(self, word_idx: int, n_words_total: int) -> Tuple[float, float]:
"""Estimate timestamps by linearly distributing words across audio duration."""
duration = max(self._total_audio_duration, 0.001)
n_total = max(n_words_total, 1)
start_time = (word_idx / n_total) * duration + self._global_time_offset
end_time = ((word_idx + 1) / n_total) * duration + self._global_time_offset
return start_time, end_time
def _extract_new_words(self) -> List[ASRToken]:
"""Extract complete words (all but the last, which may still grow)."""
with self._text_lock:
text = self._accumulated_text
if not text:
return []
words = text.split()
new_words: List[ASRToken] = []
n_words_total = len(words)
while len(words) > self._n_committed_words + 1:
word = words[self._n_committed_words]
start_time, end_time = self._time_for_word(self._n_committed_words, n_words_total)
text_out = word if self._n_committed_words == 0 else " " + word
new_words.append(ASRToken(start=start_time, end=end_time, text=text_out))
self._n_committed_words += 1
return new_words
def _flush_all_pending_words(self) -> List[ASRToken]:
"""Flush ALL words including the last partial one."""
with self._text_lock:
text = self._accumulated_text
if not text:
return []
words = text.split()
new_words: List[ASRToken] = []
n_words_total = max(len(words), 1)
while self._n_committed_words < len(words):
word = words[self._n_committed_words]
start_time, end_time = self._time_for_word(self._n_committed_words, n_words_total)
text_out = word if self._n_committed_words == 0 else " " + word
new_words.append(ASRToken(start=start_time, end=end_time, text=text_out))
self._n_committed_words += 1
return new_words
# ── Core processing ──
def _process_iter_inner(self, is_last: bool) -> Tuple[List[ASRToken], float]:
# Connect when we have enough audio buffered
if not self._connected:
if len(self._pending_audio) >= self._MIN_CONNECT_SAMPLES:
self._connect()
self._send_pending_audio()
else:
return [], self.end
# Send any new pending audio
if self._connected and not self._done:
self._send_pending_audio()
# If connection closed unexpectedly but new audio arrived, reconnect
if self._done and len(self._pending_audio) >= self._MIN_CONNECT_SAMPLES:
flush_words = self._flush_all_pending_words()
old_offset = self._global_time_offset + self._total_audio_duration
self._close_ws()
self._reset_state()
self._global_time_offset = old_offset
self._connect()
self._send_pending_audio()
return flush_words, self.end
# Extract complete words
new_words = self._extract_new_words()
if new_words:
logger.info(
"[vllm-realtime] returning %d words: %s",
len(new_words), [w.text for w in new_words],
)
self.buffer = []
return new_words, self.end

219
whisperlivekit/voxtral_hf_streaming.py
View File

@@ -102,7 +102,8 @@ class VoxtralHFStreamingOnlineProcessor:
)
def _reset_state(self):
self._pending_audio = np.zeros(0, dtype=np.float32)
self._pending_chunks: List[np.ndarray] = []
self._pending_len = 0
self._audio_queue: queue.Queue = queue.Queue()
self._streamer_texts: List[str] = []
self._generate_thread: Optional[threading.Thread] = None
@@ -110,22 +111,63 @@ class VoxtralHFStreamingOnlineProcessor:
self._generate_finished = False
self._generate_error: Optional[Exception] = None
# Text accumulation and word extraction
self._accumulated_text = ""
# Text accumulation (list of fragments, joined on demand)
self._text_fragments: List[str] = []
self._text_len = 0
# Fragment position tracking for accurate word timestamps:
# each entry is (char_offset_in_full_text, audio_tok_pos_consumed)
self._fragment_positions: List[Tuple[int, int]] = []
self._n_text_tokens_received = 0
self._n_audio_tokens_fed = 0
# Audio tokens actually consumed by the model (tracked inside generator)
self._n_audio_tokens_consumed = 0
self._n_committed_words = 0
self._global_time_offset = 0.0
# Event signalled by the generate thread when it finishes
self._generate_done = threading.Event()
# Lock for text state accessed from both generate thread and main thread
self._text_lock = threading.Lock()
# ── Audio / text helpers ──
def _get_pending_audio(self) -> np.ndarray:
"""Flatten pending audio chunks into a single array."""
if not self._pending_chunks:
return np.zeros(0, dtype=np.float32)
if len(self._pending_chunks) == 1:
return self._pending_chunks[0]
flat = np.concatenate(self._pending_chunks)
self._pending_chunks = [flat]
return flat
def _set_pending_audio(self, arr: np.ndarray):
"""Replace pending audio with a single array."""
if len(arr) == 0:
self._pending_chunks = []
self._pending_len = 0
else:
self._pending_chunks = [arr]
self._pending_len = len(arr)
def _get_accumulated_text(self) -> str:
"""Get the full accumulated text (joins fragments if needed)."""
if not self._text_fragments:
return ""
if len(self._text_fragments) == 1:
return self._text_fragments[0]
joined = "".join(self._text_fragments)
self._text_fragments = [joined]
return joined
# ── Interface methods ──
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float):
self.end = audio_stream_end_time
self._pending_audio = np.append(self._pending_audio, audio)
self.audio_buffer = self._pending_audio
self._pending_chunks.append(audio)
self._pending_len += len(audio)
self.audio_buffer = audio # diagnostic only
def process_iter(self, is_last=False) -> Tuple[List[ASRToken], float]:
try:
@@ -142,7 +184,7 @@ class VoxtralHFStreamingOnlineProcessor:
"""
self._drain_streamer()
with self._text_lock:
text = self._accumulated_text
text = self._get_accumulated_text()
if not text:
return Transcript(start=None, end=None, text="")
@@ -174,16 +216,17 @@ class VoxtralHFStreamingOnlineProcessor:
# real audio and shouldn't affect word timestamp calculations.
if self._right_pad_samples > 0:
right_pad = np.zeros(self._right_pad_samples, dtype=np.float32)
self._pending_audio = np.append(self._pending_audio, right_pad)
self._pending_chunks.append(right_pad)
self._pending_len += len(right_pad)
saved_count = self._n_audio_tokens_fed
self._feed_pending_audio()
self._n_audio_tokens_fed = saved_count
# Drain in a loop: the model may still be processing right-padding
# chunks after the first drain returns. Keep draining until no new
# text appears for two consecutive rounds.
# Drain in a loop: the model may continue producing text tokens after
# the audio queue is empty (autoregressive generation). Each iteration
# uses an event-driven blocking drain with short timeouts.
all_words: List[ASRToken] = []
for _ in range(5): # at most 5 drain+flush cycles
for _ in range(5):
self._drain_streamer_blocking(timeout=5.0)
batch = self._flush_all_pending_words()
all_words.extend(batch)
@@ -208,7 +251,8 @@ class VoxtralHFStreamingOnlineProcessor:
# Add right-padding so the model can finish decoding
if self._right_pad_samples > 0:
right_pad = np.zeros(self._right_pad_samples, dtype=np.float32)
self._pending_audio = np.append(self._pending_audio, right_pad)
self._pending_chunks.append(right_pad)
self._pending_len += len(right_pad)
# Feed remaining audio
if self._generate_started and not self._generate_finished:
@@ -218,7 +262,7 @@ class VoxtralHFStreamingOnlineProcessor:
# Wait for generate to finish
if self._generate_thread is not None:
self._generate_thread.join(timeout=30.0)
elif not self._generate_started and len(self._pending_audio) >= self._first_chunk_samples:
elif not self._generate_started and self._pending_len >= self._first_chunk_samples:
# Never started but have enough audio — start and immediately finish
self._start_generate_thread()
self._feed_pending_audio()
@@ -242,8 +286,9 @@ class VoxtralHFStreamingOnlineProcessor:
model = self.asr.model
# Extract first chunk
first_chunk_audio = self._pending_audio[:self._first_chunk_samples]
self._pending_audio = self._pending_audio[self._first_chunk_samples:]
pending = self._get_pending_audio()
first_chunk_audio = pending[:self._first_chunk_samples]
self._set_pending_audio(pending[self._first_chunk_samples:])
# First chunk covers multiple audio tokens
self._n_audio_tokens_fed += max(1, self._first_chunk_samples // self._chunk_step)
@@ -265,11 +310,14 @@ class VoxtralHFStreamingOnlineProcessor:
audio_queue = self._audio_queue
def input_features_gen():
# Track audio consumption inside the generator (runs in generate thread)
self._n_audio_tokens_consumed = max(1, self._first_chunk_samples // self._chunk_step)
yield first_inputs.input_features
while True:
chunk_audio = audio_queue.get()
if chunk_audio is None:
break
self._n_audio_tokens_consumed += 1
inputs = processor(
chunk_audio,
is_streaming=True,
@@ -298,6 +346,7 @@ class VoxtralHFStreamingOnlineProcessor:
self._generate_error = e
finally:
self._generate_finished = True
self._generate_done.set()
self._generate_thread = threading.Thread(target=run_generate, daemon=True)
self._generate_thread.start()
@@ -309,13 +358,22 @@ class VoxtralHFStreamingOnlineProcessor:
chunk_size = self._chunk_samples
step_size = self._chunk_step
while len(self._pending_audio) >= chunk_size:
chunk = self._pending_audio[:chunk_size]
pending = self._get_pending_audio()
while len(pending) >= chunk_size:
chunk = pending[:chunk_size]
self._audio_queue.put(chunk)
self._pending_audio = self._pending_audio[step_size:]
pending = pending[step_size:]
self._n_audio_tokens_fed += 1
self.audio_buffer = self._pending_audio
self._set_pending_audio(pending)
self.audio_buffer = pending
def _append_text_fragment(self, text_fragment: str):
"""Append a text fragment with its audio position (must hold _text_lock)."""
self._fragment_positions.append((self._text_len, self._n_audio_tokens_consumed))
self._text_fragments.append(text_fragment)
self._text_len += len(text_fragment)
self._n_text_tokens_received += 1
def _drain_streamer(self):
"""Non-blocking drain of all available text from the streamer."""
@@ -333,19 +391,13 @@ class VoxtralHFStreamingOnlineProcessor:
break
if text_fragment:
with self._text_lock:
self._accumulated_text += text_fragment
self._n_text_tokens_received += 1
self._append_text_fragment(text_fragment)
def _drain_streamer_blocking(self, timeout=30.0):
"""Blocking drain: wait for the generate thread to process all queued
audio and produce the corresponding text.
"""Blocking drain: wait for the generate thread to finish producing text.
Polls the text queue while the audio queue has items (model still
processing). Once the audio queue is empty, waits for trailing
tokens, then returns.
This is critical for start_silence(): without it, the non-blocking
drain races with the generate thread and the last words get stuck.
Uses the _generate_done event to know when the model is truly finished.
Falls back to text-queue polling with adaptive timeouts.
"""
if not self._generate_started or self._generate_finished:
self._drain_streamer()
@@ -353,52 +405,101 @@ class VoxtralHFStreamingOnlineProcessor:
text_queue = self._streamer.text_queue
deadline = time.time() + timeout
# Count consecutive empty polls to detect when model has caught up
empty_streak = 0
while time.time() < deadline:
# Short poll while model is still processing queued audio;
# longer wait once the audio queue is empty (trailing tokens).
wait = 2.0 if self._audio_queue.empty() else 0.1
remaining = max(deadline - time.time(), 0.01)
# If generate thread is done, do a final flush and exit
if self._generate_done.is_set() or self._generate_finished:
self._drain_streamer()
return
# Adaptive wait: short while audio is queued, longer once queue is empty
if self._audio_queue.empty():
wait = min(remaining, 0.5)
else:
wait = min(remaining, 0.1)
try:
text_fragment = text_queue.get(timeout=wait)
except queue.Empty:
if self._audio_queue.empty():
break # Audio done + no text for 2s → fully caught up
continue # Audio still queued, model still working
empty_streak += 1
# Only exit if audio queue is empty AND we've had enough empty polls
# This prevents premature exit when the model is slow
if self._audio_queue.empty() and empty_streak >= 4:
break
continue
empty_streak = 0
if text_fragment is None:
self._generate_finished = True
break
if text_fragment:
with self._text_lock:
self._accumulated_text += text_fragment
self._n_text_tokens_received += 1
self._append_text_fragment(text_fragment)
# ── Word extraction ──
def _pos_to_time(self, token_position: int) -> float:
"""Convert token position to seconds."""
"""Convert audio token position to seconds."""
return token_position * self._seconds_per_token + self._global_time_offset
def _audio_pos_for_char(self, char_idx: int) -> int:
"""Look up the audio token position for a character index in the text.
Uses the fragment position index recorded when text arrives from the
generate thread. Returns the audio position of the fragment that
contains ``char_idx``, giving much better word timestamps than the
old uniform-distribution heuristic.
"""
if not self._fragment_positions:
return 0
# _fragment_positions is sorted by char_offset — find the last entry
# whose char_offset <= char_idx (the fragment containing this char).
pos = 0
for offset, audio_tok in self._fragment_positions:
if offset > char_idx:
break
pos = audio_tok
return pos
def _word_timestamps(self, text: str, words: List[str], start_idx: int, end_idx: int) -> List[Tuple[int, int]]:
"""Compute (tok_start, tok_end) for words[start_idx:end_idx] using fragment positions."""
# Build char offsets for each word
result = []
char_pos = 0
for i, word in enumerate(words):
if i > 0:
char_pos += 1 # space separator
if start_idx <= i < end_idx:
tok_start = self._audio_pos_for_char(char_pos)
tok_end = self._audio_pos_for_char(char_pos + len(word))
result.append((tok_start, tok_end))
char_pos += len(word)
return result
def _extract_new_words(self) -> List[ASRToken]:
"""Extract complete words (all but the last, which may still be growing)."""
with self._text_lock:
text = self._accumulated_text
text = self._get_accumulated_text()
if not text:
return []
words = text.split()
new_words: List[ASRToken] = []
n_words_total = len(words)
n_audio_toks = max(self._n_audio_tokens_fed, 1)
n_to_commit = len(words) - 1 # keep last word (may still grow)
while len(words) > self._n_committed_words + 1:
if n_to_commit <= self._n_committed_words:
return []
timestamps = self._word_timestamps(text, words, self._n_committed_words, n_to_commit)
for tok_start, tok_end in timestamps:
word = words[self._n_committed_words]
word_idx = self._n_committed_words
tok_start = int(word_idx / n_words_total * n_audio_toks) if n_words_total > 0 else 0
tok_end = int((word_idx + 1) / n_words_total * n_audio_toks) if n_words_total > 0 else 0
start_time = self._pos_to_time(tok_start)
end_time = self._pos_to_time(tok_end)
end_time = self._pos_to_time(max(tok_end, tok_start + 1))
text_out = word if self._n_committed_words == 0 else " " + word
new_words.append(ASRToken(start=start_time, end=end_time, text=text_out))
@@ -409,24 +510,22 @@ class VoxtralHFStreamingOnlineProcessor:
def _flush_all_pending_words(self) -> List[ASRToken]:
"""Flush ALL words including the last partial one."""
with self._text_lock:
text = self._accumulated_text
text = self._get_accumulated_text()
if not text:
return []
words = text.split()
new_words: List[ASRToken] = []
n_words_total = max(len(words), 1)
n_audio_toks = max(self._n_audio_tokens_fed, 1)
while self._n_committed_words < len(words):
if self._n_committed_words >= len(words):
return []
timestamps = self._word_timestamps(text, words, self._n_committed_words, len(words))
for tok_start, tok_end in timestamps:
word = words[self._n_committed_words]
word_idx = self._n_committed_words
tok_start = int(word_idx / n_words_total * n_audio_toks)
tok_end = int((word_idx + 1) / n_words_total * n_audio_toks)
start_time = self._pos_to_time(tok_start)
end_time = self._pos_to_time(tok_end)
end_time = self._pos_to_time(max(tok_end, tok_start + 1))
text_out = word if self._n_committed_words == 0 else " " + word
new_words.append(ASRToken(start=start_time, end=end_time, text=text_out))
@@ -439,7 +538,7 @@ class VoxtralHFStreamingOnlineProcessor:
def _process_iter_inner(self, is_last: bool) -> Tuple[List[ASRToken], float]:
# Start generate thread when enough audio is buffered
if not self._generate_started:
if len(self._pending_audio) >= self._first_chunk_samples:
if self._pending_len >= self._first_chunk_samples:
self._start_generate_thread()
self._feed_pending_audio()
else:
@@ -450,7 +549,7 @@ class VoxtralHFStreamingOnlineProcessor:
self._feed_pending_audio()
# If generate finished unexpectedly (EOS) but new audio arrived, restart
if self._generate_finished and len(self._pending_audio) >= self._first_chunk_samples:
if self._generate_finished and self._pending_len >= self._first_chunk_samples:
self._drain_streamer()
flush_words = self._flush_all_pending_words()
# Reset for new utterance

29
whisperlivekit/voxtral_mlx/spectrogram.py
View File

@@ -91,20 +91,33 @@ def _mel_filters() -> mx.array:
# ---------------------------------------------------------------------------
# DFT helpers
# DFT helpers (cached — these are constant for a given WINDOW_SIZE)
# ---------------------------------------------------------------------------
_CACHED_WINDOW: mx.array | None = None
_CACHED_DFT_RE: mx.array | None = None
_CACHED_DFT_IM: mx.array | None = None
def _hann_window() -> mx.array:
return mx.array(np.hanning(WINDOW_SIZE + 1)[:-1].astype(np.float32))
global _CACHED_WINDOW
if _CACHED_WINDOW is None:
_CACHED_WINDOW = mx.array(np.hanning(WINDOW_SIZE + 1)[:-1].astype(np.float32))
return _CACHED_WINDOW
def _dft_matrices():
"""Pre-compute the real / imaginary DFT basis matrices."""
n_bins = WINDOW_SIZE // 2 + 1
k = mx.arange(n_bins, dtype=mx.float32)[:, None]
n = mx.arange(WINDOW_SIZE, dtype=mx.float32)[None, :]
phase = -2.0 * math.pi * (k @ n) / WINDOW_SIZE
return mx.cos(phase), mx.sin(phase)
"""Return cached real / imaginary DFT basis matrices."""
global _CACHED_DFT_RE, _CACHED_DFT_IM
if _CACHED_DFT_RE is None:
n_bins = WINDOW_SIZE // 2 + 1
k = mx.arange(n_bins, dtype=mx.float32)[:, None]
n = mx.arange(WINDOW_SIZE, dtype=mx.float32)[None, :]
phase = -2.0 * math.pi * (k @ n) / WINDOW_SIZE
_CACHED_DFT_RE = mx.cos(phase)
_CACHED_DFT_IM = mx.sin(phase)
mx.eval(_CACHED_DFT_RE, _CACHED_DFT_IM)
return _CACHED_DFT_RE, _CACHED_DFT_IM
def _stft_frames(audio: mx.array, window: mx.array) -> mx.array:

255
whisperlivekit/voxtral_mlx_asr.py
View File

@@ -34,6 +34,11 @@ logger = logging.getLogger(__name__)
# Decoder sliding-window size (matches the model's training configuration).
_DECODER_WINDOW = 8192
# Maximum continuous decoding positions before forcing a reset.
# Beyond ~20s of continuous audio the autoregressive context drifts and
# produces hallucination. 20s / 80ms per token = 250 tokens.
_MAX_CONTINUOUS_POSITIONS = 250
def _prompt_tokens(tokenizer, n_left_pad=LEFT_PAD_TOKENS, n_delay=6):
"""Build the prompt token sequence and return ``(token_ids, n_delay)``."""
@@ -135,8 +140,9 @@ class VoxtralMLXOnlineProcessor:
def _reset_state(self):
"""Reset all incremental state for a fresh utterance."""
# Audio accumulation
self._pending = np.zeros(0, dtype=np.float32)
# Audio accumulation (list of chunks, concatenated on demand)
self._pending_chunks: list[np.ndarray] = []
self._pending_len = 0
# Mel overlap
self._mel_overlap: np.ndarray | None = None
# Encoder incremental state
@@ -151,6 +157,7 @@ class VoxtralMLXOnlineProcessor:
self._last_token: mx.array | None = None
# Bookkeeping
self._samples_encoded = 0
self._real_samples_encoded = 0 # only real audio, excludes silence padding
self._positions_decoded = 0
self._prefilled = False
self._first_chunk = True
@@ -167,10 +174,31 @@ class VoxtralMLXOnlineProcessor:
# -- audio ingestion --
def _get_pending(self) -> np.ndarray:
"""Flatten pending chunks into a single array."""
if not self._pending_chunks:
return np.zeros(0, dtype=np.float32)
if len(self._pending_chunks) == 1:
return self._pending_chunks[0]
flat = np.concatenate(self._pending_chunks)
self._pending_chunks = [flat]
return flat
def _set_pending(self, arr: np.ndarray):
"""Replace pending audio with a single array."""
if len(arr) == 0:
self._pending_chunks = []
self._pending_len = 0
else:
self._pending_chunks = [arr]
self._pending_len = len(arr)
def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float):
self.end = audio_stream_end_time
self._pending = np.append(self._pending, audio)
self.audio_buffer = self._pending
self._pending_chunks.append(audio)
self._pending_len += len(audio)
self._real_samples_encoded += len(audio)
self.audio_buffer = audio # diagnostic only
# -- core processing --
@@ -182,14 +210,28 @@ class VoxtralMLXOnlineProcessor:
return [], self.end
def _step(self, is_last: bool) -> Tuple[List[ASRToken], float]:
# 0. Safety cap: if continuous decoding exceeds the limit, force a
# flush+reset to prevent hallucination even without VAD silence.
if self._prefilled and self._positions_decoded >= _MAX_CONTINUOUS_POSITIONS + self._prefix_len:
logger.info(
"[voxtral-mlx] continuous decoding cap hit at %d positions — "
"forcing flush+reset",
self._positions_decoded,
)
words = self._flush_and_reset()
return words, self.end
# 1. Encode any new audio
self._encode_pending()
if self._audio_embeds is None:
return [], self.end
# 2. Compute how many positions we can safely decode
total_safe = LEFT_PAD_TOKENS + self._samples_encoded // SAMPLES_PER_TOKEN
# 2. Compute how many positions we can safely decode.
# The safe boundary prevents the decoder from running ahead of the
# audio encoder. _samples_encoded tracks only real audio (not
# silence padding), so positions beyond this produce hallucination.
total_safe = LEFT_PAD_TOKENS + self._real_samples_encoded // SAMPLES_PER_TOKEN
n_available = self._audio_embeds.shape[0]
n_decodable = min(n_available, total_safe - self._positions_decoded)
@@ -208,11 +250,19 @@ class VoxtralMLXOnlineProcessor:
if n_decodable <= 0 or self._audio_embeds is None:
return [], self.end
# Clamp to the continuous decoding cap so we don't overshoot
max_left = _MAX_CONTINUOUS_POSITIONS + self._prefix_len - self._positions_decoded
if max_left > 0:
n_decodable = min(n_decodable, max_left)
else:
# Will be caught by the cap check on the next call
return self._extract_committed_words(), self.end
# 4. Decode available positions
hit_eos = self._decode_positions(n_decodable)
if hit_eos:
# Flush words, reset for next utterance
# Flush words, then full reset for next utterance
words = self._flush_all_words()
logger.debug(
"[voxtral-mlx] EOS hit during stream: flushed %d words, "
@@ -221,9 +271,12 @@ class VoxtralMLXOnlineProcessor:
self._samples_encoded / self.SAMPLING_RATE,
self._full_text[-60:] if self._full_text else "",
)
saved_offset = self._time_offset
new_offset = self._time_offset + self._real_samples_encoded / self.SAMPLING_RATE
saved_end = self.end
self._reset_state()
self._time_offset = saved_offset
self._time_offset = new_offset
self.end = saved_end
mx.clear_cache()
return words, self.end
# 5. Extract committed words (all but the last, which may still grow)
@@ -231,22 +284,24 @@ class VoxtralMLXOnlineProcessor:
def _encode_pending(self):
"""Feed pending audio through the incremental encoder."""
available = len(self._pending)
if available < SAMPLES_PER_TOKEN:
if self._pending_len < SAMPLES_PER_TOKEN:
return
pending = self._get_pending()
available = len(pending)
if self._first_chunk:
# First chunk: prepend silence for left-padding
n_take = (available // SAMPLES_PER_TOKEN) * SAMPLES_PER_TOKEN
left_pad = np.zeros(LEFT_PAD_TOKENS * SAMPLES_PER_TOKEN, dtype=np.float32)
chunk = np.concatenate([left_pad, self._pending[:n_take]])
self._pending = self._pending[n_take:]
chunk = np.concatenate([left_pad, pending[:n_take]])
self._set_pending(pending[n_take:])
self._samples_encoded += n_take
self._first_chunk = False
else:
n_take = (available // SAMPLES_PER_TOKEN) * SAMPLES_PER_TOKEN
chunk = self._pending[:n_take]
self._pending = self._pending[n_take:]
chunk = pending[:n_take]
self._set_pending(pending[n_take:])
self._samples_encoded += n_take
mel, self._mel_overlap = compute_mel_streaming(chunk, self._mel_overlap)
@@ -261,11 +316,10 @@ class VoxtralMLXOnlineProcessor:
mx.eval(embeds)
if self._audio_embeds is not None:
self._audio_embeds = mx.concatenate([self._audio_embeds, embeds])
mx.eval(self._audio_embeds)
else:
self._audio_embeds = embeds
self.audio_buffer = self._pending
def _do_prefill(self):
"""Run the decoder prefill pass over the prompt + first audio embeddings."""
n_dec_layers = len(self._model.decoder.blocks)
@@ -429,8 +483,114 @@ class VoxtralMLXOnlineProcessor:
return Transcript(start=self.end, end=self.end, text=" ".join(remaining))
return Transcript(start=None, end=None, text="")
def start_silence(self) -> Tuple[List[ASRToken], float]:
def _safe_decode_remaining(self):
"""Decode remaining audio embeddings, respecting the safe boundary.
Uses the same guard as ``_step`` to avoid decoding positions that
are beyond the real audio frontier, which causes hallucination.
"""
if self._audio_embeds is None or not self._prefilled:
return
# Use the same formula as _step() — this excludes padding positions
total_safe = LEFT_PAD_TOKENS + self._samples_encoded // SAMPLES_PER_TOKEN
n_available = self._audio_embeds.shape[0]
n_decodable = min(n_available, max(0, total_safe - self._positions_decoded))
# Cap at RIGHT_PAD_TOKENS to only decode the padding needed for
# the model to emit final tokens, not all accumulated padding
n_decodable = min(n_decodable, RIGHT_PAD_TOKENS)
if n_decodable > 0:
self._decode_positions(n_decodable)
def _flush_last_token_text(self):
"""Add the last pending token's text (if not EOS) to _full_text."""
if self._last_token is None:
return
tid = self._last_token.item()
if tid == self._eos_id:
return
text = self._tokenizer.decode(
[tid], special_token_policy=SpecialTokenPolicy.IGNORE
)
if not text:
return
last_pos = self._positions_decoded - self._prefix_len
if text.lstrip() != text or not self._full_text:
if self._current_word_pos is not None:
self._word_audio_ends.append(last_pos)
self._word_audio_starts.append(last_pos)
self._current_word_pos = last_pos
elif self._current_word_pos is None:
self._word_audio_starts.append(last_pos)
self._current_word_pos = last_pos
self._full_text += text
self._n_text_tokens += 1
def _close_current_word(self):
"""Close the last word if one is being built."""
if self._current_word_pos is not None:
last_pos = self._positions_decoded - self._prefix_len
self._word_audio_ends.append(last_pos)
self._current_word_pos = None
def _flush_and_reset(self) -> List[ASRToken]:
"""Flush pending audio, decode remaining, extract all words, then
fully reset both encoder and decoder state.
Used at silence boundaries and when the continuous decoding cap is
hit. A full reset (encoder + decoder) is necessary because the
encoder's incremental state (conv tails, KV caches) contains history
that would produce embeddings incompatible with a freshly-initialised
decoder. After reset ``_first_chunk=True``, so the next audio chunk
receives proper left-padding and both encoder and decoder start in
sync.
"""
# Align pending audio to SAMPLES_PER_TOKEN boundary
remainder = self._pending_len % SAMPLES_PER_TOKEN
align_pad = (SAMPLES_PER_TOKEN - remainder) if remainder > 0 else 0
# Add alignment + right-padding silence to provide future context
total_pad = align_pad + RIGHT_PAD_TOKENS * SAMPLES_PER_TOKEN
if total_pad > 0:
self._pending_chunks.append(np.zeros(total_pad, dtype=np.float32))
self._pending_len += total_pad
# Encode remaining audio (including right-padding)
self._encode_pending()
# Decode only positions backed by real audio
self._safe_decode_remaining()
self._flush_last_token_text()
self._close_current_word()
words = self._flush_all_words()
# Compute time offset: the decoded audio covers up to this point
new_offset = self._time_offset + self._real_samples_encoded / self.SAMPLING_RATE
saved_end = self.end
# Full reset — encoder AND decoder. The encoder's incremental
# state (conv tails, transformer KV caches) carries history from
# the previous segment; keeping it would make the next set of
# embeddings incompatible with a fresh decoder prefill.
self._reset_state()
self._time_offset = new_offset
self.end = saved_end
# Free MLX caches eagerly
mx.clear_cache()
return words
def start_silence(self) -> Tuple[List[ASRToken], float]:
"""Flush all pending words when silence starts, then fully reset.
Adds right-padding silence and forces a decode pass so the
decoder emits tokens for the last words of speech. After flushing,
resets both encoder and decoder state to prevent hallucination from
accumulated autoregressive context drift on long audio.
"""
words = self._flush_and_reset()
logger.info("[voxtral-mlx] start_silence: flushed %d words", len(words))
return words, self.end
@@ -448,7 +608,7 @@ class VoxtralMLXOnlineProcessor:
logger.debug(
"[voxtral-mlx] finish: pending=%d samples, audio_embeds=%s, "
"samples_encoded=%d, positions_decoded=%d, prefilled=%s, text so far='%s'",
len(self._pending),
self._pending_len,
self._audio_embeds.shape if self._audio_embeds is not None else None,
self._samples_encoded,
self._positions_decoded,
@@ -457,64 +617,23 @@ class VoxtralMLXOnlineProcessor:
)
# Align pending audio to SAMPLES_PER_TOKEN boundary so nothing is lost
remainder = len(self._pending) % SAMPLES_PER_TOKEN
if remainder > 0:
align_pad = SAMPLES_PER_TOKEN - remainder
else:
align_pad = 0
remainder = self._pending_len % SAMPLES_PER_TOKEN
align_pad = (SAMPLES_PER_TOKEN - remainder) if remainder > 0 else 0
# Add alignment + right-padding silence
total_pad = align_pad + RIGHT_PAD_TOKENS * SAMPLES_PER_TOKEN
if total_pad > 0:
self._pending = np.append(
self._pending, np.zeros(total_pad, dtype=np.float32)
)
self._pending_chunks.append(np.zeros(total_pad, dtype=np.float32))
self._pending_len += total_pad
# Encode remaining audio (including right-padding)
self._encode_pending()
logger.debug(
"[voxtral-mlx] finish after encode: audio_embeds=%s, pending=%d",
self._audio_embeds.shape if self._audio_embeds is not None else None,
len(self._pending),
)
# Decode only positions backed by real audio
self._safe_decode_remaining()
hit_eos = False
# Decode everything that's left from right-padding
if self._audio_embeds is not None and self._prefilled:
hit_eos = self._decode_positions(self._audio_embeds.shape[0])
logger.debug(
"[voxtral-mlx] finish decode: hit_eos=%s, text='%s'",
hit_eos, self._full_text[-80:] if self._full_text else "",
)
# Flush last token if it wasn't EOS
if self._last_token is not None:
tid = self._last_token.item()
if tid != self._eos_id:
text = self._tokenizer.decode(
[tid], special_token_policy=SpecialTokenPolicy.IGNORE
)
if text:
last_pos = self._positions_decoded - self._prefix_len
# Check if this starts a new word
if text.lstrip() != text or not self._full_text:
if self._current_word_pos is not None:
self._word_audio_ends.append(last_pos)
self._word_audio_starts.append(last_pos)
self._current_word_pos = last_pos
elif self._current_word_pos is None:
self._word_audio_starts.append(last_pos)
self._current_word_pos = last_pos
self._full_text += text
self._n_text_tokens += 1
# Close the last word if still open
if self._current_word_pos is not None:
last_pos = self._positions_decoded - self._prefix_len
self._word_audio_ends.append(last_pos)
self._current_word_pos = None
self._flush_last_token_text()
self._close_current_word()
words = self._flush_all_words()
logger.info("[voxtral-mlx] finish: flushed %d words", len(words))