42 Commits

Author SHA1 Message Date
Vineet Suryan 582d5426d6 Bump version v0.9.0 2026-06-02 11:54:07 +05:30
nightcityblade 1814dd9bfa fix: include server API dependencies in package 2026-06-01 19:50:34 -04:00
nightcityblade 0747529910 test: strengthen wheel install smoke test 2026-06-01 19:50:34 -04:00
nightcityblade cef3352c5b test: add clean virtualenv install smoke test 2026-06-01 19:50:34 -04:00
nightcityblade 3516f663e5 fix: support Python 3.13 server installs 2026-06-01 19:24:17 -04:00
nightcityblade d006d4abb3 fix: make caption overlay line count configurable
Fixes collabora/WhisperLive#486
2026-06-01 19:18:15 -04:00
nightcityblade 6242ad7f81 fix: add av to package requirements 2026-06-01 18:58:28 -04:00
Aaron Boxer 5334ea0f7a Add WebSocket authentication via api_key
- When --api_key is set, WebSocket connections require auth too
- Supports Authorization: Bearer <key> header or ?token=<key> query param
- Unauthenticated connections receive HTTP 401 before upgrade
- Uses websockets process_request callback (no resource allocation before auth)
- Added 5 unit tests for WebSocket auth handler
2026-05-26 23:07:50 -04:00
Aaron Boxer b648bcb2a4 Add optional API key auth and rate limiting for REST API
- api_key param: requires 'Authorization: Bearer <key>' header
- rate_limit_rpm param: per-IP sliding-window rate limit (requests/min)
- Both are off by default (backward compatible)
- CLI flags: --api_key, --rate_limit_rpm
- Added 5 unit tests for auth and rate limiting
2026-05-26 23:02:43 -04:00
Aaron Boxer e86d98dd80 Improve REST API unsupported param warnings
- Enumerate each ignored param individually in log message
- Add warnings for 'include' param (was previously silent)
- Added 6 unit tests for REST API param validation
2026-05-26 22:59:12 -04:00
Aaron Boxer c5ec7f4a99 Add reconnect logic to WebSocket client
- New params: max_retries (default 0), retry_delay (default 5s)
- On unexpected close, retries up to max_retries times
- Does not retry on server_error (server rejected connection)
- Extracted _create_websocket() helper for reuse
- Added 4 unit tests for reconnect behavior
2026-05-26 05:45:37 -04:00
Vineet Suryan 52005b94cb Merge pull request #443 from boxerab/configurable-constants
Extract hardcoded buffer constants into class attributes
2026-05-26 13:55:13 +05:30
Vineet Suryan f2f769532b Merge pull request #442 from boxerab/unify-translate-flags
Clarify --translate vs --enable_translation CLI flags
2026-05-26 13:25:23 +05:30
Vineet Suryan cdc661ce28 Merge pull request #434 from nightcityblade/fix/issue-405
docs: fix setup instructions reported in #405
2026-05-26 13:17:37 +05:30
Aaron Boxer 8396763444 fix: set proper mock_info attributes in SSE streaming tests
MagicMock auto-attributes are not JSON serializable. Set language,
language_probability, and duration explicitly. Also exclude metadata
events from segment count assertion.
2026-05-25 09:44:21 -04:00
Aaron Boxer 8ac98dceec feat: add SSE streaming for REST transcription endpoint
- stream=true now returns text/event-stream with per-segment SSE events
- Each segment yields 'data: {json}' followed by 'data: [DONE]'
- Error events streamed as 'data: {"error": ...}'
- Temp files cleaned up in finally block
- 5 new tests in test_server_extended.py (183 total passing)
2026-05-25 09:44:21 -04:00
Aaron Boxer 8bde966c1e docs: remove features from README that moved to Aavaaz
Remove authentication, rate limiting, and auto-reconnect documentation
since these features now live in the Aavaaz project.
2026-05-15 10:45:31 -04:00
Aaron Boxer dc4a707f9a chore: add .gitignore to exclude __pycache__, virtualenvs, and build artifacts 2026-05-15 10:45:31 -04:00
Aaron Boxer c028c4b584 feat: add segment_post_processor hook for external plugins
Add a minimal, non-breaking hook to WhisperLive that allows external
projects to post-process transcription segments before they are sent
to the client.

Changes:
- ServeClientBase: add segment_post_processor attribute (default None)
- ServeClientBase.send_transcription_to_client: apply post_processor
  per-segment with error handling (falls back to original segment)
- TranscriptionServer: add segment_post_processor parameter to run()
  and wire it to each client on creation

This enables downstream projects to plug in custom processing
(e.g. formatting, PII redaction, diarization tagging) without
modifying WhisperLive core code.
2026-05-15 10:45:31 -04:00
Aaron Boxer 4e31f8c61b Add word-level timestamps and confidence scores
- New word_timestamps option (default False) in client handshake
- When enabled, each segment includes 'words' array with per-word
  start/end times and probability scores
- Wired through entire pipeline: client → server → backend → transcribe()
- Words include timestamp_offset for accurate absolute times
- REST API already supported word timestamps; now WebSocket does too
- Added 9 unit tests for word timestamp extraction and formatting
2026-05-15 10:35:27 -04:00
Aaron Boxer 18de3eacc7 Document all new features in README
- Added 'Advanced Features' section with 8 subsections
- Word-level timestamps: WebSocket JSON example
- Custom vocabulary / hotwords: usage and REST API support
- Speaker diarization: setup, pyannote dependency, output format
- Authentication: API key for REST + WebSocket
- Rate limiting: per-IP RPM configuration
- Auto-reconnect: max_retries / retry_delay
- Batch inference: CLI flags
- Raw PCM input: int16 normalization
- Updated table of contents
2026-05-15 09:01:48 -04:00
Aaron Boxer 18b897277f Add real-time speaker diarization support
- New whisper_live/diarization.py: SpeakerDiarizer with online clustering
- Uses pyannote.audio speaker embeddings (optional dependency)
- Cosine similarity threshold for speaker matching (default 0.55)
- Running average embedding update for speaker stability
- Configurable max_speakers limit (default 10)
- Client options: enable_diarization, max_speakers
- Segments include 'speaker' field when diarization is active
- Graceful fallback: logs warning if pyannote not installed
- Added 12 unit tests (mock-based, no GPU required)
2026-05-13 10:50:59 -04:00
Aaron Boxer 3d63e82571 fix: add skip decorator to TestStartMetricsServer for CI without prometheus_client 2026-05-13 10:45:40 -04:00
Aaron Boxer ced4bdb737 feat: add Prometheus metrics instrumentation
- New whisper_live/metrics.py with Counter, Gauge, Histogram metrics
- Track connections (opened/closed/rejected), transcription latency,
  audio processed, segments emitted, REST requests, and errors
- All metric helpers are no-ops when prometheus_client not installed
- --metrics_port CLI flag to expose /metrics endpoint (0 = disabled)
- Metrics integrated into server.py, base.py at key instrumentation points
- 17 new tests in tests/test_metrics.py (178 total passing)
2026-05-13 10:45:40 -04:00
Aaron Boxer 4210697ca6 Add custom vocabulary / hotwords support 2026-05-13 10:32:28 -04:00
Aaron Boxer 445bf26e85 Extract hardcoded buffer constants into class attributes
- MAX_BUFFER_DURATION_S (45): max audio buffer before trimming
- BUFFER_TRIM_DURATION_S (30): duration to discard on trim
- CLIP_THRESHOLD_DURATION_S (25): stale audio clip threshold
- CLIP_TAIL_DURATION_S (5): audio tail to keep after clipping
- All values can now be overridden by subclasses
2026-05-11 20:05:38 -04:00
Aaron Boxer b534f9d249 Clarify --translate vs --enable_translation CLI flags
- --translate: Whisper built-in to-English translation (task=translate)
- --enable_translation: M2M100 any-to-any translation backend
- Added warning when both flags are used simultaneously
- Updated help text to distinguish the two features
2026-05-11 20:05:34 -04:00
Vineet Suryan 9a71a95ca8 Merge pull request #441 from boxerab/fix-clear-screen-shell-injection
Avoid clear screen shell injection by using ANSI escape codes
2026-05-08 17:36:11 +02:00
Vineet Suryan 485b211072 Merge pull request #440 from boxerab/input-validation-server-params
Validate server parameters on startup
2026-05-08 17:35:26 +02:00
Vineet Suryan 19847784ae Merge pull request #439 from boxerab/bounded-transcript-memory
Bound transcript memory and translation queue size
2026-05-08 17:34:41 +02:00
Vineet Suryan 52d94bf1fb Merge pull request #438 from boxerab/thread-safety-client-manager
Add thread safety to client manager with threading lock
2026-04-21 13:29:54 +02:00
Vineet Suryan 1c663d0bba Merge pull request #437 from boxerab/rawpcm
audio: add support for raw pcm input via server flag
2026-04-21 13:17:47 +02:00
Vineet Suryan 298a01f1b0 Merge pull request #436 from boxerab/testing
CI: expand test suite coverage
2026-04-20 17:58:48 +02:00
Aaron Boxer a6147a6745 Replace os.system() in clear_screen() with ANSI escape codes
- Eliminates shell injection risk from os.system('clear'/'cls')
- Uses ANSI escape sequence \033[H\033[2J instead
- Removed unused os import
- Added test verifying ANSI codes are used
2026-04-17 09:31:00 -04:00
Aaron Boxer 18bce1864a Validate server parameters on startup
- max_clients must be >= 1
- max_connection_time must be > 0
- batch_max_size must be >= 1 (when batch enabled)
- batch_window_ms must be >= 0 (when batch enabled)
- Added 5 new tests for parameter validation
2026-04-17 09:30:19 -04:00
Aaron Boxer 9e5e4a9970 Bound transcript memory and translation queue size
- Add MAX_TRANSCRIPT_LENGTH (500) and MAX_TRANSLATION_QUEUE_SIZE (100)
  class constants to ServeClientBase
- Trim transcript and text lists after each update_segments() call
- Create translation queue with maxsize to prevent unbounded growth
- Added tests for _trim_transcript()
2026-04-17 09:29:38 -04:00
Aaron Boxer 81cdbbca95 Add thread safety to ClientManager with threading.Lock
- All ClientManager methods (add_client, get_client, remove_client,
  get_wait_time, is_server_full, is_client_timeout) now protected by
  a threading.Lock
- cleanup() called outside the lock to avoid holding it during I/O
- is_server_full() computes wait time inline under lock instead of
  calling get_wait_time() to avoid nested lock acquisition
- Added concurrent thread safety tests for add/remove and get operations
2026-04-17 09:27:37 -04:00
Aaron Boxer f5340ddf1e audio: add support for raw pcm input via server flag
fixes #
2026-04-17 09:21:06 -04:00
Aaron Boxer b1cd51ac8a CI: expand test suite coverage
these new test cover issues such as thread safety, VAD thresholding,
message routing, error handling etc. that weren't covered by existing
tests. Mocking is used to avoid dependencies on GPU, ONNX etc.
2026-04-17 09:05:34 -04:00
Vineet Suryan e41324bf03 Merge pull request #435 from nightcityblade/fix/issue-327
fix: render transcript text safely in browser extensions
2026-04-16 12:37:33 +02:00
nightcityblade 31efff9330 fix: render transcript text safely in browser extensions 2026-04-15 23:09:39 +08:00
nightcityblade 68a8b57e66 docs: fix setup instructions errors reported in #405
- Clarify that setup.sh installs portaudio system dependency and list
  per-distro package names
- Add missing --gpus all flag to TensorRT Docker run command
- Fix Docker TensorRT example showing multiple --trt_model_path on one
  command (should be separate alternatives)
- Document --trt_py_session flag as workaround for TensorRT C++ session
  crashes (CrossAttentionMask warnings)

Closes #405

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-13 11:09:40 +08:00
31 changed files with 3082 additions and 133 deletions
+28 -2
View File
@@ -74,8 +74,34 @@ jobs:
# exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
venv-install-smoke-test:
runs-on: ubuntu-22.04
steps:
- uses: actions/checkout@v2
- name: Set up Python 3.12
uses: actions/setup-python@v2
with:
python-version: '3.12'
- name: Install system dependencies
run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
- name: Build package artifacts
run: |
python -m pip install --upgrade pip
python -m pip install build
python -m build --sdist --wheel
- name: Verify install in a clean virtualenv
run: |
python -m venv smoke-test-venv
source smoke-test-venv/bin/activate
pip install dist/*.whl
python -c "import whisper_live.client; import whisper_live.server"
build-and-push-docker-cpu: build-and-push-docker-cpu:
needs: [run-tests, check-code-format] needs: [run-tests, check-code-format, venv-install-smoke-test]
runs-on: ubuntu-22.04 runs-on: ubuntu-22.04
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/')) if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
steps: steps:
@@ -158,7 +184,7 @@ jobs:
tags: ghcr.io/collabora/whisperlive-openvino:latest tags: ghcr.io/collabora/whisperlive-openvino:latest
publish-to-pypi: publish-to-pypi:
needs: [run-tests, check-code-format] needs: [run-tests, check-code-format, venv-install-smoke-test]
runs-on: ubuntu-22.04 runs-on: ubuntu-22.04
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags') if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
steps: steps:
+23
View File
@@ -0,0 +1,23 @@
__pycache__/
*.pyc
*.pyo
*.egg-info/
dist/
build/
*.egg
.eggs/
whisper_env/
venv/
.venv/
env/
.env
*.so
*.o
.pytest_cache/
.mypy_cache/
.ruff_cache/
output*.srt
transcript*.srt
translation*.srt
*.wav
docs/site/
+26 -23
View File
@@ -3,6 +3,7 @@ var elem_text = null;
var segments = []; var segments = [];
var text_segments = []; var text_segments = [];
var captionLineCount = 3;
var allSegments = []; var allSegments = [];
var lastIncompleteSegment = null; var lastIncompleteSegment = null;
@@ -87,22 +88,23 @@ function showPopup(customText) {
} }
function init_element() { function init_element(lines = 3) {
captionLineCount = Math.min(Math.max(parseInt(lines, 10) || 3, 1), 8);
if (document.getElementById('transcription')) { if (document.getElementById('transcription')) {
return; return;
} }
elem_container = document.createElement('div'); elem_container = document.createElement('div');
elem_container.id = "transcription"; elem_container.id = "transcription";
elem_container.style.cssText = 'padding-top:16px;font-size:18px;position: fixed; top: 85%; left: 50%; transform: translate(-50%, -50%);line-height:18px;width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;'; elem_container.style.cssText = 'padding-top:16px;font-size:18px;position: fixed; top: 85%; left: 50%; transform: translate(-50%, -50%);line-height:18px;width:500px;height:' + (captionLineCount * 30) + 'px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
for (var i = 0; i < 4; i++) { for (var i = 0; i <= captionLineCount; i++) {
elem_text = document.createElement('span'); elem_text = document.createElement('span');
elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;'; elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;';
elem_text.id = "t" + i; elem_text.id = "t" + i;
elem_container.appendChild(elem_text); elem_container.appendChild(elem_text);
if (i == 3) { if (i == captionLineCount) {
elem_text.style.top = "-1000px" elem_text.style.top = "-1000px"
} }
} }
@@ -164,7 +166,7 @@ function get_lines(elem, line_height) {
var divHeight = elem.offsetHeight; var divHeight = elem.offsetHeight;
var lines = divHeight / line_height; var lines = divHeight / line_height;
var original_text = elem.innerHTML; var original_text = elem.textContent;
var words = original_text.split(' '); var words = original_text.split(' ');
var segments = []; var segments = [];
@@ -174,7 +176,7 @@ function get_lines(elem, line_height) {
for (var i = 0; i < words.length; i++) for (var i = 0; i < words.length; i++)
{ {
segment += words[i] + ' '; segment += words[i] + ' ';
elem.innerHTML = segment; elem.textContent = segment;
divHeight = elem.offsetHeight; divHeight = elem.offsetHeight;
if ((divHeight / line_height) > current_lines) { if ((divHeight / line_height) > current_lines) {
@@ -188,7 +190,7 @@ function get_lines(elem, line_height) {
var line_segment = segment.substring(segment_len, segment.length - 1) var line_segment = segment.substring(segment_len, segment.length - 1)
segments.push(line_segment); segments.push(line_segment);
elem.innerHTML = original_text; elem.textContent = original_text;
return segments; return segments;
@@ -196,7 +198,7 @@ function get_lines(elem, line_height) {
function remove_element() { function remove_element() {
var elem = document.getElementById('transcription') var elem = document.getElementById('transcription')
for (var i = 0; i < 4; i++) { for (var i = 0; i <= captionLineCount; i++) {
document.getElementById("t" + i).remove(); document.getElementById("t" + i).remove();
} }
elem.remove() elem.remove()
@@ -205,6 +207,7 @@ function remove_element() {
chrome.runtime.onMessage.addListener((request, sender, sendResponse) => { chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
const { type, data } = request; const { type, data } = request;
const saveCaptions = data.saveCaptions; const saveCaptions = data.saveCaptions;
const captionLines = data.captionLines || captionLineCount;
if (type === "STOP") { if (type === "STOP") {
if (saveCaptions === true) { if (saveCaptions === true) {
@@ -234,7 +237,7 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
return true; return true;
} }
init_element(); init_element(captionLines);
try { try {
const message = JSON.parse(data.data); const message = JSON.parse(data.data);
@@ -262,11 +265,11 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
} }
text = text.replace(/(\r\n|\n|\r)/gm, ""); text = text.replace(/(\r\n|\n|\r)/gm, "");
var elem = document.getElementById('t3'); var elem = document.getElementById('t' + captionLineCount);
if (elem) { if (elem) {
elem.innerHTML = text; elem.textContent = text;
var line_height_style = getStyle('t3', 'line-height'); var line_height_style = getStyle('t' + captionLineCount, 'line-height');
var line_height = parseInt(line_height_style.substring(0, line_height_style.length - 2)); var line_height = parseInt(line_height_style.substring(0, line_height_style.length - 2));
var divHeight = elem.offsetHeight; var divHeight = elem.offsetHeight;
var lines = divHeight / line_height; var lines = divHeight / line_height;
@@ -274,29 +277,29 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
text_segments = []; text_segments = [];
text_segments = get_lines(elem, line_height); text_segments = get_lines(elem, line_height);
elem.innerHTML = ''; elem.textContent = '';
if (text_segments.length > 2) { if (text_segments.length > captionLineCount - 1) {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i]; document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
} }
} else { } else {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = ''; document.getElementById('t' + i).textContent = '';
} }
} }
if (text_segments.length <= 2) { if (text_segments.length <= captionLineCount - 1) {
for (var i = 0; i < text_segments.length; i++) { for (var i = 0; i < text_segments.length; i++) {
document.getElementById('t' + i).innerHTML = text_segments[i]; document.getElementById('t' + i).textContent = text_segments[i];
} }
} else { } else {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i]; document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
} }
} }
for (var i = 1; i < 3; i++) for (var i = 1; i < captionLineCount; i++)
{ {
var parent_elem = document.getElementById('t' + (i - 1)); var parent_elem = document.getElementById('t' + (i - 1));
var elem = document.getElementById('t' + i); var elem = document.getElementById('t' + i);
+8
View File
@@ -23,6 +23,14 @@
<input type="checkbox" id="saveCaptionsCheckbox"> <input type="checkbox" id="saveCaptionsCheckbox">
<label for="saveCaptions">Download SRT file at Stop Capture</label> <label for="saveCaptions">Download SRT file at Stop Capture</label>
</div> </div>
<div class="dropdown-container">
<label for="captionLinesDropdown">Caption Lines:</label>
<select id="captionLinesDropdown">
<option value="3" selected>3 lines</option>
<option value="5">5 lines</option>
<option value="8">8 lines</option>
</select>
</div>
<div class="dropdown-container"> <div class="dropdown-container">
<label for="languageDropdown">Select Language:</label> <label for="languageDropdown">Select Language:</label>
<select id="languageDropdown"> <select id="languageDropdown">
+16
View File
@@ -9,9 +9,11 @@ document.addEventListener("DOMContentLoaded", function () {
const languageDropdown = document.getElementById('languageDropdown'); const languageDropdown = document.getElementById('languageDropdown');
const taskDropdown = document.getElementById('taskDropdown'); const taskDropdown = document.getElementById('taskDropdown');
const modelSizeDropdown = document.getElementById('modelSizeDropdown'); const modelSizeDropdown = document.getElementById('modelSizeDropdown');
const captionLinesDropdown = document.getElementById('captionLinesDropdown');
let selectedLanguage = null; let selectedLanguage = null;
let selectedTask = taskDropdown.value; let selectedTask = taskDropdown.value;
let selectedModelSize = modelSizeDropdown.value; let selectedModelSize = modelSizeDropdown.value;
let selectedCaptionLines = captionLinesDropdown.value;
// Add click event listeners to the buttons // Add click event listeners to the buttons
startButton.addEventListener("click", startCapture); startButton.addEventListener("click", startCapture);
@@ -66,6 +68,13 @@ document.addEventListener("DOMContentLoaded", function () {
} }
}); });
chrome.storage.local.get("selectedCaptionLines", ({ selectedCaptionLines: storedCaptionLines }) => {
if (storedCaptionLines !== undefined) {
captionLinesDropdown.value = storedCaptionLines;
selectedCaptionLines = storedCaptionLines;
}
});
// Function to handle the start capture button click event // Function to handle the start capture button click event
async function startCapture() { async function startCapture() {
// Ignore click if the button is disabled // Ignore click if the button is disabled
@@ -96,6 +105,7 @@ document.addEventListener("DOMContentLoaded", function () {
modelSize: selectedModelSize, modelSize: selectedModelSize,
useVad: useVadCheckbox.checked, useVad: useVadCheckbox.checked,
saveCaptions: saveCaptionsCheckbox.checked, saveCaptions: saveCaptionsCheckbox.checked,
captionLines: Number(selectedCaptionLines),
}, () => { }, () => {
// Update capturing state in storage and toggle the buttons // Update capturing state in storage and toggle the buttons
chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => { chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => {
@@ -144,6 +154,7 @@ document.addEventListener("DOMContentLoaded", function () {
modelSizeDropdown.disabled = isCapturing; modelSizeDropdown.disabled = isCapturing;
languageDropdown.disabled = isCapturing; languageDropdown.disabled = isCapturing;
taskDropdown.disabled = isCapturing; taskDropdown.disabled = isCapturing;
captionLinesDropdown.disabled = isCapturing;
startButton.classList.toggle("disabled", isCapturing); startButton.classList.toggle("disabled", isCapturing);
stopButton.classList.toggle("disabled", !isCapturing); stopButton.classList.toggle("disabled", !isCapturing);
} }
@@ -183,6 +194,11 @@ document.addEventListener("DOMContentLoaded", function () {
chrome.storage.local.set({ selectedModelSize }); chrome.storage.local.set({ selectedModelSize });
}); });
captionLinesDropdown.addEventListener('change', function() {
selectedCaptionLines = captionLinesDropdown.value;
chrome.storage.local.set({ selectedCaptionLines });
});
chrome.runtime.onMessage.addListener(async (request, sender, sendResponse) => { chrome.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
if (request.action === "updateSelectedLanguage") { if (request.action === "updateSelectedLanguage") {
const detectedLanguage = request.detectedLanguage; const detectedLanguage = request.detectedLanguage;
+26 -23
View File
@@ -162,6 +162,7 @@ var elem_text = null;
var segments = []; var segments = [];
var text_segments = []; var text_segments = [];
var captionLineCount = 3;
function initPopupElement() { function initPopupElement() {
if (document.getElementById('popupElement')) { if (document.getElementById('popupElement')) {
@@ -209,22 +210,23 @@ function showPopup(customText) {
} }
function init_element() { function init_element(lines = 3) {
captionLineCount = Math.min(Math.max(parseInt(lines, 10) || 3, 1), 8);
if (document.getElementById('transcription')) { if (document.getElementById('transcription')) {
return; return;
} }
elem_container = document.createElement('div'); elem_container = document.createElement('div');
elem_container.id = "transcription"; elem_container.id = "transcription";
elem_container.style.cssText = 'padding-top:16px;font-size:18px;line-height:18px;position:fixed;top:85%;left:50%;transform:translate(-50%,-50%);width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;'; elem_container.style.cssText = 'padding-top:16px;font-size:18px;line-height:18px;position:fixed;top:85%;left:50%;transform:translate(-50%,-50%);width:500px;height:' + (captionLineCount * 30) + 'px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
for (var i = 0; i < 4; i++) { for (var i = 0; i <= captionLineCount; i++) {
elem_text = document.createElement('span'); elem_text = document.createElement('span');
elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;'; elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;';
elem_text.id = "t" + i; elem_text.id = "t" + i;
elem_container.appendChild(elem_text); elem_container.appendChild(elem_text);
if (i == 3) { if (i == captionLineCount) {
elem_text.style.top = "-1000px" elem_text.style.top = "-1000px"
} }
} }
@@ -286,7 +288,7 @@ function get_lines(elem, line_height) {
var divHeight = elem.offsetHeight; var divHeight = elem.offsetHeight;
var lines = divHeight / line_height; var lines = divHeight / line_height;
var original_text = elem.innerHTML; var original_text = elem.textContent;
var words = original_text.split(' '); var words = original_text.split(' ');
var segments = []; var segments = [];
@@ -296,7 +298,7 @@ function get_lines(elem, line_height) {
for (var i = 0; i < words.length; i++) for (var i = 0; i < words.length; i++)
{ {
segment += words[i] + ' '; segment += words[i] + ' ';
elem.innerHTML = segment; elem.textContent = segment;
divHeight = elem.offsetHeight; divHeight = elem.offsetHeight;
if ((divHeight / line_height) > current_lines) { if ((divHeight / line_height) > current_lines) {
@@ -310,7 +312,7 @@ function get_lines(elem, line_height) {
var line_segment = segment.substring(segment_len, segment.length - 1) var line_segment = segment.substring(segment_len, segment.length - 1)
segments.push(line_segment); segments.push(line_segment);
elem.innerHTML = original_text; elem.textContent = original_text;
return segments; return segments;
@@ -318,7 +320,7 @@ function get_lines(elem, line_height) {
function remove_element() { function remove_element() {
var elem = document.getElementById('transcription') var elem = document.getElementById('transcription')
for (var i = 0; i < 4; i++) { for (var i = 0; i <= captionLineCount; i++) {
document.getElementById("t" + i).remove(); document.getElementById("t" + i).remove();
} }
elem.remove() elem.remove()
@@ -327,6 +329,7 @@ function remove_element() {
browser.runtime.onMessage.addListener((request, sender, sendResponse) => { browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
const { action, data } = request; const { action, data } = request;
const saveCaption = data.saveCaption || false; const saveCaption = data.saveCaption || false;
const captionLines = data.captionLines || captionLineCount;
if (action === "startCapture") { if (action === "startCapture") {
isCapturing = true; isCapturing = true;
@@ -364,7 +367,7 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
} else if (action === "show_transcript"){ } else if (action === "show_transcript"){
if (!isCapturing) return; if (!isCapturing) return;
init_element(); init_element(captionLines);
message = JSON.parse(data.data); message = JSON.parse(data.data);
message = message["segments"]; message = message["segments"];
@@ -391,10 +394,10 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
} }
text = text.replace(/(\r\n|\n|\r)/gm, ""); text = text.replace(/(\r\n|\n|\r)/gm, "");
var elem = document.getElementById('t3'); var elem = document.getElementById('t' + captionLineCount);
elem.innerHTML = text; elem.textContent = text;
var line_height_style = getStyle('t3', 'line-height'); var line_height_style = getStyle('t' + captionLineCount, 'line-height');
var line_height = parseInt(line_height_style.substring(0, line_height_style.length - 2)); var line_height = parseInt(line_height_style.substring(0, line_height_style.length - 2));
var divHeight = elem.offsetHeight; var divHeight = elem.offsetHeight;
var lines = divHeight / line_height; var lines = divHeight / line_height;
@@ -402,29 +405,29 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
text_segments = []; text_segments = [];
text_segments = get_lines(elem, line_height); text_segments = get_lines(elem, line_height);
elem.innerHTML = ''; elem.textContent = '';
if (text_segments.length > 2) { if (text_segments.length > captionLineCount - 1) {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i]; document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
} }
} else { } else {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = ''; document.getElementById('t' + i).textContent = '';
} }
} }
if (text_segments.length <= 2) { if (text_segments.length <= captionLineCount - 1) {
for (var i = 0; i < text_segments.length; i++) { for (var i = 0; i < text_segments.length; i++) {
document.getElementById('t' + i).innerHTML = text_segments[i]; document.getElementById('t' + i).textContent = text_segments[i];
} }
} else { } else {
for (var i = 0; i < 3; i++) { for (var i = 0; i < captionLineCount; i++) {
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i]; document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
} }
} }
for (var i = 1; i < 3; i++) for (var i = 1; i < captionLineCount; i++)
{ {
var parent_elem = document.getElementById('t' + (i - 1)); var parent_elem = document.getElementById('t' + (i - 1));
var elem = document.getElementById('t' + i); var elem = document.getElementById('t' + i);
+8
View File
@@ -24,6 +24,14 @@
<label for="saveCaption">Download SRT file at Stop Capture</label> <label for="saveCaption">Download SRT file at Stop Capture</label>
</div> </div>
<textarea id="waitTextBox" style="display: none;"></textarea> <textarea id="waitTextBox" style="display: none;"></textarea>
<div class="dropdown-container">
<label for="captionLinesDropdown">Caption Lines:</label>
<select id="captionLinesDropdown">
<option value="3" selected>3 lines</option>
<option value="5">5 lines</option>
<option value="8">8 lines</option>
</select>
</div>
<div class="dropdown-container"> <div class="dropdown-container">
<label for="languageDropdown">Select Language:</label> <label for="languageDropdown">Select Language:</label>
<select id="languageDropdown"> <select id="languageDropdown">
+16
View File
@@ -8,9 +8,11 @@ document.addEventListener("DOMContentLoaded", function() {
const languageDropdown = document.getElementById('languageDropdown'); const languageDropdown = document.getElementById('languageDropdown');
const taskDropdown = document.getElementById('taskDropdown'); const taskDropdown = document.getElementById('taskDropdown');
const modelSizeDropdown = document.getElementById('modelSizeDropdown'); const modelSizeDropdown = document.getElementById('modelSizeDropdown');
const captionLinesDropdown = document.getElementById('captionLinesDropdown');
let selectedLanguage = null; let selectedLanguage = null;
let selectedTask = taskDropdown.value; let selectedTask = taskDropdown.value;
let selectedModelSize = modelSizeDropdown.value; let selectedModelSize = modelSizeDropdown.value;
let selectedCaptionLines = captionLinesDropdown.value;
browser.storage.local.get("capturingState") browser.storage.local.get("capturingState")
@@ -69,6 +71,13 @@ document.addEventListener("DOMContentLoaded", function() {
} }
}); });
browser.storage.local.get("selectedCaptionLines", ({ selectedCaptionLines: storedCaptionLines }) => {
if (storedCaptionLines !== undefined) {
captionLinesDropdown.value = storedCaptionLines;
selectedCaptionLines = storedCaptionLines;
}
});
startButton.addEventListener("click", function() { startButton.addEventListener("click", function() {
let host = "localhost"; let host = "localhost";
let port = "9090"; let port = "9090";
@@ -93,6 +102,7 @@ document.addEventListener("DOMContentLoaded", function() {
modelSize: selectedModelSize, modelSize: selectedModelSize,
useVad: useVadCheckbox.checked, useVad: useVadCheckbox.checked,
saveCaption: saveCaptionCheckbox.checked, saveCaption: saveCaptionCheckbox.checked,
captionLines: Number(selectedCaptionLines),
} }
}); });
toggleCaptureButtons(true); toggleCaptureButtons(true);
@@ -136,6 +146,7 @@ document.addEventListener("DOMContentLoaded", function() {
modelSizeDropdown.disabled = isCapturing; modelSizeDropdown.disabled = isCapturing;
languageDropdown.disabled = isCapturing; languageDropdown.disabled = isCapturing;
taskDropdown.disabled = isCapturing; taskDropdown.disabled = isCapturing;
captionLinesDropdown.disabled = isCapturing;
startButton.classList.toggle("disabled", isCapturing); startButton.classList.toggle("disabled", isCapturing);
stopButton.classList.toggle("disabled", !isCapturing); stopButton.classList.toggle("disabled", !isCapturing);
} }
@@ -175,6 +186,11 @@ document.addEventListener("DOMContentLoaded", function() {
browser.storage.local.set({ selectedModelSize }); browser.storage.local.set({ selectedModelSize });
}); });
captionLinesDropdown.addEventListener('change', function() {
selectedCaptionLines = captionLinesDropdown.value;
browser.storage.local.set({ selectedCaptionLines });
});
browser.runtime.onMessage.addListener((request, sender, sendResponse) => { browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
if (request.action === "updateSelectedLanguage") { if (request.action === "updateSelectedLanguage") {
const detectedLanguage = request.data; const detectedLanguage = request.data;
+79 -5
View File
@@ -17,6 +17,12 @@ input from microphone and pre-recorded audio files.
- [Getting Started](#getting-started) - [Getting Started](#getting-started)
- [Running the Server](#running-the-server) - [Running the Server](#running-the-server)
- [Running the Client](#running-the-client) - [Running the Client](#running-the-client)
- [Advanced Features](#advanced-features)
- [Word-Level Timestamps](#word-level-timestamps)
- [Custom Vocabulary / Hotwords](#custom-vocabulary--hotwords)
- [Speaker Diarization](#speaker-diarization)
- [Batch Inference](#batch-inference)
- [Raw PCM Input](#raw-pcm-input)
- [Browser Extensions](#browser-extensions) - [Browser Extensions](#browser-extensions)
- [Whisper Live Server in Docker](#whisper-live-server-in-docker) - [Whisper Live Server in Docker](#whisper-live-server-in-docker)
- [Future Work](#future-work) - [Future Work](#future-work)
@@ -25,10 +31,11 @@ input from microphone and pre-recorded audio files.
- [Citations](#citations) - [Citations](#citations)
## Installation ## Installation
- Install PortAudio - Install PortAudio (required system dependency for microphone input via PyAudio)
```bash ```bash
bash scripts/setup.sh bash scripts/setup.sh
``` ```
On Debian/Ubuntu this installs `portaudio19-dev`, on Fedora `portaudio-devel`, on macOS it uses Homebrew (`portaudio`).
- Install whisper-live from pip - Install whisper-live from pip
```bash ```bash
@@ -101,6 +108,7 @@ python3 run_server.py -p 9090 \
--max_clients 4 \ --max_clients 4 \
--max_connection_time 600 --max_connection_time 600
``` ```
> **Note:** The TensorRT backend uses a C++ session by default. If you experience issues (e.g. repeated `CrossAttentionMask` warnings or crashes), add the `--trt_py_session` flag to use the Python session instead.
- Use `--max_clients` option to restrict the number of clients the server should allow. Defaults to 4. - Use `--max_clients` option to restrict the number of clients the server should allow. Defaults to 4.
- Use `--max_connection_time` options to limit connection time for a client in seconds. Defaults to 600. - Use `--max_connection_time` options to limit connection time for a client in seconds. Defaults to 600.
- WhisperLive now supports the [OpenVINO](https://github.com/openvinotoolkit/openvino) backend for efficient inference on Intel CPUs, iGPU and dGPUs. Currently, we tested the models uploaded to [huggingface by OpenVINO](https://huggingface.co/OpenVINO?search_models=whisper). - WhisperLive now supports the [OpenVINO](https://github.com/openvinotoolkit/openvino) backend for efficient inference on Intel CPUs, iGPU and dGPUs. Currently, we tested the models uploaded to [huggingface by OpenVINO](https://huggingface.co/OpenVINO?search_models=whisper).
@@ -186,6 +194,71 @@ client(rtsp_url="rtsp://admin:admin@192.168.0.1/rtsp")
client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/bbc_1xtra.isml/bbc_1xtra-audio%3d96000.norewind.m3u8") client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/bbc_1xtra.isml/bbc_1xtra-audio%3d96000.norewind.m3u8")
``` ```
## Advanced Features
#### Word-Level Timestamps
Enable per-word timing and confidence scores in transcription segments:
```python
client = TranscriptionClient(
"localhost", 9090,
word_timestamps=True,
)
```
When enabled, each segment in the WebSocket response includes a `words` array:
```json
{
"segments": [{
"start": "0.000", "end": "2.500", "text": "Hello world",
"words": [
{"word": "Hello", "start": "0.000", "end": "0.800", "probability": 0.95},
{"word": " world", "start": "0.900", "end": "2.500", "probability": 0.88}
]
}]
}
```
#### Custom Vocabulary / Hotwords
Boost recognition of specific terms (product names, acronyms, domain jargon):
```python
client = TranscriptionClient(
"localhost", 9090,
hotwords="WhisperLive,TensorRT,OpenVINO",
)
```
The `hotwords` parameter is a comma-separated string passed directly to faster-whisper's keyword boosting. Also available in the REST API via the `hotwords` form field.
#### Speaker Diarization
Real-time speaker identification using pyannote.audio embeddings (optional dependency):
```bash
pip install pyannote.audio
```
```python
client = TranscriptionClient(
"localhost", 9090,
enable_diarization=True,
max_speakers=4,
)
```
When enabled, completed segments include a `speaker` field:
```json
{"start": "0.000", "end": "2.500", "text": "Hello", "speaker": "SPEAKER_00", "completed": true}
```
Diarization uses online cosine-similarity clustering of speaker embeddings. If `pyannote.audio` is not installed, the server logs a warning and continues without diarization.
#### Batch Inference
Batch multiple client sessions into single GPU calls for higher throughput:
```bash
python3 run_server.py --port 9090 --backend faster_whisper \
--batch_inference --batch_max_size 8 --batch_window_ms 50
```
#### Raw PCM Input
Accept raw PCM int16 audio from clients (useful for embedded devices):
```bash
python3 run_server.py --port 9090 --backend faster_whisper --raw_pcm_input
```
Audio is automatically normalized to float32 range [-1.0, 1.0].
## Browser Extensions ## Browser Extensions
- Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server). - Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server).
- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md - Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md
@@ -206,19 +279,20 @@ Refer to [`ios-client`](https://github.com/collabora/WhisperLive/tree/main/Audio
- TensorRT. Refer to [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) for setup and more tensorrt backend configurations. - TensorRT. Refer to [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) for setup and more tensorrt backend configurations.
```bash ```bash
docker build . -f docker/Dockerfile.tensorrt -t whisperlive-tensorrt docker build . -f docker/Dockerfile.tensorrt -t whisperlive-tensorrt
docker run -p 9090:9090 --runtime=nvidia --entrypoint /bin/bash -it whisperlive-tensorrt docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it whisperlive-tensorrt
# Build small.en engine # Build small.en engine
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en # float16 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en # float16
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8 # int8 weight only quantization bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8 # int8 weight only quantization
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4 # int4 weight only quantization bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4 # int4 weight only quantization
# Run server with small.en # Run server with small.en (pick one engine)
python3 run_server.py --port 9090 \ python3 run_server.py --port 9090 \
--backend tensorrt \ --backend tensorrt \
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16" --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8" # or int8 / int4:
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4" # --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
# --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
``` ```
- OpenVINO - OpenVINO
+25
View File
@@ -0,0 +1,25 @@
import sys
from whisper_live.client import TranscriptionClient
if len(sys.argv) < 2:
print("Usage: python transcribe_file.py <path_to_audio_file>")
sys.exit(1)
audio_file = sys.argv[1]
client = TranscriptionClient(
"localhost",
9090,
lang="en",
translate=False,
model="small", # also support hf_model => `Systran/faster-whisper-small`
use_vad=False,
save_output_recording=True, # Only used for microphone input, False by Default
output_recording_filename="./output_recording.wav", # Only used for microphone input
mute_audio_playback=False, # Only used for file input, False by Default
enable_translation=True,
target_language="hi",
)
# Transcribe the offline audio file
client(audio_file)
+5
View File
@@ -0,0 +1,5 @@
[pytest]
testpaths = tests
python_files = test_*.py
python_classes = Test*
python_functions = test_*
+2 -1
View File
@@ -1,6 +1,7 @@
faster-whisper==1.2.0 faster-whisper==1.2.0
websockets websockets
onnxruntime==1.17.0 onnxruntime>=1.17.0,<1.20.0; python_version < "3.10"
onnxruntime>=1.20.0,<2; python_version >= "3.10"
numba numba
kaldialign kaldialign
soundfile soundfile
+9 -2
View File
@@ -33,7 +33,8 @@ if __name__ == '__main__':
help='Language code for transcription, e.g., "en" for English.') help='Language code for transcription, e.g., "en" for English.')
parser.add_argument('--translate', '-t', parser.add_argument('--translate', '-t',
action='store_true', action='store_true',
help='Enable translation of the transcription output.') help='Use Whisper built-in translation to English (sets task=translate). '
'For any-to-any translation, use --enable_translation instead.')
parser.add_argument('--mute_audio_playback', '-a', parser.add_argument('--mute_audio_playback', '-a',
action='store_true', action='store_true',
help='Mute audio playback during transcription.') help='Mute audio playback during transcription.')
@@ -42,7 +43,7 @@ if __name__ == '__main__':
help='Save the output recording, only used for microphone input.') help='Save the output recording, only used for microphone input.')
parser.add_argument('--enable_translation', parser.add_argument('--enable_translation',
action='store_true', action='store_true',
help='Enable translation of the transcription output.') help='Enable any-to-any translation via M2M100 model (separate from Whisper --translate).')
parser.add_argument('--target_language', '-tl', parser.add_argument('--target_language', '-tl',
type=str, type=str,
default='fr', default='fr',
@@ -57,6 +58,12 @@ if __name__ == '__main__':
args = parser.parse_args() args = parser.parse_args()
if args.translate and args.enable_translation:
print("[WARN]: Both --translate and --enable_translation are set. "
"--translate uses Whisper's built-in to-English translation, "
"while --enable_translation uses M2M100 for any-to-any. "
"Both will be active.")
client = TranscriptionClient( client = TranscriptionClient(
args.server, args.server,
args.port, args.port,
+29
View File
@@ -84,6 +84,31 @@ if __name__ == "__main__":
default=50, default=50,
help='Maximum time in ms to wait for batch to fill (default: 50).' help='Maximum time in ms to wait for batch to fill (default: 50).'
) )
parser.add_argument(
'--raw_pcm_input',
action='store_true',
help='Expect raw PCM int16 audio from clients instead of float32. '
'Audio will be normalized to float32 range [-1.0, 1.0].'
)
parser.add_argument(
'--metrics_port',
type=int,
default=0,
help='Port for Prometheus /metrics endpoint. 0 = disabled (default). Requires prometheus_client.'
)
parser.add_argument(
'--api_key',
type=str,
default=None,
help='Optional API key for authenticating REST API and WebSocket connections. '
'Clients must send "Authorization: Bearer <key>" header or "?token=<key>" query parameter.'
)
parser.add_argument(
'--rate_limit_rpm',
type=int,
default=0,
help='Maximum REST API requests per minute per client IP. 0 = unlimited (default).'
)
args = parser.parse_args() args = parser.parse_args()
if args.backend == "tensorrt": if args.backend == "tensorrt":
@@ -113,4 +138,8 @@ if __name__ == "__main__":
batch_enabled=args.batch_inference, batch_enabled=args.batch_inference,
batch_max_size=args.batch_max_size, batch_max_size=args.batch_max_size,
batch_window_ms=args.batch_window_ms, batch_window_ms=args.batch_window_ms,
raw_pcm_input=args.raw_pcm_input,
metrics_port=args.metrics_port,
api_key=args.api_key,
rate_limit_rpm=args.rate_limit_rpm,
) )
+10 -2
View File
@@ -28,8 +28,11 @@ setup(
"License :: OSI Approved :: MIT License", "License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3", "Programming Language :: Python :: 3",
"Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3 :: Only",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Artificial Intelligence",
], ],
packages=find_packages( packages=find_packages(
@@ -43,11 +46,13 @@ setup(
), ),
install_requires=[ install_requires=[
"PyAudio", "PyAudio",
"av",
"faster-whisper==1.2.0", "faster-whisper==1.2.0",
"torch", "torch",
"torchaudio", "torchaudio",
"websockets", "websockets",
"onnxruntime==1.17.0", "onnxruntime>=1.17.0,<1.20.0; python_version < '3.10'",
"onnxruntime>=1.20.0,<2; python_version >= '3.10'",
"scipy", "scipy",
"websocket-client", "websocket-client",
"numba", "numba",
@@ -62,6 +67,9 @@ setup(
"openvino-tokenizers", "openvino-tokenizers",
"optimum", "optimum",
"optimum-intel", "optimum-intel",
"fastapi",
"uvicorn",
"python-multipart",
], ],
python_requires=">=3.9" python_requires=">=3.9"
) )
+519
View File
@@ -0,0 +1,519 @@
import json
import queue
import threading
import time
import unittest
from unittest.mock import MagicMock, patch
import numpy as np
from whisper_live.backend.base import ServeClientBase
class ConcreteServeClient(ServeClientBase):
"""Concrete subclass for testing the abstract base class."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.language = "en"
def transcribe_audio(self, input_sample):
return None
def handle_transcription_output(self, result, duration):
pass
class TestServeClientBaseInit(unittest.TestCase):
def test_default_values(self):
ws = MagicMock()
client = ConcreteServeClient(client_uid="test-uid", websocket=ws)
self.assertEqual(client.client_uid, "test-uid")
self.assertEqual(client.send_last_n_segments, 10)
self.assertAlmostEqual(client.no_speech_thresh, 0.45)
self.assertFalse(client.clip_audio)
self.assertEqual(client.same_output_threshold, 10)
self.assertIsNone(client.frames_np)
self.assertAlmostEqual(client.timestamp_offset, 0.0)
self.assertFalse(client.exit)
self.assertEqual(client.transcript, [])
def test_custom_values(self):
ws = MagicMock()
q = queue.Queue()
client = ConcreteServeClient(
client_uid="uid2",
websocket=ws,
send_last_n_segments=5,
no_speech_thresh=0.6,
clip_audio=True,
same_output_threshold=20,
translation_queue=q,
)
self.assertEqual(client.send_last_n_segments, 5)
self.assertAlmostEqual(client.no_speech_thresh, 0.6)
self.assertTrue(client.clip_audio)
self.assertEqual(client.same_output_threshold, 20)
self.assertIs(client.translation_queue, q)
class TestAddFrames(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
def test_first_frame_initializes_buffer(self):
frame = np.array([0.1, 0.2, 0.3], dtype=np.float32)
self.client.add_frames(frame)
np.testing.assert_array_equal(self.client.frames_np, frame)
def test_subsequent_frames_concatenated(self):
frame1 = np.array([0.1, 0.2], dtype=np.float32)
frame2 = np.array([0.3, 0.4], dtype=np.float32)
self.client.add_frames(frame1)
self.client.add_frames(frame2)
expected = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
np.testing.assert_array_equal(self.client.frames_np, expected)
def test_buffer_trimmed_at_45_seconds(self):
# 45 seconds + 1 sample at 16kHz = 720001 samples
self.client.frames_np = np.zeros(45 * 16000 + 1, dtype=np.float32)
self.client.add_frames(np.array([1.0], dtype=np.float32))
# after trimming 30s, buffer should be ~15s + 1 original + 1 new
expected_len = (45 * 16000 + 1) - (30 * 16000) + 1
self.assertEqual(self.client.frames_np.shape[0], expected_len)
self.assertAlmostEqual(self.client.frames_offset, 30.0)
def test_timestamp_offset_updated_on_trim(self):
self.client.frames_np = np.zeros(45 * 16000 + 1, dtype=np.float32)
self.client.timestamp_offset = 5.0 # behind frames_offset after trim
self.client.add_frames(np.array([1.0], dtype=np.float32))
# timestamp_offset should be bumped to at least frames_offset
self.assertGreaterEqual(self.client.timestamp_offset, self.client.frames_offset)
class TestAddFramesThreadSafety(unittest.TestCase):
def test_concurrent_add_frames(self):
ws = MagicMock()
client = ConcreteServeClient(client_uid="test", websocket=ws)
errors = []
def add_many():
try:
for _ in range(100):
client.add_frames(np.random.randn(160).astype(np.float32))
except Exception as e:
errors.append(e)
threads = [threading.Thread(target=add_many) for _ in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(errors, [])
self.assertIsNotNone(client.frames_np)
class TestGetAudioChunkForProcessing(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
def test_empty_buffer_returns_empty(self):
self.client.frames_np = np.array([], dtype=np.float32)
chunk, duration = self.client.get_audio_chunk_for_processing()
self.assertEqual(duration, 0.0)
self.assertEqual(chunk.shape[0], 0)
def test_full_buffer_no_offset(self):
audio = np.random.randn(16000).astype(np.float32) # 1 second
self.client.frames_np = audio
chunk, duration = self.client.get_audio_chunk_for_processing()
self.assertAlmostEqual(duration, 1.0)
np.testing.assert_array_equal(chunk, audio)
def test_with_offset(self):
audio = np.random.randn(32000).astype(np.float32) # 2 seconds
self.client.frames_np = audio
self.client.timestamp_offset = 1.0 # skip first second
chunk, duration = self.client.get_audio_chunk_for_processing()
self.assertAlmostEqual(duration, 1.0)
self.assertEqual(chunk.shape[0], 16000)
class TestClipAudioIfNoValidSegment(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(
client_uid="test", websocket=self.ws, clip_audio=True
)
def test_clips_when_chunk_exceeds_25s(self):
# 30 seconds of audio with no valid segments
self.client.frames_np = np.zeros(30 * 16000, dtype=np.float32)
self.client.timestamp_offset = 0.0
self.client.frames_offset = 0.0
self.client.clip_audio_if_no_valid_segment()
# offset should have advanced to leave ~5s of remaining audio
expected_offset = (30 * 16000 / 16000) - 5
self.assertAlmostEqual(self.client.timestamp_offset, expected_offset, places=1)
def test_no_clip_when_short(self):
self.client.frames_np = np.zeros(10 * 16000, dtype=np.float32)
self.client.timestamp_offset = 0.0
self.client.frames_offset = 0.0
self.client.clip_audio_if_no_valid_segment()
self.assertAlmostEqual(self.client.timestamp_offset, 0.0)
class TestPrepareSegments(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(
client_uid="test", websocket=self.ws, send_last_n_segments=3
)
def test_empty_transcript_no_last(self):
segments = self.client.prepare_segments()
self.assertEqual(segments, [])
def test_empty_transcript_with_last(self):
last = {"start": "0.000", "end": "1.000", "text": "hello", "completed": False}
segments = self.client.prepare_segments(last_segment=last)
self.assertEqual(len(segments), 1)
self.assertEqual(segments[0]["text"], "hello")
def test_fewer_than_n_segments(self):
self.client.transcript = [
{"start": "0.000", "end": "1.000", "text": "a", "completed": True},
{"start": "1.000", "end": "2.000", "text": "b", "completed": True},
]
segments = self.client.prepare_segments()
self.assertEqual(len(segments), 2)
def test_more_than_n_segments_truncated(self):
self.client.transcript = [
{"start": f"{i}.000", "end": f"{i+1}.000", "text": f"seg{i}", "completed": True}
for i in range(10)
]
segments = self.client.prepare_segments()
self.assertEqual(len(segments), 3)
self.assertEqual(segments[0]["text"], "seg7")
def test_last_segment_appended(self):
self.client.transcript = [
{"start": "0.000", "end": "1.000", "text": "a", "completed": True},
]
last = {"start": "1.000", "end": "2.000", "text": "in progress", "completed": False}
segments = self.client.prepare_segments(last_segment=last)
self.assertEqual(len(segments), 2)
self.assertEqual(segments[-1]["text"], "in progress")
class TestFormatSegment(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
def test_format(self):
seg = self.client.format_segment(1.234, 5.678, "hello world", completed=True)
self.assertEqual(seg["start"], "1.234")
self.assertEqual(seg["end"], "5.678")
self.assertEqual(seg["text"], "hello world")
self.assertTrue(seg["completed"])
def test_format_not_completed(self):
seg = self.client.format_segment(0.0, 1.0, "text")
self.assertFalse(seg["completed"])
class TestSendTranscriptionToClient(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test-uid", websocket=self.ws)
def test_sends_json(self):
segments = [{"start": "0.000", "end": "1.000", "text": "hi", "completed": True}]
self.client.send_transcription_to_client(segments)
self.ws.send.assert_called_once()
sent = json.loads(self.ws.send.call_args[0][0])
self.assertEqual(sent["uid"], "test-uid")
self.assertEqual(len(sent["segments"]), 1)
def test_send_failure_logged_not_raised(self):
self.ws.send.side_effect = ConnectionError("broken pipe")
# should not raise
self.client.send_transcription_to_client([])
class TestDisconnect(unittest.TestCase):
def test_sends_disconnect_message(self):
ws = MagicMock()
client = ConcreteServeClient(client_uid="uid1", websocket=ws)
client.disconnect()
sent = json.loads(ws.send.call_args[0][0])
self.assertEqual(sent["uid"], "uid1")
self.assertEqual(sent["message"], "DISCONNECT")
class TestCleanup(unittest.TestCase):
def test_sets_exit_flag(self):
ws = MagicMock()
client = ConcreteServeClient(client_uid="uid1", websocket=ws)
self.assertFalse(client.exit)
client.cleanup()
self.assertTrue(client.exit)
class TestTrimTranscript(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
def test_transcript_trimmed_when_over_max(self):
self.client.transcript = [
{"start": f"{i}.000", "end": f"{i+1}.000", "text": f"seg{i}", "completed": True}
for i in range(self.client.MAX_TRANSCRIPT_LENGTH + 100)
]
self.client._trim_transcript()
self.assertEqual(len(self.client.transcript), self.client.MAX_TRANSCRIPT_LENGTH)
self.assertEqual(self.client.transcript[0]["text"], "seg100")
def test_transcript_not_trimmed_when_under_max(self):
self.client.transcript = [
{"start": "0.000", "end": "1.000", "text": "a", "completed": True}
]
self.client._trim_transcript()
self.assertEqual(len(self.client.transcript), 1)
def test_text_list_trimmed(self):
self.client.text = ["word"] * (self.client.MAX_TRANSCRIPT_LENGTH + 50)
self.client._trim_transcript()
self.assertEqual(len(self.client.text), self.client.MAX_TRANSCRIPT_LENGTH)
class TestUpdateSegments(unittest.TestCase):
"""Tests for the core update_segments() logic."""
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(
client_uid="test",
websocket=self.ws,
no_speech_thresh=0.45,
same_output_threshold=3,
)
self.client.frames_np = np.zeros(16000 * 5, dtype=np.float32)
def _make_segment(self, start, end, text, no_speech_prob=0.0):
seg = MagicMock()
seg.start = start
seg.end = end
seg.text = text
seg.no_speech_prob = no_speech_prob
return seg
def test_single_segment_becomes_last(self):
segs = [self._make_segment(0.0, 1.0, " hello")]
last = self.client.update_segments(segs, duration=2.0)
self.assertIsNotNone(last)
self.assertIn("hello", last["text"])
self.assertFalse(last["completed"])
self.assertEqual(len(self.client.transcript), 0)
def test_multiple_segments_completes_all_but_last(self):
segs = [
self._make_segment(0.0, 1.0, " first"),
self._make_segment(1.0, 2.0, " second"),
]
last = self.client.update_segments(segs, duration=3.0)
self.assertEqual(len(self.client.transcript), 1)
self.assertTrue(self.client.transcript[0]["completed"])
self.assertIn("first", self.client.transcript[0]["text"])
self.assertIsNotNone(last)
self.assertIn("second", last["text"])
def test_high_no_speech_prob_skipped(self):
segs = [
self._make_segment(0.0, 1.0, " noise", no_speech_prob=0.9),
self._make_segment(1.0, 2.0, " also noise", no_speech_prob=0.9),
]
last = self.client.update_segments(segs, duration=3.0)
self.assertEqual(len(self.client.transcript), 0)
self.assertIsNone(last)
def test_segment_with_start_gte_end_skipped(self):
segs = [
self._make_segment(1.0, 0.5, " backwards"),
self._make_segment(1.5, 2.0, " normal"),
]
last = self.client.update_segments(segs, duration=3.0)
self.assertEqual(len(self.client.transcript), 0)
self.assertIsNotNone(last)
def test_repeated_output_triggers_completion(self):
seg = self._make_segment(0.0, 1.0, " repeated")
for _ in range(self.client.same_output_threshold + 2):
last = self.client.update_segments([seg], duration=2.0)
# after enough repeats, should be added to transcript
self.assertTrue(len(self.client.transcript) >= 1)
def test_translation_queue_receives_completed(self):
q = queue.Queue()
self.client.translation_queue = q
segs = [
self._make_segment(0.0, 1.0, " first"),
self._make_segment(1.0, 2.0, " second"),
]
self.client.update_segments(segs, duration=3.0)
self.assertFalse(q.empty())
item = q.get_nowait()
self.assertIn("first", item["text"])
def test_timestamp_offset_advances(self):
segs = [
self._make_segment(0.0, 1.0, " first"),
self._make_segment(1.0, 2.0, " second"),
]
self.client.update_segments(segs, duration=3.0)
self.assertGreater(self.client.timestamp_offset, 0.0)
class TestGetSegmentHelpers(unittest.TestCase):
def setUp(self):
self.ws = MagicMock()
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
def test_get_segment_no_speech_prob_attr(self):
seg = MagicMock()
seg.no_speech_prob = 0.3
self.assertAlmostEqual(self.client.get_segment_no_speech_prob(seg), 0.3)
def test_get_segment_no_speech_prob_fallback(self):
seg = MagicMock(spec=[]) # no attributes
self.assertEqual(self.client.get_segment_no_speech_prob(seg), 0)
def test_get_segment_start_uses_start(self):
seg = MagicMock()
seg.start = 1.5
self.assertAlmostEqual(self.client.get_segment_start(seg), 1.5)
def test_get_segment_end_uses_end(self):
seg = MagicMock()
seg.end = 3.0
self.assertAlmostEqual(self.client.get_segment_end(seg), 3.0)
def test_get_segment_start_fallback_to_start_ts(self):
seg = MagicMock(spec=["start_ts"])
seg.start_ts = 2.0
self.assertAlmostEqual(self.client.get_segment_start(seg), 2.0)
class TestWordTimestamps(unittest.TestCase):
"""Tests for word-level timestamp extraction."""
def _make_client(self, word_timestamps=False):
ws = MagicMock()
return ConcreteServeClient(
client_uid="wt-uid", websocket=ws, word_timestamps=word_timestamps
)
def _make_word(self, word, start, end, prob):
w = MagicMock()
w.word = word
w.start = start
w.end = end
w.probability = prob
return w
def _make_segment(self, text, start, end, no_speech_prob=0.0, words=None):
seg = MagicMock()
seg.text = text
seg.start = start
seg.end = end
seg.no_speech_prob = no_speech_prob
seg.words = words
return seg
def test_word_timestamps_disabled_by_default(self):
client = self._make_client()
self.assertFalse(client.word_timestamps)
def test_word_timestamps_enabled(self):
client = self._make_client(word_timestamps=True)
self.assertTrue(client.word_timestamps)
def test_extract_words_when_disabled(self):
client = self._make_client(word_timestamps=False)
seg = self._make_segment("hello", 0.0, 1.0, words=[self._make_word("hello", 0.0, 0.5, 0.99)])
result = client._extract_words(seg, 0.0)
self.assertIsNone(result)
def test_extract_words_when_enabled(self):
client = self._make_client(word_timestamps=True)
words = [
self._make_word("hello", 0.0, 0.3, 0.95),
self._make_word("world", 0.4, 0.8, 0.88),
]
seg = self._make_segment("hello world", 0.0, 1.0, words=words)
result = client._extract_words(seg, 10.0)
self.assertEqual(len(result), 2)
self.assertEqual(result[0]["word"], "hello")
self.assertEqual(result[0]["start"], "10.000")
self.assertEqual(result[0]["end"], "10.300")
self.assertEqual(result[0]["probability"], 0.95)
self.assertEqual(result[1]["word"], "world")
self.assertEqual(result[1]["start"], "10.400")
def test_extract_words_no_words_on_segment(self):
client = self._make_client(word_timestamps=True)
seg = self._make_segment("hello", 0.0, 1.0, words=None)
result = client._extract_words(seg, 0.0)
self.assertIsNone(result)
def test_format_segment_without_words(self):
client = self._make_client()
seg = client.format_segment(0.0, 1.0, "hello")
self.assertNotIn("words", seg)
def test_format_segment_with_words(self):
client = self._make_client(word_timestamps=True)
words = [{"word": "hello", "start": "0.000", "end": "0.500", "probability": 0.95}]
seg = client.format_segment(0.0, 1.0, "hello", words=words)
self.assertIn("words", seg)
self.assertEqual(len(seg["words"]), 1)
self.assertEqual(seg["words"][0]["word"], "hello")
def test_update_segments_includes_words(self):
client = self._make_client(word_timestamps=True)
words1 = [self._make_word("hello", 0.0, 0.5, 0.9)]
words2 = [self._make_word("world", 0.6, 1.0, 0.85)]
segments = [
self._make_segment(" hello", 0.0, 0.5, words=words1),
self._make_segment(" world", 0.6, 1.0, words=words2),
]
last = client.update_segments(segments, 2.0)
# First segment should be completed (in transcript) with words
self.assertTrue(len(client.transcript) > 0)
self.assertIn("words", client.transcript[-1])
# Last segment should be in-progress with words
self.assertIsNotNone(last)
self.assertIn("words", last)
def test_update_segments_no_words_when_disabled(self):
client = self._make_client(word_timestamps=False)
words1 = [self._make_word("hello", 0.0, 0.5, 0.9)]
words2 = [self._make_word("world", 0.6, 1.0, 0.85)]
segments = [
self._make_segment(" hello", 0.0, 0.5, words=words1),
self._make_segment(" world", 0.6, 1.0, words=words2),
]
last = client.update_segments(segments, 2.0)
self.assertTrue(len(client.transcript) > 0)
self.assertNotIn("words", client.transcript[-1])
self.assertNotIn("words", last)
if __name__ == "__main__":
unittest.main()
+18 -14
View File
@@ -43,21 +43,25 @@ class TestClientWebSocketCommunication(BaseTestCase):
class TestClientCallbacks(BaseTestCase): class TestClientCallbacks(BaseTestCase):
def test_on_open(self): def test_on_open(self):
expected_message = json.dumps({
"uid": self.client.uid,
"language": self.client.language,
"task": self.client.task,
"model": self.client.model,
"use_vad": True,
"send_last_n_segments": 10,
"no_speech_thresh": 0.45,
"clip_audio": False,
"same_output_threshold": 10,
"enable_translation": False,
"target_language": "fr",
})
self.client.on_open(self.mock_ws_app) self.client.on_open(self.mock_ws_app)
self.mock_ws_app.send.assert_called_with(expected_message) self.mock_ws_app.send.assert_called_once()
sent_message = json.loads(self.mock_ws_app.send.call_args[0][0])
self.assertEqual(sent_message["uid"], self.client.uid)
self.assertEqual(sent_message["language"], self.client.language)
self.assertEqual(sent_message["task"], self.client.task)
self.assertEqual(sent_message["model"], self.client.model)
self.assertTrue(sent_message["use_vad"])
self.assertEqual(sent_message["send_last_n_segments"], 10)
self.assertAlmostEqual(sent_message["no_speech_thresh"], 0.45)
self.assertFalse(sent_message["clip_audio"])
self.assertEqual(sent_message["same_output_threshold"], 10)
self.assertFalse(sent_message["enable_translation"])
self.assertEqual(sent_message["target_language"], "fr")
self.assertIsNone(sent_message["hotwords"])
self.assertFalse(sent_message["enable_diarization"])
self.assertEqual(sent_message["max_speakers"], 10)
self.assertFalse(sent_message["word_timestamps"])
def test_on_message(self): def test_on_message(self):
message = json.dumps( message = json.dumps(
+305
View File
@@ -0,0 +1,305 @@
import json
import time
import unittest
from unittest.mock import patch, MagicMock, PropertyMock
from whisper_live.client import Client, TranscriptionTeeClient
class TestClientStatusMessages(unittest.TestCase):
"""Tests for Client.handle_status_messages() and on_message() branches."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def setUp(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
self.client = Client(host="localhost", port=9090, lang="en")
def tearDown(self):
self.client.close_websocket()
def test_wait_status(self):
msg = {"uid": self.client.uid, "status": "WAIT", "message": 5.0}
self.client.handle_status_messages(msg)
self.assertTrue(self.client.waiting)
def test_error_status(self):
msg = {"uid": self.client.uid, "status": "ERROR", "message": "model not found"}
self.client.handle_status_messages(msg)
self.assertTrue(self.client.server_error)
def test_warning_status_no_side_effects(self):
msg = {"uid": self.client.uid, "status": "WARNING", "message": "fallback backend"}
self.client.handle_status_messages(msg)
self.assertFalse(self.client.server_error)
self.assertFalse(self.client.waiting)
def test_on_message_wrong_uid_ignored(self):
msg = json.dumps({"uid": "wrong-uid", "segments": [{"start": 0, "end": 1, "text": "hi", "completed": True}]})
self.client.on_message(MagicMock(), msg)
self.assertEqual(len(self.client.transcript), 0)
def test_on_message_disconnect(self):
self.client.recording = True
msg = json.dumps({"uid": self.client.uid, "message": "DISCONNECT"})
self.client.on_message(MagicMock(), msg)
self.assertFalse(self.client.recording)
def test_on_message_server_ready(self):
msg = json.dumps({
"uid": self.client.uid,
"message": "SERVER_READY",
"backend": "faster_whisper",
})
self.client.on_message(MagicMock(), msg)
self.assertTrue(self.client.recording)
self.assertEqual(self.client.server_backend, "faster_whisper")
def test_on_message_language_detection(self):
msg = json.dumps({
"uid": self.client.uid,
"language": "fr",
"language_prob": 0.95,
})
self.client.on_message(MagicMock(), msg)
self.assertEqual(self.client.language, "fr")
class TestClientTranslationFlow(unittest.TestCase):
"""Tests for the translation-related client functionality."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def setUp(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
self.client = Client(
host="localhost",
port=9090,
lang="en",
enable_translation=True,
target_language="es",
)
# simulate SERVER_READY so server_backend is set
ready_msg = json.dumps({
"uid": self.client.uid,
"message": "SERVER_READY",
"backend": "faster_whisper",
})
self.client.on_message(MagicMock(), ready_msg)
def tearDown(self):
self.client.close_websocket()
def test_on_open_includes_translation_fields(self):
mock_ws = MagicMock()
self.client.on_open(mock_ws)
sent = json.loads(mock_ws.send.call_args[0][0])
self.assertTrue(sent["enable_translation"])
self.assertEqual(sent["target_language"], "es")
def test_translated_segments_processed(self):
msg = json.dumps({
"uid": self.client.uid,
"translated_segments": [
{"start": "0.000", "end": "1.000", "text": "Hola mundo", "completed": True},
],
})
self.client.on_message(MagicMock(), msg)
self.assertEqual(len(self.client.translated_transcript), 1)
self.assertEqual(self.client.translated_transcript[0]["text"], "Hola mundo")
def test_translation_callback_invoked(self):
callback = MagicMock()
self.client.translation_callback = callback
msg = json.dumps({
"uid": self.client.uid,
"translated_segments": [
{"start": "0.000", "end": "1.000", "text": "Hola", "completed": True},
],
})
self.client.on_message(MagicMock(), msg)
callback.assert_called_once()
def test_translation_callback_exception_handled(self):
callback = MagicMock(side_effect=RuntimeError("callback broke"))
self.client.translation_callback = callback
msg = json.dumps({
"uid": self.client.uid,
"translated_segments": [
{"start": "0.000", "end": "1.000", "text": "Hola", "completed": True},
],
})
# should not raise
self.client.on_message(MagicMock(), msg)
class TestClientTranscriptionCallback(unittest.TestCase):
"""Tests for the transcription callback feature."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def setUp(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
self.callback = MagicMock()
self.client = Client(
host="localhost",
port=9090,
lang="en",
transcription_callback=self.callback,
)
ready_msg = json.dumps({
"uid": self.client.uid,
"message": "SERVER_READY",
"backend": "faster_whisper",
})
self.client.on_message(MagicMock(), ready_msg)
def tearDown(self):
self.client.close_websocket()
def test_callback_receives_text_and_segments(self):
msg = json.dumps({
"uid": self.client.uid,
"segments": [
{"start": "0.000", "end": "1.000", "text": "Hello", "completed": True},
],
})
self.client.on_message(MagicMock(), msg)
self.callback.assert_called_once()
text_arg, segments_arg = self.callback.call_args[0]
self.assertIn("Hello", text_arg)
self.assertIsInstance(segments_arg, list)
def test_callback_exception_does_not_crash(self):
self.callback.side_effect = ValueError("boom")
msg = json.dumps({
"uid": self.client.uid,
"segments": [
{"start": "0.000", "end": "1.000", "text": "Test", "completed": True},
],
})
# should not raise
self.client.on_message(MagicMock(), msg)
class TestClientSrtWriting(unittest.TestCase):
"""Tests for Client.write_srt_file() edge cases."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def setUp(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
self.client = Client(host="localhost", port=9090, lang="en")
self.client.server_backend = "faster_whisper"
def tearDown(self):
self.client.close_websocket()
import os
for f in ["test_out.srt"]:
if os.path.exists(f):
os.remove(f)
def test_write_srt_empty_transcript_with_last_segment(self):
self.client.transcript = []
self.client.last_segment = {"start": "0.000", "end": "1.000", "text": "final"}
self.client.write_srt_file("test_out.srt")
self.assertEqual(len(self.client.transcript), 1)
self.assertEqual(self.client.transcript[0]["text"], "final")
def test_write_srt_appends_last_segment_if_different(self):
self.client.transcript = [{"start": "0.000", "end": "1.000", "text": "first"}]
self.client.last_segment = {"start": "1.000", "end": "2.000", "text": "second"}
self.client.write_srt_file("test_out.srt")
self.assertEqual(len(self.client.transcript), 2)
def test_write_srt_no_duplicate_last_segment(self):
self.client.transcript = [{"start": "0.000", "end": "1.000", "text": "same"}]
self.client.last_segment = {"start": "0.000", "end": "1.000", "text": "same"}
self.client.write_srt_file("test_out.srt")
self.assertEqual(len(self.client.transcript), 1)
class TestWaitBeforeDisconnect(unittest.TestCase):
"""Tests for Client.wait_before_disconnect()."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def setUp(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
self.client = Client(host="localhost", port=9090, lang="en")
def tearDown(self):
self.client.close_websocket()
def test_raises_if_no_response(self):
self.client.last_response_received = None
with self.assertRaises(AssertionError):
self.client.wait_before_disconnect()
def test_returns_immediately_if_timeout_elapsed(self):
self.client.last_response_received = time.time() - 100
self.client.disconnect_if_no_response_for = 15
start = time.time()
self.client.wait_before_disconnect()
elapsed = time.time() - start
self.assertLess(elapsed, 1.0)
class TestTeeClientEdgeCases(unittest.TestCase):
"""Edge cases for TranscriptionTeeClient."""
def test_empty_clients_raises(self):
with self.assertRaises(Exception):
TranscriptionTeeClient([])
class TestClientReconnect(unittest.TestCase):
"""Tests for reconnection logic."""
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def test_reconnect_on_close(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
client = Client(host="localhost", port=9090, lang="en", max_retries=2, retry_delay=0)
initial_socket = client.client_socket
client.on_close(MagicMock(), 1006, "abnormal closure")
self.assertEqual(client._retry_count, 1)
# A new websocket should have been created
self.assertIsNotNone(client.client_socket)
client.close_websocket()
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def test_no_reconnect_on_server_error(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
client = Client(host="localhost", port=9090, lang="en", max_retries=2, retry_delay=0)
client.server_error = True
client.on_close(MagicMock(), 1000, "normal")
self.assertEqual(client._retry_count, 0)
client.close_websocket()
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def test_no_reconnect_when_max_retries_zero(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
client = Client(host="localhost", port=9090, lang="en", max_retries=0, retry_delay=0)
client.on_close(MagicMock(), 1006, "abnormal closure")
self.assertEqual(client._retry_count, 0)
client.close_websocket()
@patch("whisper_live.client.websocket.WebSocketApp")
@patch("whisper_live.client.pyaudio.PyAudio")
def test_stops_after_max_retries(self, mock_pyaudio, mock_websocket):
mock_pyaudio.return_value.open.return_value = MagicMock()
client = Client(host="localhost", port=9090, lang="en", max_retries=2, retry_delay=0)
client.on_close(MagicMock(), 1006, "closed")
client.on_close(MagicMock(), 1006, "closed")
self.assertEqual(client._retry_count, 2)
# third close should NOT retry
client.on_close(MagicMock(), 1006, "closed")
self.assertEqual(client._retry_count, 2)
client.close_websocket()
if __name__ == "__main__":
unittest.main()
+152
View File
@@ -0,0 +1,152 @@
import unittest
from unittest.mock import MagicMock, patch
import numpy as np
class TestSpeakerDiarizer(unittest.TestCase):
"""Tests for SpeakerDiarizer with mocked embedding model."""
def _make_diarizer(self, **kwargs):
from whisper_live.diarization import SpeakerDiarizer
d = SpeakerDiarizer(**kwargs)
# Mock the embedding model to return deterministic embeddings
d._model = MagicMock()
return d
def _set_embedding(self, diarizer, embedding):
"""Configure mock model to return a specific embedding."""
emb = np.array(embedding, dtype=np.float32)
emb = emb / np.linalg.norm(emb)
diarizer._model.return_value = emb
def test_first_speaker_creates_new(self):
d = self._make_diarizer()
self._set_embedding(d, [1.0, 0.0, 0.0])
audio = np.zeros(16000, dtype=np.float32) # 1 second of audio
speaker = d.identify_speaker(audio)
self.assertEqual(speaker, "SPEAKER_00")
self.assertEqual(len(d.speakers), 1)
def test_same_speaker_matches(self):
d = self._make_diarizer(similarity_threshold=0.8)
self._set_embedding(d, [1.0, 0.0, 0.0])
audio = np.zeros(16000, dtype=np.float32)
d.identify_speaker(audio) # SPEAKER_00
# Same embedding should match
self._set_embedding(d, [0.99, 0.01, 0.0])
speaker = d.identify_speaker(audio)
self.assertEqual(speaker, "SPEAKER_00")
self.assertEqual(len(d.speakers), 1)
def test_different_speaker_creates_new(self):
d = self._make_diarizer(similarity_threshold=0.8)
self._set_embedding(d, [1.0, 0.0, 0.0])
audio = np.zeros(16000, dtype=np.float32)
d.identify_speaker(audio) # SPEAKER_00
# Very different embedding
self._set_embedding(d, [0.0, 1.0, 0.0])
speaker = d.identify_speaker(audio)
self.assertEqual(speaker, "SPEAKER_01")
self.assertEqual(len(d.speakers), 2)
def test_max_speakers_limit(self):
d = self._make_diarizer(similarity_threshold=0.95, max_speakers=2)
audio = np.zeros(16000, dtype=np.float32)
self._set_embedding(d, [1.0, 0.0, 0.0])
d.identify_speaker(audio) # SPEAKER_00
self._set_embedding(d, [0.0, 1.0, 0.0])
d.identify_speaker(audio) # SPEAKER_01
# Third distinct speaker should be assigned to closest existing
self._set_embedding(d, [0.0, 0.0, 1.0])
speaker = d.identify_speaker(audio)
self.assertIn(speaker, ["SPEAKER_00", "SPEAKER_01"])
self.assertEqual(len(d.speakers), 2)
def test_short_audio_returns_none(self):
d = self._make_diarizer()
# Less than 0.3 seconds
audio = np.zeros(3000, dtype=np.float32)
speaker = d.identify_speaker(audio)
self.assertIsNone(speaker)
def test_reset_clears_state(self):
d = self._make_diarizer()
self._set_embedding(d, [1.0, 0.0, 0.0])
audio = np.zeros(16000, dtype=np.float32)
d.identify_speaker(audio)
self.assertEqual(len(d.speakers), 1)
d.reset()
self.assertEqual(len(d.speakers), 0)
self.assertEqual(d._speaker_count, 0)
def test_import_error_without_pyannote(self):
from whisper_live.diarization import SpeakerDiarizer
d = SpeakerDiarizer()
with patch.dict("sys.modules", {"pyannote": None, "pyannote.audio": None}):
with self.assertRaises(ImportError):
d._load_model()
class TestDiarizationInBase(unittest.TestCase):
"""Test diarization integration in ServeClientBase."""
def _make_client(self, diarization=None):
from whisper_live.backend.base import ServeClientBase
class ConcreteClient(ServeClientBase):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.language = "en"
def transcribe_audio(self, input_sample):
return None
def handle_transcription_output(self, result, duration):
pass
ws = MagicMock()
return ConcreteClient(
client_uid="test-uid", websocket=ws, diarization=diarization
)
def test_no_diarization_by_default(self):
client = self._make_client()
self.assertIsNone(client.diarization)
def test_format_segment_with_speaker(self):
client = self._make_client()
seg = client.format_segment(0.0, 1.0, "hello", speaker="SPEAKER_00")
self.assertEqual(seg["speaker"], "SPEAKER_00")
def test_format_segment_without_speaker(self):
client = self._make_client()
seg = client.format_segment(0.0, 1.0, "hello")
self.assertNotIn("speaker", seg)
def test_identify_speaker_disabled(self):
client = self._make_client(diarization=None)
seg = MagicMock()
seg.start = 0.0
seg.end = 1.0
result = client._identify_speaker(seg)
self.assertIsNone(result)
def test_identify_speaker_calls_diarizer(self):
mock_diarizer = MagicMock()
mock_diarizer.identify_speaker.return_value = "SPEAKER_01"
client = self._make_client(diarization=mock_diarizer)
# Set up audio buffer
client.frames_np = np.zeros(48000, dtype=np.float32)
client.frames_offset = 0.0
client.timestamp_offset = 0.0
seg = MagicMock()
seg.start = 0.5
seg.end = 1.5
result = client._identify_speaker(seg)
self.assertEqual(result, "SPEAKER_01")
mock_diarizer.identify_speaker.assert_called_once()
if __name__ == "__main__":
unittest.main()
+138
View File
@@ -0,0 +1,138 @@
import unittest
from unittest.mock import patch, MagicMock
from whisper_live import metrics as wl_metrics
_skip_no_prometheus = unittest.skipUnless(
wl_metrics.is_available(), "prometheus_client not installed"
)
class TestMetricsAvailability(unittest.TestCase):
def test_is_available_returns_bool(self):
self.assertIsInstance(wl_metrics.is_available(), bool)
@_skip_no_prometheus
class TestTrackConnectionOpened(unittest.TestCase):
def test_increments_total_and_active(self):
total_before = wl_metrics.CONNECTIONS_TOTAL._value.get()
active_before = wl_metrics.CONNECTIONS_ACTIVE._value.get()
wl_metrics.track_connection_opened()
self.assertEqual(wl_metrics.CONNECTIONS_TOTAL._value.get(), total_before + 1)
self.assertEqual(wl_metrics.CONNECTIONS_ACTIVE._value.get(), active_before + 1)
@_skip_no_prometheus
class TestTrackConnectionClosed(unittest.TestCase):
def test_decrements_active(self):
wl_metrics.track_connection_opened()
active_before = wl_metrics.CONNECTIONS_ACTIVE._value.get()
wl_metrics.track_connection_closed()
self.assertEqual(wl_metrics.CONNECTIONS_ACTIVE._value.get(), active_before - 1)
@_skip_no_prometheus
class TestTrackConnectionRejected(unittest.TestCase):
def test_rejected_full(self):
before = wl_metrics.CONNECTIONS_REJECTED.labels(reason="full")._value.get()
wl_metrics.track_connection_rejected(reason="full")
self.assertEqual(wl_metrics.CONNECTIONS_REJECTED.labels(reason="full")._value.get(), before + 1)
def test_rejected_auth(self):
before = wl_metrics.CONNECTIONS_REJECTED.labels(reason="auth")._value.get()
wl_metrics.track_connection_rejected(reason="auth")
self.assertEqual(wl_metrics.CONNECTIONS_REJECTED.labels(reason="auth")._value.get(), before + 1)
@_skip_no_prometheus
class TestTrackTranscriptionLatency(unittest.TestCase):
def test_observe_records_value(self):
count_before = wl_metrics.TRANSCRIPTION_LATENCY._sum.get()
wl_metrics.track_transcription_latency(0.5)
self.assertAlmostEqual(wl_metrics.TRANSCRIPTION_LATENCY._sum.get(), count_before + 0.5, places=3)
@_skip_no_prometheus
class TestTrackAudioProcessed(unittest.TestCase):
def test_increments_by_duration(self):
before = wl_metrics.AUDIO_PROCESSED._value.get()
wl_metrics.track_audio_processed(3.5)
self.assertAlmostEqual(wl_metrics.AUDIO_PROCESSED._value.get(), before + 3.5, places=3)
@_skip_no_prometheus
class TestTrackSegmentEmitted(unittest.TestCase):
def test_completed_true(self):
before = wl_metrics.SEGMENTS_EMITTED.labels(completed="true")._value.get()
wl_metrics.track_segment_emitted(completed=True)
self.assertEqual(wl_metrics.SEGMENTS_EMITTED.labels(completed="true")._value.get(), before + 1)
def test_completed_false(self):
before = wl_metrics.SEGMENTS_EMITTED.labels(completed="false")._value.get()
wl_metrics.track_segment_emitted(completed=False)
self.assertEqual(wl_metrics.SEGMENTS_EMITTED.labels(completed="false")._value.get(), before + 1)
@_skip_no_prometheus
class TestTrackRestRequest(unittest.TestCase):
def test_tracks_200(self):
before = wl_metrics.REST_REQUESTS.labels(endpoint="transcriptions", status="200")._value.get()
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
self.assertEqual(wl_metrics.REST_REQUESTS.labels(endpoint="transcriptions", status="200")._value.get(), before + 1)
def test_tracks_500(self):
before = wl_metrics.REST_REQUESTS.labels(endpoint="transcriptions", status="500")._value.get()
wl_metrics.track_rest_request(endpoint="transcriptions", status=500)
self.assertEqual(wl_metrics.REST_REQUESTS.labels(endpoint="transcriptions", status="500")._value.get(), before + 1)
@_skip_no_prometheus
class TestTrackError(unittest.TestCase):
def test_tracks_transcription_error(self):
before = wl_metrics.ERRORS.labels(type="transcription")._value.get()
wl_metrics.track_error("transcription")
self.assertEqual(wl_metrics.ERRORS.labels(type="transcription")._value.get(), before + 1)
def test_tracks_rest_error(self):
before = wl_metrics.ERRORS.labels(type="rest_transcription")._value.get()
wl_metrics.track_error("rest_transcription")
self.assertEqual(wl_metrics.ERRORS.labels(type="rest_transcription")._value.get(), before + 1)
@_skip_no_prometheus
class TestStartMetricsServer(unittest.TestCase):
@patch("whisper_live.metrics.start_http_server")
def test_starts_on_given_port(self, mock_start):
wl_metrics.start_metrics_server(9999)
mock_start.assert_called_once_with(9999)
@patch("whisper_live.metrics.start_http_server", side_effect=OSError("port in use"))
def test_logs_error_on_failure(self, mock_start):
with self.assertLogs(level="ERROR") as cm:
wl_metrics.start_metrics_server(9999)
self.assertTrue(any("Failed to start" in msg for msg in cm.output))
class TestNoOpWhenUnavailable(unittest.TestCase):
"""Verify helper functions are no-ops when _AVAILABLE is False."""
def test_all_helpers_are_noop(self):
original = wl_metrics._AVAILABLE
try:
wl_metrics._AVAILABLE = False
# None of these should raise
wl_metrics.track_connection_opened()
wl_metrics.track_connection_closed()
wl_metrics.track_connection_rejected("full")
wl_metrics.track_transcription_latency(1.0)
wl_metrics.track_audio_processed(1.0)
wl_metrics.track_segment_emitted()
wl_metrics.track_rest_request()
wl_metrics.track_error()
finally:
wl_metrics._AVAILABLE = original
if __name__ == "__main__":
unittest.main()
+700
View File
@@ -0,0 +1,700 @@
import json
import time
import threading
import collections
import unittest
from unittest import mock
from unittest.mock import MagicMock, patch
from whisper_live.server import TranscriptionServer, BackendType, ClientManager
class TestClientManagerAddRemove(unittest.TestCase):
def setUp(self):
self.cm = ClientManager(max_clients=2, max_connection_time=60)
def test_add_and_get_client(self):
ws = MagicMock()
client = MagicMock()
self.cm.add_client(ws, client)
self.assertIs(self.cm.get_client(ws), client)
def test_get_nonexistent_client(self):
ws = MagicMock()
self.assertFalse(self.cm.get_client(ws))
def test_remove_client_calls_cleanup(self):
ws = MagicMock()
client = MagicMock()
self.cm.add_client(ws, client)
self.cm.remove_client(ws)
client.cleanup.assert_called_once()
self.assertNotIn(ws, self.cm.clients)
self.assertNotIn(ws, self.cm.start_times)
def test_remove_nonexistent_client_no_error(self):
ws = MagicMock()
self.cm.remove_client(ws) # should not raise
class TestClientManagerThreadSafety(unittest.TestCase):
def test_concurrent_add_remove(self):
cm = ClientManager(max_clients=100, max_connection_time=600)
errors = []
def add_clients(start_idx):
try:
for i in range(50):
ws = MagicMock(name=f"ws-{start_idx}-{i}")
client = MagicMock(name=f"client-{start_idx}-{i}")
cm.add_client(ws, client)
except Exception as e:
errors.append(e)
def remove_clients():
try:
for _ in range(25):
with cm.lock:
if cm.clients:
ws = next(iter(cm.clients))
else:
continue
cm.remove_client(ws)
except Exception as e:
errors.append(e)
threads = [
threading.Thread(target=add_clients, args=(0,)),
threading.Thread(target=add_clients, args=(1,)),
threading.Thread(target=remove_clients),
threading.Thread(target=remove_clients),
]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(errors, [])
def test_concurrent_get_client(self):
cm = ClientManager(max_clients=100, max_connection_time=600)
ws = MagicMock()
client = MagicMock()
cm.add_client(ws, client)
errors = []
results = []
def get_many():
try:
for _ in range(100):
results.append(cm.get_client(ws))
except Exception as e:
errors.append(e)
threads = [threading.Thread(target=get_many) for _ in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(errors, [])
self.assertTrue(all(r is client for r in results))
class TestClientManagerServerFull(unittest.TestCase):
def setUp(self):
self.cm = ClientManager(max_clients=1, max_connection_time=60)
def test_not_full_returns_false(self):
ws = MagicMock()
options = {"uid": "test"}
self.assertFalse(self.cm.is_server_full(ws, options))
def test_full_sends_wait_and_returns_true(self):
ws1 = MagicMock()
self.cm.add_client(ws1, MagicMock())
ws2 = MagicMock()
options = {"uid": "new-client"}
self.assertTrue(self.cm.is_server_full(ws2, options))
ws2.send.assert_called_once()
sent = json.loads(ws2.send.call_args[0][0])
self.assertEqual(sent["status"], "WAIT")
self.assertEqual(sent["uid"], "new-client")
class TestClientManagerTimeout(unittest.TestCase):
def setUp(self):
self.cm = ClientManager(max_clients=4, max_connection_time=10)
def test_not_timed_out(self):
ws = MagicMock()
client = MagicMock()
self.cm.add_client(ws, client)
self.assertFalse(self.cm.is_client_timeout(ws))
def test_timed_out(self):
ws = MagicMock()
client = MagicMock()
self.cm.add_client(ws, client)
self.cm.start_times[ws] = time.time() - 20
self.assertTrue(self.cm.is_client_timeout(ws))
client.disconnect.assert_called_once()
class TestClientManagerGetWaitTime(unittest.TestCase):
def test_no_clients_returns_zero(self):
cm = ClientManager(max_clients=4, max_connection_time=600)
self.assertEqual(cm.get_wait_time(), 0)
def test_single_client_wait_time(self):
cm = ClientManager(max_clients=4, max_connection_time=600)
ws = MagicMock()
cm.add_client(ws, MagicMock())
cm.start_times[ws] = time.time() - 300
wait = cm.get_wait_time()
self.assertAlmostEqual(wait, 5.0, places=0)
def test_multiple_clients_returns_minimum(self):
cm = ClientManager(max_clients=4, max_connection_time=600)
ws1, ws2 = MagicMock(), MagicMock()
cm.add_client(ws1, MagicMock())
cm.add_client(ws2, MagicMock())
cm.start_times[ws1] = time.time() - 100
cm.start_times[ws2] = time.time() - 500
wait = cm.get_wait_time()
# ws2 has 100s remaining = ~1.67 minutes
self.assertAlmostEqual(wait, 100 / 60, places=0)
class TestBackendType(unittest.TestCase):
def test_valid_types(self):
valid = BackendType.valid_types()
self.assertIn("faster_whisper", valid)
self.assertIn("tensorrt", valid)
self.assertIn("openvino", valid)
def test_is_valid(self):
self.assertTrue(BackendType.is_valid("faster_whisper"))
self.assertFalse(BackendType.is_valid("nonexistent"))
def test_type_checks(self):
self.assertTrue(BackendType.FASTER_WHISPER.is_faster_whisper())
self.assertFalse(BackendType.FASTER_WHISPER.is_tensorrt())
self.assertTrue(BackendType.TENSORRT.is_tensorrt())
self.assertTrue(BackendType.OPENVINO.is_openvino())
def test_enum_from_string(self):
bt = BackendType("faster_whisper")
self.assertEqual(bt, BackendType.FASTER_WHISPER)
def test_invalid_enum_raises(self):
with self.assertRaises(ValueError):
BackendType("invalid_backend")
class TestTranscriptionServerInit(unittest.TestCase):
def test_defaults(self):
server = TranscriptionServer()
self.assertIsNone(server.client_manager)
self.assertTrue(server.use_vad)
self.assertFalse(server.single_model)
self.assertIsNone(server.batch_config)
def test_run_invalid_backend_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, backend="nonexistent")
def test_run_invalid_trt_path_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(
host="localhost",
port=9090,
backend="tensorrt",
whisper_tensorrt_path="/nonexistent/path",
)
def test_run_max_clients_zero_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, max_clients=0)
def test_run_max_clients_negative_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, max_clients=-1)
def test_run_max_connection_time_zero_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, max_connection_time=0)
def test_run_batch_max_size_zero_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, batch_enabled=True, batch_max_size=0)
def test_run_batch_window_ms_negative_raises(self):
server = TranscriptionServer()
with self.assertRaises(ValueError):
server.run(host="localhost", port=9090, batch_enabled=True, batch_window_ms=-1)
class TestTranscriptionServerGetAudio(unittest.TestCase):
def setUp(self):
self.server = TranscriptionServer()
def test_end_of_audio_returns_false(self):
ws = MagicMock()
ws.recv.return_value = b"END_OF_AUDIO"
result = self.server.get_audio_from_websocket(ws)
self.assertFalse(result)
def test_valid_audio_returns_numpy(self):
import numpy as np
ws = MagicMock()
audio = np.array([0.1, 0.2, 0.3], dtype=np.float32)
ws.recv.return_value = audio.tobytes()
result = self.server.get_audio_from_websocket(ws)
np.testing.assert_array_almost_equal(result, audio)
def test_raw_pcm_input_normalizes_int16(self):
import numpy as np
self.server.raw_pcm_input = True
ws = MagicMock()
pcm = np.array([0, 16384, -16384, 32767], dtype=np.int16)
ws.recv.return_value = pcm.tobytes()
result = self.server.get_audio_from_websocket(ws)
expected = pcm.astype(np.float32) / 32768.0
np.testing.assert_array_almost_equal(result, expected)
self.assertTrue(result.dtype == np.float32)
self.assertTrue(np.all(result >= -1.0))
self.assertTrue(np.all(result <= 1.0))
def test_raw_pcm_input_off_reads_float32(self):
import numpy as np
self.server.raw_pcm_input = False
ws = MagicMock()
audio = np.array([0.5, -0.5], dtype=np.float32)
ws.recv.return_value = audio.tobytes()
result = self.server.get_audio_from_websocket(ws)
np.testing.assert_array_almost_equal(result, audio)
class TestTranscriptionServerHandleNewConnection(unittest.TestCase):
def setUp(self):
self.server = TranscriptionServer()
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
self.server.cache_path = "~/.cache/whisper-live/"
self.server.backend = BackendType.FASTER_WHISPER
@mock.patch("websockets.WebSocketCommonProtocol")
def test_invalid_json_returns_false(self, mock_ws):
mock_ws.recv.return_value = "not valid json {{"
result = self.server.handle_new_connection(mock_ws, None, None, False)
self.assertFalse(result)
@mock.patch("websockets.WebSocketCommonProtocol")
def test_server_full_returns_false(self, mock_ws):
# Fill server
for i in range(4):
self.server.client_manager.add_client(MagicMock(), MagicMock())
mock_ws.recv.return_value = json.dumps({
"uid": "test",
"language": "en",
"task": "transcribe",
"model": "tiny.en",
})
result = self.server.handle_new_connection(mock_ws, None, None, False)
self.assertFalse(result)
class TestTranscriptionServerCleanup(unittest.TestCase):
def setUp(self):
self.server = TranscriptionServer()
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
def test_cleanup_removes_client(self):
ws = MagicMock()
client = MagicMock()
self.server.client_manager.add_client(ws, client)
self.cleanup_server = self.server
self.server.cleanup(ws)
self.assertNotIn(ws, self.server.client_manager.clients)
client.cleanup.assert_called_once()
class TestStreamTranscription(unittest.TestCase):
"""Tests for the SSE streaming endpoint (stream=true)."""
def _make_app(self):
"""Create a FastAPI app with the transcribe endpoint that has streaming support."""
from fastapi import FastAPI, UploadFile, Form
from fastapi.testclient import TestClient
from starlette.responses import StreamingResponse
import os
import tempfile
import shutil
app = FastAPI()
server = TranscriptionServer()
@app.post("/v1/audio/transcriptions")
async def transcribe(
file: UploadFile,
stream: bool = Form(default=False),
language: str = Form(default=None),
response_format: str = Form(default="json"),
):
if stream:
return server._stream_transcription(
file, language, None, 0.0, None, None
)
return {"text": "non-streamed"}
return app
@patch("whisper_live.server.WhisperModel")
def test_stream_returns_sse_content_type(self, mock_model_cls):
mock_seg = MagicMock()
mock_seg.id = 0
mock_seg.start = 0.0
mock_seg.end = 1.0
mock_seg.text = " hello "
mock_seg.words = []
mock_info = MagicMock()
mock_info.language = "en"
mock_info.language_probability = 0.98
mock_info.duration = 1.0
mock_model = MagicMock()
mock_model.transcribe.return_value = (iter([mock_seg]), mock_info)
mock_model_cls.return_value = mock_model
import io
from fastapi.testclient import TestClient
app = self._make_app()
client = TestClient(app)
resp = client.post(
"/v1/audio/transcriptions",
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
data={"stream": "true"},
)
self.assertEqual(resp.status_code, 200)
self.assertIn("text/event-stream", resp.headers.get("content-type", ""))
@patch("whisper_live.server.WhisperModel")
def test_stream_yields_segment_and_done(self, mock_model_cls):
mock_seg = MagicMock()
mock_seg.id = 0
mock_seg.start = 0.0
mock_seg.end = 1.5
mock_seg.text = " hello world "
mock_seg.words = []
mock_info = MagicMock()
mock_info.language = "en"
mock_info.language_probability = 0.95
mock_info.duration = 1.5
mock_model = MagicMock()
mock_model.transcribe.return_value = (iter([mock_seg]), mock_info)
mock_model_cls.return_value = mock_model
import io
from fastapi.testclient import TestClient
app = self._make_app()
client = TestClient(app)
resp = client.post(
"/v1/audio/transcriptions",
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
data={"stream": "true"},
)
body = resp.text
self.assertIn('"text": "hello world"', body)
self.assertIn("[DONE]", body)
@patch("whisper_live.server.WhisperModel")
def test_stream_multiple_segments(self, mock_model_cls):
segs = []
for i in range(3):
s = MagicMock()
s.id = i
s.start = float(i)
s.end = float(i + 1)
s.text = f" segment {i} "
s.words = []
segs.append(s)
mock_info = MagicMock()
mock_info.language = "en"
mock_info.language_probability = 0.99
mock_info.duration = 3.0
mock_model = MagicMock()
mock_model.transcribe.return_value = (iter(segs), mock_info)
mock_model_cls.return_value = mock_model
import io
from fastapi.testclient import TestClient
app = self._make_app()
client = TestClient(app)
resp = client.post(
"/v1/audio/transcriptions",
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
data={"stream": "true"},
)
body = resp.text
events = [line for line in body.split("\n") if line.startswith("data: ") and "[DONE]" not in line and '"type": "metadata"' not in line]
self.assertEqual(len(events), 3)
for i, event in enumerate(events):
data = json.loads(event.removeprefix("data: "))
self.assertEqual(data["text"], f"segment {i}")
@patch("whisper_live.server.WhisperModel", side_effect=RuntimeError("model error"))
def test_stream_error_yields_error_event(self, mock_model_cls):
import io
from fastapi.testclient import TestClient
app = self._make_app()
client = TestClient(app)
resp = client.post(
"/v1/audio/transcriptions",
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
data={"stream": "true"},
)
body = resp.text
self.assertIn('"error"', body)
self.assertIn("model error", body)
def test_non_stream_still_works(self):
import io
from fastapi.testclient import TestClient
app = self._make_app()
client = TestClient(app)
resp = client.post(
"/v1/audio/transcriptions",
files={"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")},
data={"stream": "false"},
)
self.assertEqual(resp.status_code, 200)
self.assertEqual(resp.json()["text"], "non-streamed")
class TestRESTAPIParamWarnings(unittest.TestCase):
"""Test that unsupported OpenAI-compatible REST params produce warnings."""
@classmethod
def setUpClass(cls):
"""Build a FastAPI test app by extracting the endpoint definition."""
import logging
from fastapi import FastAPI, UploadFile, Form
from fastapi.testclient import TestClient
from typing import Optional, List
from starlette.responses import PlainTextResponse, JSONResponse
app = FastAPI()
@app.post("/v1/audio/transcriptions")
async def transcribe(
file: UploadFile,
model: str = Form(default="whisper-1"),
language: Optional[str] = Form(default=None),
prompt: Optional[str] = Form(default=None),
response_format: str = Form(default="json"),
temperature: float = Form(default=0.0),
timestamp_granularities: Optional[List[str]] = Form(default=None),
chunking_strategy: Optional[str] = Form(default=None),
include: Optional[List[str]] = Form(default=None),
known_speaker_names: Optional[List[str]] = Form(default=None),
known_speaker_references: Optional[List[str]] = Form(default=None),
stream: bool = Form(default=False),
):
ignored_params = []
if chunking_strategy:
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
if known_speaker_names:
ignored_params.append("known_speaker_names")
if known_speaker_references:
ignored_params.append("known_speaker_references")
if include:
ignored_params.append(f"include={include}")
if ignored_params:
logging.warning(f"Unsupported OpenAI params ignored: {', '.join(ignored_params)}")
# Return a JSON response with the ignored list for testing
return {"text": "test", "ignored": ignored_params}
cls.test_client = TestClient(app)
def _post(self, **extra_fields):
import io
data = {**extra_fields}
files = {"file": ("test.wav", io.BytesIO(b"\x00" * 100), "audio/wav")}
return self.test_client.post("/v1/audio/transcriptions", data=data, files=files)
def test_no_warnings_when_no_extra_params(self):
resp = self._post()
self.assertEqual(resp.status_code, 200)
self.assertEqual(resp.json()["ignored"], [])
def test_chunking_strategy_warning(self):
resp = self._post(chunking_strategy="auto")
self.assertEqual(resp.status_code, 200)
ignored = resp.json()["ignored"]
self.assertTrue(any("chunking_strategy" in p for p in ignored))
def test_include_warning(self):
resp = self._post(include="logprobs")
self.assertEqual(resp.status_code, 200)
ignored = resp.json()["ignored"]
self.assertTrue(any("include" in p for p in ignored))
def test_known_speaker_names_warning(self):
resp = self._post(known_speaker_names="alice")
self.assertEqual(resp.status_code, 200)
ignored = resp.json()["ignored"]
self.assertTrue(any("known_speaker_names" in p for p in ignored))
def test_multiple_ignored_params(self):
resp = self._post(chunking_strategy="auto", known_speaker_names="bob")
self.assertEqual(resp.status_code, 200)
ignored = resp.json()["ignored"]
self.assertGreaterEqual(len(ignored), 2)
class TestAPIKeyAuth(unittest.TestCase):
"""Test optional API key authentication middleware."""
@classmethod
def setUpClass(cls):
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from fastapi.responses import JSONResponse as JSONR
app = FastAPI()
@app.middleware("http")
async def _check_api_key(request: Request, call_next):
auth = request.headers.get("Authorization", "")
if auth != "Bearer test-secret":
return JSONR({"error": "Invalid or missing API key"}, status_code=401)
return await call_next(request)
@app.get("/ping")
async def ping():
return {"status": "ok"}
cls.test_client = TestClient(app)
def test_missing_key_returns_401(self):
resp = self.test_client.get("/ping")
self.assertEqual(resp.status_code, 401)
def test_wrong_key_returns_401(self):
resp = self.test_client.get("/ping", headers={"Authorization": "Bearer wrong"})
self.assertEqual(resp.status_code, 401)
def test_correct_key_returns_200(self):
resp = self.test_client.get("/ping", headers={"Authorization": "Bearer test-secret"})
self.assertEqual(resp.status_code, 200)
self.assertEqual(resp.json()["status"], "ok")
class TestRateLimiting(unittest.TestCase):
"""Test per-IP rate limiting middleware."""
def _make_app(self, rpm_limit=3):
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from fastapi.responses import JSONResponse as JSONR
_rate_lock = threading.Lock()
_rate_buckets: dict = {}
app = FastAPI()
@app.middleware("http")
async def _rate_limit(request: Request, call_next):
client_ip = request.client.host if request.client else "unknown"
now = time.time()
with _rate_lock:
bucket = _rate_buckets.setdefault(client_ip, collections.deque())
while bucket and bucket[0] < now - 60:
bucket.popleft()
if len(bucket) >= rpm_limit:
return JSONR({"error": "Rate limit exceeded"}, status_code=429)
bucket.append(now)
return await call_next(request)
@app.get("/ping")
async def ping():
return {"status": "ok"}
return TestClient(app)
def test_within_limit_succeeds(self):
client = self._make_app(rpm_limit=3)
for _ in range(3):
resp = client.get("/ping")
self.assertEqual(resp.status_code, 200)
def test_exceeding_limit_returns_429(self):
client = self._make_app(rpm_limit=3)
for _ in range(3):
client.get("/ping")
resp = client.get("/ping")
self.assertEqual(resp.status_code, 429)
self.assertIn("Rate limit", resp.json()["error"])
class TestWebSocketAuth(unittest.TestCase):
"""Tests for the WebSocket process_request auth callback."""
def _make_auth_handler(self, api_key):
"""Build the same auth function the server creates."""
def _ws_auth(path, request_headers):
auth = request_headers.get("Authorization", "")
token_param = None
if "?" in path:
from urllib.parse import urlparse, parse_qs
parsed = urlparse(path)
token_param = parse_qs(parsed.query).get("token", [None])[0]
if auth == f"Bearer {api_key}" or token_param == api_key:
return None
return (401, [("Content-Type", "text/plain")], b"Unauthorized\n")
return _ws_auth
def test_valid_bearer_token(self):
handler = self._make_auth_handler("my-secret")
result = handler("/", {"Authorization": "Bearer my-secret"})
self.assertIsNone(result)
def test_invalid_bearer_token(self):
handler = self._make_auth_handler("my-secret")
result = handler("/", {"Authorization": "Bearer wrong"})
self.assertEqual(result[0], 401)
def test_missing_auth_header(self):
handler = self._make_auth_handler("my-secret")
result = handler("/", {})
self.assertEqual(result[0], 401)
def test_valid_query_token(self):
handler = self._make_auth_handler("my-secret")
result = handler("/?token=my-secret", {})
self.assertIsNone(result)
def test_invalid_query_token(self):
handler = self._make_auth_handler("my-secret")
result = handler("/?token=wrong", {})
self.assertEqual(result[0], 401)
if __name__ == "__main__":
unittest.main()
+140
View File
@@ -0,0 +1,140 @@
import os
import tempfile
import unittest
from io import StringIO
from unittest.mock import patch
from whisper_live.utils import format_time, create_srt_file, print_transcript, clear_screen
class TestFormatTime(unittest.TestCase):
def test_zero(self):
self.assertEqual(format_time(0), "00:00:00,000")
def test_seconds_only(self):
self.assertEqual(format_time(5.0), "00:00:05,000")
def test_fractional_seconds(self):
self.assertEqual(format_time(1.5), "00:00:01,500")
def test_minutes(self):
self.assertEqual(format_time(65.0), "00:01:05,000")
def test_hours(self):
self.assertEqual(format_time(3661.123), "01:01:01,123")
def test_millisecond_precision(self):
self.assertEqual(format_time(0.001), "00:00:00,001")
def test_large_value(self):
# float precision: int((86399.999 - 86399) * 1000) may be 998 or 999
result = format_time(86399.999)
self.assertIn(result, ("23:59:59,998", "23:59:59,999"))
def test_rounding_edge(self):
result = format_time(0.9999)
# 0.9999 -> int(s%60)=0, milliseconds=int(0.9999*1000)=999
self.assertEqual(result, "00:00:00,999")
class TestCreateSrtFile(unittest.TestCase):
def test_single_segment(self):
segments = [{"start": "0.000", "end": "1.500", "text": "Hello world"}]
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
path = f.name
try:
create_srt_file(segments, path)
with open(path, "r", encoding="utf-8") as f:
content = f.read()
self.assertIn("1\n", content)
self.assertIn("00:00:00,000 --> 00:00:01,500", content)
self.assertIn("Hello world", content)
finally:
os.remove(path)
def test_multiple_segments(self):
segments = [
{"start": "0.000", "end": "1.000", "text": "First"},
{"start": "1.000", "end": "2.500", "text": "Second"},
{"start": "2.500", "end": "4.000", "text": "Third"},
]
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
path = f.name
try:
create_srt_file(segments, path)
with open(path, "r", encoding="utf-8") as f:
content = f.read()
self.assertIn("1\n", content)
self.assertIn("2\n", content)
self.assertIn("3\n", content)
self.assertIn("First", content)
self.assertIn("Third", content)
finally:
os.remove(path)
def test_empty_segments(self):
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
path = f.name
try:
create_srt_file([], path)
with open(path, "r", encoding="utf-8") as f:
content = f.read()
self.assertEqual(content, "")
finally:
os.remove(path)
def test_unicode_text(self):
segments = [{"start": "0.000", "end": "1.000", "text": "日本語テスト"}]
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
path = f.name
try:
create_srt_file(segments, path)
with open(path, "r", encoding="utf-8") as f:
content = f.read()
self.assertIn("日本語テスト", content)
finally:
os.remove(path)
class TestPrintTranscript(unittest.TestCase):
@patch("sys.stdout", new_callable=StringIO)
def test_clear_screen_uses_ansi(self, mock_stdout):
clear_screen()
output = mock_stdout.getvalue()
self.assertIn("\033[H\033[2J", output)
@patch("sys.stdout", new_callable=StringIO)
def test_print_plain_text(self, mock_stdout):
text = ["Hello", " world"]
print_transcript(text)
output = mock_stdout.getvalue()
self.assertIn("Hello world", output)
@patch("sys.stdout", new_callable=StringIO)
def test_print_with_timestamps(self, mock_stdout):
text = [
{"start": 0.0, "end": 1.0, "text": "Hello"},
{"start": 1.0, "end": 2.0, "text": "world"},
]
print_transcript(text, timestamps=True)
output = mock_stdout.getvalue()
self.assertIn("[0.0 -> 1.0]", output)
self.assertIn("Hello", output)
@patch("sys.stdout", new_callable=StringIO)
def test_print_translated(self, mock_stdout):
text = ["Bonjour", "le monde"]
print_transcript(text, translated=True)
output = mock_stdout.getvalue()
self.assertIn("Bonjour le monde", output)
@patch("sys.stdout", new_callable=StringIO)
def test_print_empty(self, mock_stdout):
print_transcript([])
output = mock_stdout.getvalue()
# empty text joined is empty string, should not crash
self.assertEqual(output.strip(), "")
if __name__ == "__main__":
unittest.main()
+131
View File
@@ -0,0 +1,131 @@
import unittest
from unittest.mock import patch, MagicMock
import numpy as np
import torch
from whisper_live.vad import VoiceActivityDetection, VoiceActivityDetector
class TestVoiceActivityDetectionValidation(unittest.TestCase):
"""Tests for VoiceActivityDetection input validation without requiring the ONNX model."""
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
def setUp(self):
self.vad = VoiceActivityDetection()
self.vad.sample_rates = [8000, 16000]
def test_1d_input_unsqueezed(self):
x = torch.randn(512)
x_out, sr_out = self.vad._validate_input(x, 16000)
self.assertEqual(x_out.dim(), 2)
self.assertEqual(sr_out, 16000)
def test_3d_input_raises(self):
x = torch.randn(1, 1, 512)
with self.assertRaises(ValueError):
self.vad._validate_input(x, 16000)
def test_unsupported_sample_rate_raises(self):
x = torch.randn(1, 512)
with self.assertRaises(ValueError):
self.vad._validate_input(x, 44100)
def test_too_short_audio_raises(self):
x = torch.randn(1, 1)
with self.assertRaises(ValueError):
self.vad._validate_input(x, 16000)
def test_downsample_multiple_of_16k(self):
x = torch.randn(1, 512 * 3)
x_out, sr_out = self.vad._validate_input(x, 48000)
self.assertEqual(sr_out, 16000)
self.assertEqual(x_out.shape[1], 512)
class TestVoiceActivityDetectionStateReset(unittest.TestCase):
"""Tests for VoiceActivityDetection.reset_states()."""
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
def setUp(self):
self.vad = VoiceActivityDetection()
def test_reset_creates_correct_shapes(self):
self.vad.reset_states(batch_size=4)
self.assertEqual(self.vad._state.shape, (2, 4, 128))
self.assertEqual(self.vad._context.shape[0], 0)
self.assertEqual(self.vad._last_sr, 0)
self.assertEqual(self.vad._last_batch_size, 0)
def test_reset_default_batch_size(self):
self.vad.reset_states()
self.assertEqual(self.vad._state.shape, (2, 1, 128))
class TestVoiceActivityDetectionDownload(unittest.TestCase):
"""Tests for the model download function."""
@patch("os.path.exists", return_value=True)
def test_skips_download_if_exists(self, mock_exists):
path = VoiceActivityDetection.download()
self.assertTrue(path.endswith("silero_vad.onnx"))
@patch("os.path.exists", return_value=False)
@patch("subprocess.run")
@patch("os.makedirs")
def test_downloads_if_missing(self, mock_makedirs, mock_run, mock_exists):
path = VoiceActivityDetection.download()
mock_run.assert_called_once()
self.assertIn("silero_vad.onnx", path)
@patch("os.path.exists", return_value=False)
@patch("subprocess.run", side_effect=Exception("wget not found"))
@patch("os.makedirs")
def test_handles_download_failure(self, mock_makedirs, mock_run, mock_exists):
# should not raise, just prints an error
with self.assertRaises(Exception):
VoiceActivityDetection.download()
class TestVoiceActivityDetectorThreshold(unittest.TestCase):
"""Tests for VoiceActivityDetector threshold behavior."""
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
def test_above_threshold_returns_true(self):
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
detector.model = VoiceActivityDetection()
detector.threshold = 0.5
detector.frame_rate = 16000
mock_probs = torch.tensor([[0.9, 0.8, 0.7]])
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
result = detector(np.random.randn(16000).astype(np.float32))
self.assertTrue(result)
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
def test_below_threshold_returns_false(self):
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
detector.model = VoiceActivityDetection()
detector.threshold = 0.5
detector.frame_rate = 16000
mock_probs = torch.tensor([[0.1, 0.2, 0.3]])
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
result = detector(np.random.randn(16000).astype(np.float32))
self.assertFalse(result)
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
def test_custom_threshold(self):
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
detector.model = VoiceActivityDetection()
detector.threshold = 0.95
detector.frame_rate = 16000
mock_probs = torch.tensor([[0.9]])
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
result = detector(np.random.randn(16000).astype(np.float32))
self.assertFalse(result)
if __name__ == "__main__":
unittest.main()
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.8.0" __version__ = "0.9.0"
+113 -10
View File
@@ -5,12 +5,23 @@ import time
import queue import queue
import numpy as np import numpy as np
from whisper_live import metrics as wl_metrics
class ServeClientBase(object): class ServeClientBase(object):
RATE = 16000 RATE = 16000
SERVER_READY = "SERVER_READY" SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT" DISCONNECT = "DISCONNECT"
MAX_BUFFER_DURATION_S = 45
"""Maximum audio buffer duration in seconds before trimming."""
BUFFER_TRIM_DURATION_S = 30
"""Duration in seconds to trim from the buffer when it exceeds MAX_BUFFER_DURATION_S."""
CLIP_THRESHOLD_DURATION_S = 25
"""Duration threshold in seconds for clipping audio with no valid segments."""
CLIP_TAIL_DURATION_S = 5
"""Duration in seconds of audio to keep after clipping."""
client_uid: str client_uid: str
"""A unique identifier for the client.""" """A unique identifier for the client."""
websocket: object websocket: object
@@ -24,6 +35,9 @@ class ServeClientBase(object):
same_output_threshold: int same_output_threshold: int
"""Number of repeated outputs before considering it as a valid segment.""" """Number of repeated outputs before considering it as a valid segment."""
MAX_TRANSCRIPT_LENGTH = 500
MAX_TRANSLATION_QUEUE_SIZE = 100
def __init__( def __init__(
self, self,
client_uid, client_uid,
@@ -33,6 +47,8 @@ class ServeClientBase(object):
clip_audio=False, clip_audio=False,
same_output_threshold=10, same_output_threshold=10,
translation_queue=None, translation_queue=None,
diarization=None,
word_timestamps=False,
): ):
self.client_uid = client_uid self.client_uid = client_uid
self.websocket = websocket self.websocket = websocket
@@ -40,6 +56,8 @@ class ServeClientBase(object):
self.no_speech_thresh = no_speech_thresh self.no_speech_thresh = no_speech_thresh
self.clip_audio = clip_audio self.clip_audio = clip_audio
self.same_output_threshold = same_output_threshold self.same_output_threshold = same_output_threshold
self.diarization = diarization
self.word_timestamps = word_timestamps
self.frames = b"" self.frames = b""
self.timestamp_offset = 0.0 self.timestamp_offset = 0.0
@@ -54,6 +72,13 @@ class ServeClientBase(object):
self.end_time_for_same_output = None self.end_time_for_same_output = None
self.translation_queue = translation_queue self.translation_queue = translation_queue
# Optional post-processing callable for segments.
# If set, called with a segment dict and must return a segment dict.
# Allows external projects to plug in custom post-processing
# (e.g. PII redaction, formatting, diarization) without modifying
# WhisperLive's core code.
self.segment_post_processor = None
# threading # threading
self.lock = threading.Lock() self.lock = threading.Lock()
@@ -89,16 +114,20 @@ class ServeClientBase(object):
continue continue
try: try:
input_sample = input_bytes.copy() input_sample = input_bytes.copy()
t0 = time.time()
result = self.transcribe_audio(input_sample) result = self.transcribe_audio(input_sample)
if result is None or self.language is None: if result is None or self.language is None:
self.timestamp_offset += duration self.timestamp_offset += duration
time.sleep(0.25) # wait for voice activity, result is None when no voice activity time.sleep(0.25) # wait for voice activity, result is None when no voice activity
continue continue
wl_metrics.track_transcription_latency(time.time() - t0)
wl_metrics.track_audio_processed(duration)
self.handle_transcription_output(result, duration) self.handle_transcription_output(result, duration)
except Exception as e: except Exception as e:
logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}") logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
wl_metrics.track_error("transcription")
time.sleep(0.01) time.sleep(0.01)
def transcribe_audio(self): def transcribe_audio(self):
@@ -107,7 +136,7 @@ class ServeClientBase(object):
def handle_transcription_output(self, result, duration): def handle_transcription_output(self, result, duration):
raise NotImplementedError raise NotImplementedError
def format_segment(self, start, end, text, completed=False): def format_segment(self, start, end, text, completed=False, speaker=None, words=None):
""" """
Formats a transcription segment with precise start and end times alongside the transcribed text. Formats a transcription segment with precise start and end times alongside the transcribed text.
@@ -115,18 +144,25 @@ class ServeClientBase(object):
start (float): The start time of the transcription segment in seconds. start (float): The start time of the transcription segment in seconds.
end (float): The end time of the transcription segment in seconds. end (float): The end time of the transcription segment in seconds.
text (str): The transcribed text corresponding to the segment. text (str): The transcribed text corresponding to the segment.
speaker (str, optional): Speaker label from diarization.
words (list, optional): Word-level timestamps and probabilities.
Returns: Returns:
dict: A dictionary representing the formatted transcription segment, including dict: A dictionary representing the formatted transcription segment, including
'start' and 'end' times as strings with three decimal places and the 'text' 'start' and 'end' times as strings with three decimal places and the 'text'
of the transcription. of the transcription.
""" """
return { seg = {
'start': "{:.3f}".format(start), 'start': "{:.3f}".format(start),
'end': "{:.3f}".format(end), 'end': "{:.3f}".format(end),
'text': text, 'text': text,
'completed': completed 'completed': completed,
} }
if speaker is not None:
seg['speaker'] = speaker
if words is not None:
seg['words'] = words
return seg
def add_frames(self, frame_np): def add_frames(self, frame_np):
""" """
@@ -145,9 +181,9 @@ class ServeClientBase(object):
""" """
self.lock.acquire() self.lock.acquire()
if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE: if self.frames_np is not None and self.frames_np.shape[0] > self.MAX_BUFFER_DURATION_S*self.RATE:
self.frames_offset += 30.0 self.frames_offset += float(self.BUFFER_TRIM_DURATION_S)
self.frames_np = self.frames_np[int(30*self.RATE):] self.frames_np = self.frames_np[int(self.BUFFER_TRIM_DURATION_S*self.RATE):]
# check timestamp offset(should be >= self.frame_offset) # check timestamp offset(should be >= self.frame_offset)
# this basically means that there is no speech as timestamp offset hasnt updated # this basically means that there is no speech as timestamp offset hasnt updated
# and is less than frame_offset # and is less than frame_offset
@@ -166,9 +202,9 @@ class ServeClientBase(object):
no valid segment for the last 30 seconds from whisper no valid segment for the last 30 seconds from whisper
""" """
with self.lock: with self.lock:
if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE: if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > self.CLIP_THRESHOLD_DURATION_S * self.RATE:
duration = self.frames_np.shape[0] / self.RATE duration = self.frames_np.shape[0] / self.RATE
self.timestamp_offset = self.frames_offset + duration - 5 self.timestamp_offset = self.frames_offset + duration - self.CLIP_TAIL_DURATION_S
def get_audio_chunk_for_processing(self): def get_audio_chunk_for_processing(self):
""" """
@@ -234,9 +270,23 @@ class ServeClientBase(object):
This method formats the transcription segments into a JSON object and attempts to send This method formats the transcription segments into a JSON object and attempts to send
this object to the client. If an error occurs during the send operation, it logs the error. this object to the client. If an error occurs during the send operation, it logs the error.
If a ``segment_post_processor`` callable is set, each segment is passed through it
before sending. The callable receives a segment dict and must return a segment dict.
Returns: Returns:
segments (list): A list of transcription segments to be sent to the client. segments (list): A list of transcription segments to be sent to the client.
""" """
if self.segment_post_processor is not None:
processed = []
for seg in segments:
try:
result = self.segment_post_processor(seg)
processed.append(result if result is not None else seg)
except Exception as e:
logging.error(f"[ERROR]: segment_post_processor failed: {e}")
processed.append(seg)
segments = processed
try: try:
self.websocket.send( self.websocket.send(
json.dumps({ json.dumps({
@@ -244,6 +294,8 @@ class ServeClientBase(object):
"segments": segments, "segments": segments,
}) })
) )
for seg in segments:
wl_metrics.track_segment_emitted(completed=seg.get("completed", False))
except Exception as e: except Exception as e:
logging.error(f"[ERROR]: Sending data to client: {e}") logging.error(f"[ERROR]: Sending data to client: {e}")
@@ -281,6 +333,45 @@ class ServeClientBase(object):
def get_segment_end(self, segment): def get_segment_end(self, segment):
return getattr(segment, "end", getattr(segment, "end_ts", 0)) return getattr(segment, "end", getattr(segment, "end_ts", 0))
def _identify_speaker(self, segment):
"""Run diarization on a segment's audio slice if diarization is enabled.
Returns:
str or None: Speaker label, or None if diarization is disabled or audio unavailable.
"""
if self.diarization is None or self.frames_np is None:
return None
try:
seg_start = self.get_segment_start(segment)
seg_end = self.get_segment_end(segment)
start_sample = int(seg_start * self.RATE)
end_sample = int(seg_end * self.RATE)
samples_offset = max(0, int((self.timestamp_offset - self.frames_offset) * self.RATE))
audio_slice = self.frames_np[samples_offset + start_sample:samples_offset + end_sample]
if len(audio_slice) < self.RATE * 0.3:
return None
return self.diarization.identify_speaker(audio_slice, self.RATE)
except Exception as e:
logging.error(f"Diarization error: {e}")
return None
def _extract_words(self, segment, time_offset):
"""Extracts word-level timestamps from a segment if word_timestamps is enabled."""
if not self.word_timestamps:
return None
words = getattr(segment, "words", None)
if not words:
return None
return [
{
"word": w.word,
"start": "{:.3f}".format(time_offset + w.start),
"end": "{:.3f}".format(time_offset + w.end),
"probability": round(w.probability, 4),
}
for w in words
]
def update_segments(self, segments, duration): def update_segments(self, segments, duration):
""" """
Processes the segments from Whisper and updates the transcript. Processes the segments from Whisper and updates the transcript.
@@ -310,7 +401,9 @@ class ServeClientBase(object):
continue continue
if self.get_segment_no_speech_prob(s) > self.no_speech_thresh: if self.get_segment_no_speech_prob(s) > self.no_speech_thresh:
continue continue
completed_segment = self.format_segment(start, end, text_, completed=True) speaker = self._identify_speaker(s)
words = self._extract_words(s, self.timestamp_offset)
completed_segment = self.format_segment(start, end, text_, completed=True, speaker=speaker, words=words)
self.transcript.append(completed_segment) self.transcript.append(completed_segment)
if self.translation_queue: if self.translation_queue:
@@ -323,12 +416,14 @@ class ServeClientBase(object):
# Process the last segment if its no_speech_prob is acceptable. # Process the last segment if its no_speech_prob is acceptable.
if self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh: if self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh:
self.current_out += segments[-1].text self.current_out += segments[-1].text
words = self._extract_words(segments[-1], self.timestamp_offset)
with self.lock: with self.lock:
last_segment = self.format_segment( last_segment = self.format_segment(
self.timestamp_offset + self.get_segment_start(segments[-1]), self.timestamp_offset + self.get_segment_start(segments[-1]),
self.timestamp_offset + min(duration, self.get_segment_end(segments[-1])), self.timestamp_offset + min(duration, self.get_segment_end(segments[-1])),
self.current_out, self.current_out,
completed=False completed=False,
words=words
) )
# Handle repeated output logic. # Handle repeated output logic.
@@ -376,4 +471,12 @@ class ServeClientBase(object):
with self.lock: with self.lock:
self.timestamp_offset += offset self.timestamp_offset += offset
self._trim_transcript()
return last_segment return last_segment
def _trim_transcript(self):
"""Trims transcript and text lists to prevent unbounded memory growth."""
if len(self.transcript) > self.MAX_TRANSCRIPT_LENGTH:
self.transcript = self.transcript[-self.MAX_TRANSCRIPT_LENGTH:]
if len(self.text) > self.MAX_TRANSCRIPT_LENGTH:
self.text = self.text[-self.MAX_TRANSCRIPT_LENGTH:]
+11 -2
View File
@@ -34,6 +34,9 @@ class ServeClientFasterWhisper(ServeClientBase):
same_output_threshold=7, same_output_threshold=7,
cache_path="~/.cache/whisper-live/", cache_path="~/.cache/whisper-live/",
translation_queue=None, translation_queue=None,
hotwords=None,
diarization=None,
word_timestamps=False,
): ):
""" """
Initialize a ServeClient instance. Initialize a ServeClient instance.
@@ -63,7 +66,9 @@ class ServeClientFasterWhisper(ServeClientBase):
no_speech_thresh, no_speech_thresh,
clip_audio, clip_audio,
same_output_threshold, same_output_threshold,
translation_queue translation_queue,
diarization,
word_timestamps,
) )
self.cache_path = cache_path self.cache_path = cache_path
self.model_sizes = [ self.model_sizes = [
@@ -78,6 +83,7 @@ class ServeClientFasterWhisper(ServeClientBase):
self.task = task self.task = task
self.initial_prompt = initial_prompt self.initial_prompt = initial_prompt
self.vad_parameters = vad_parameters or {"threshold": 0.5} self.vad_parameters = vad_parameters or {"threshold": 0.5}
self.hotwords = hotwords
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda": if device == "cuda":
@@ -213,6 +219,7 @@ class ServeClientFasterWhisper(ServeClientBase):
initial_prompt=self.initial_prompt, initial_prompt=self.initial_prompt,
use_vad=self.use_vad, use_vad=self.use_vad,
vad_parameters=self.vad_parameters if self.use_vad else None, vad_parameters=self.vad_parameters if self.use_vad else None,
word_timestamps=self.word_timestamps,
) )
ServeClientFasterWhisper.BATCH_WORKER.submit(request) ServeClientFasterWhisper.BATCH_WORKER.submit(request)
request.future.wait(timeout=30) request.future.wait(timeout=30)
@@ -231,7 +238,9 @@ class ServeClientFasterWhisper(ServeClientBase):
language=self.language, language=self.language,
task=self.task, task=self.task,
vad_filter=self.use_vad, vad_filter=self.use_vad,
vad_parameters=self.vad_parameters if self.use_vad else None) vad_parameters=self.vad_parameters if self.use_vad else None,
hotwords=self.hotwords,
word_timestamps=self.word_timestamps)
if ServeClientFasterWhisper.SINGLE_MODEL: if ServeClientFasterWhisper.SINGLE_MODEL:
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release() ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
+51 -11
View File
@@ -43,6 +43,12 @@ class Client:
translation_srt_file_path="output_translated.srt", translation_srt_file_path="output_translated.srt",
enable_timestamps=False, enable_timestamps=False,
display_segments=4, display_segments=4,
hotwords=None,
enable_diarization=False,
max_speakers=10,
word_timestamps=False,
max_retries=0,
retry_delay=5,
): ):
""" """
Initializes a Client instance for audio recording and streaming to a server. Initializes a Client instance for audio recording and streaming to a server.
@@ -101,21 +107,21 @@ class Client:
self.task = "translate" self.task = "translate"
self.enable_timestamps = enable_timestamps self.enable_timestamps = enable_timestamps
self.display_segments = display_segments self.display_segments = display_segments
self.hotwords = hotwords
self.enable_diarization = enable_diarization
self.max_speakers = max_speakers
self.word_timestamps = word_timestamps
self.max_retries = max_retries
self.retry_delay = retry_delay
self._retry_count = 0
self.audio_bytes = None self.audio_bytes = None
if host is not None and port is not None: if host is not None and port is not None:
self.host = host
self.port = port
socket_protocol = 'wss' if self.use_wss else "ws" socket_protocol = 'wss' if self.use_wss else "ws"
socket_url = f"{socket_protocol}://{host}:{port}" self.socket_url = f"{socket_protocol}://{host}:{port}"
self.client_socket = websocket.WebSocketApp( self._create_websocket()
socket_url,
on_open=lambda ws: self.on_open(ws),
on_message=lambda ws, message: self.on_message(ws, message),
on_error=lambda ws, error: self.on_error(ws, error),
on_close=lambda ws, close_status_code, close_msg: self.on_close(
ws, close_status_code, close_msg
),
)
else: else:
print("[ERROR]: No host or port specified.") print("[ERROR]: No host or port specified.")
return return
@@ -131,6 +137,18 @@ class Client:
self.translated_transcript = [] self.translated_transcript = []
print("[INFO]: * recording") print("[INFO]: * recording")
def _create_websocket(self):
"""Creates a new WebSocketApp instance."""
self.client_socket = websocket.WebSocketApp(
self.socket_url,
on_open=lambda ws: self.on_open(ws),
on_message=lambda ws, message: self.on_message(ws, message),
on_error=lambda ws, error: self.on_error(ws, error),
on_close=lambda ws, close_status_code, close_msg: self.on_close(
ws, close_status_code, close_msg
),
)
def handle_status_messages(self, message_data): def handle_status_messages(self, message_data):
"""Handles server status messages.""" """Handles server status messages."""
status = message_data["status"] status = message_data["status"]
@@ -273,6 +291,15 @@ class Client:
self.recording = False self.recording = False
self.waiting = False self.waiting = False
if self.max_retries > 0 and self._retry_count < self.max_retries and not self.server_error:
self._retry_count += 1
print(f"[INFO]: Reconnecting ({self._retry_count}/{self.max_retries}) in {self.retry_delay}s...")
time.sleep(self.retry_delay)
self._create_websocket()
self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
self.ws_thread.daemon = True
self.ws_thread.start()
def on_open(self, ws): def on_open(self, ws):
""" """
Callback function called when the WebSocket connection is successfully opened. Callback function called when the WebSocket connection is successfully opened.
@@ -299,6 +326,10 @@ class Client:
"same_output_threshold": self.same_output_threshold, "same_output_threshold": self.same_output_threshold,
"enable_translation": self.enable_translation, "enable_translation": self.enable_translation,
"target_language": self.target_language, "target_language": self.target_language,
"hotwords": self.hotwords,
"enable_diarization": self.enable_diarization,
"max_speakers": self.max_speakers,
"word_timestamps": self.word_timestamps,
} }
) )
) )
@@ -820,7 +851,12 @@ class TranscriptionClient(TranscriptionTeeClient):
translation_srt_file_path="./output_translated.srt", translation_srt_file_path="./output_translated.srt",
enable_timestamps=False, enable_timestamps=False,
display_segments=4, display_segments=4,
hotwords=None,
enable_diarization=False,
max_speakers=10,
word_timestamps=False,
): ):
self.client = Client( self.client = Client(
host, host,
port, port,
@@ -842,6 +878,10 @@ class TranscriptionClient(TranscriptionTeeClient):
translation_srt_file_path=translation_srt_file_path, translation_srt_file_path=translation_srt_file_path,
enable_timestamps=enable_timestamps, enable_timestamps=enable_timestamps,
display_segments=display_segments, display_segments=display_segments,
hotwords=hotwords,
enable_diarization=enable_diarization,
max_speakers=max_speakers,
word_timestamps=word_timestamps,
) )
if save_output_recording and not output_recording_filename.endswith(".wav"): if save_output_recording and not output_recording_filename.endswith(".wav"):
+142
View File
@@ -0,0 +1,142 @@
"""
Optional speaker diarization module for WhisperLive.
Uses speaker embeddings and online clustering to assign speaker labels
to transcription segments in real-time. Requires pyannote.audio as an
optional dependency.
Install: pip install pyannote.audio
"""
import logging
import numpy as np
class SpeakerDiarizer:
"""Real-time speaker diarization using speaker embeddings and online clustering.
Each completed transcription segment's audio is passed through a speaker
embedding model. The embedding is compared against known speakers using
cosine similarity. If no match exceeds the threshold, a new speaker is
created.
Args:
similarity_threshold (float): Minimum cosine similarity to match an
existing speaker. Lower values merge speakers more aggressively.
Default 0.55.
max_speakers (int): Maximum number of distinct speakers to track.
Once reached, new segments are assigned to the closest existing
speaker. Default 10.
embedding_model (str): The pyannote embedding model to use.
Default "pyannote/wespeaker-voxceleb-resnet34-LM".
hf_token (str or None): HuggingFace token for gated model access.
"""
def __init__(
self,
similarity_threshold=0.55,
max_speakers=10,
embedding_model="pyannote/wespeaker-voxceleb-resnet34-LM",
hf_token=None,
):
self.similarity_threshold = similarity_threshold
self.max_speakers = max_speakers
self.speakers = {} # speaker_id -> embedding (averaged)
self._speaker_count = 0
self._model = None
self._embedding_model_name = embedding_model
self._hf_token = hf_token
def _load_model(self):
"""Lazy-load the embedding model on first use."""
if self._model is not None:
return
try:
from pyannote.audio import Model, Inference
import torch
model = Model.from_pretrained(
self._embedding_model_name,
use_auth_token=self._hf_token,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
self._model = Inference(model, window="whole", device=torch.device(device))
logging.info(f"Speaker embedding model loaded on {device}")
except ImportError:
raise ImportError(
"pyannote.audio is required for speaker diarization. "
"Install it with: pip install pyannote.audio"
)
def _compute_embedding(self, audio_np, sample_rate=16000):
"""Compute a speaker embedding from an audio numpy array.
Args:
audio_np (np.ndarray): 1-D float32 audio samples.
sample_rate (int): Sample rate of the audio.
Returns:
np.ndarray: Speaker embedding vector, or None if audio is too short.
"""
self._load_model()
if len(audio_np) < sample_rate * 0.3:
return None
waveform = {
"waveform": __import__("torch").tensor(audio_np).unsqueeze(0),
"sample_rate": sample_rate,
}
embedding = self._model(waveform)
return embedding / np.linalg.norm(embedding)
@staticmethod
def _cosine_similarity(a, b):
"""Compute cosine similarity between two vectors."""
return float(np.dot(a, b))
def identify_speaker(self, audio_np, sample_rate=16000):
"""Identify or create a speaker from an audio segment.
Args:
audio_np (np.ndarray): 1-D float32 audio for the segment.
sample_rate (int): Sample rate. Default 16000.
Returns:
str or None: Speaker label (e.g. "SPEAKER_00"), or None if
the audio is too short to embed.
"""
embedding = self._compute_embedding(audio_np, sample_rate)
if embedding is None:
return None
best_speaker = None
best_sim = -1.0
for speaker_id, stored_emb in self.speakers.items():
sim = self._cosine_similarity(embedding, stored_emb)
if sim > best_sim:
best_sim = sim
best_speaker = speaker_id
if best_sim >= self.similarity_threshold:
# Update running average for the matched speaker
self.speakers[best_speaker] = (
self.speakers[best_speaker] * 0.9 + embedding * 0.1
)
# Re-normalize
self.speakers[best_speaker] /= np.linalg.norm(self.speakers[best_speaker])
return best_speaker
if len(self.speakers) >= self.max_speakers:
# Assign to closest speaker
return best_speaker if best_speaker else f"SPEAKER_{self._speaker_count:02d}"
# Create a new speaker
speaker_id = f"SPEAKER_{self._speaker_count:02d}"
self._speaker_count += 1
self.speakers[speaker_id] = embedding
return speaker_id
def reset(self):
"""Reset all speaker state."""
self.speakers.clear()
self._speaker_count = 0
+122
View File
@@ -0,0 +1,122 @@
"""
Prometheus metrics for WhisperLive server.
Exposes a /metrics HTTP endpoint on a configurable port for Prometheus scraping.
All metrics are optional — the server works fine without prometheus_client installed.
"""
import logging
import threading
try:
from prometheus_client import (
Counter,
Gauge,
Histogram,
start_http_server,
)
CONNECTIONS_TOTAL = Counter(
"whisperlive_connections_total",
"Total WebSocket connections accepted",
)
CONNECTIONS_ACTIVE = Gauge(
"whisperlive_connections_active",
"Currently active WebSocket connections",
)
CONNECTIONS_REJECTED = Counter(
"whisperlive_connections_rejected_total",
"Connections rejected (server full or auth failure)",
["reason"],
)
TRANSCRIPTION_LATENCY = Histogram(
"whisperlive_transcription_latency_seconds",
"Time to transcribe a single audio chunk",
buckets=(0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0),
)
AUDIO_PROCESSED = Counter(
"whisperlive_audio_processed_seconds_total",
"Total seconds of audio processed",
)
SEGMENTS_EMITTED = Counter(
"whisperlive_segments_emitted_total",
"Total transcription segments sent to clients",
["completed"],
)
REST_REQUESTS = Counter(
"whisperlive_rest_requests_total",
"Total REST API requests",
["endpoint", "status"],
)
ERRORS = Counter(
"whisperlive_errors_total",
"Total errors by type",
["type"],
)
_AVAILABLE = True
except ImportError:
_AVAILABLE = False
def is_available():
"""Check if prometheus_client is installed."""
return _AVAILABLE
def start_metrics_server(port=9091):
"""Start the Prometheus metrics HTTP server on the given port.
Args:
port (int): Port to serve /metrics on. Default 9091.
"""
if not _AVAILABLE:
logging.warning("prometheus_client not installed; metrics endpoint disabled")
return
try:
start_http_server(port)
logging.info(f"Prometheus metrics available at http://0.0.0.0:{port}/metrics")
except Exception as e:
logging.error(f"Failed to start metrics server: {e}")
def track_connection_opened():
if _AVAILABLE:
CONNECTIONS_TOTAL.inc()
CONNECTIONS_ACTIVE.inc()
def track_connection_closed():
if _AVAILABLE:
CONNECTIONS_ACTIVE.dec()
def track_connection_rejected(reason="full"):
if _AVAILABLE:
CONNECTIONS_REJECTED.labels(reason=reason).inc()
def track_transcription_latency(seconds):
if _AVAILABLE:
TRANSCRIPTION_LATENCY.observe(seconds)
def track_audio_processed(seconds):
if _AVAILABLE:
AUDIO_PROCESSED.inc(seconds)
def track_segment_emitted(completed=True):
if _AVAILABLE:
SEGMENTS_EMITTED.labels(completed=str(completed).lower()).inc()
def track_rest_request(endpoint="/v1/audio/transcriptions", status="200"):
if _AVAILABLE:
REST_REQUESTS.labels(endpoint=endpoint, status=str(status)).inc()
def track_error(error_type="transcription"):
if _AVAILABLE:
ERRORS.labels(type=error_type).inc()
+226 -33
View File
@@ -1,6 +1,7 @@
import os import os
import time import time
import threading import threading
import collections
import queue import queue
import json import json
import functools import functools
@@ -8,14 +9,17 @@ import logging
import shutil import shutil
import tempfile import tempfile
from typing import Optional, List from typing import Optional, List
from fastapi import FastAPI, UploadFile, Form from fastapi import FastAPI, UploadFile, Form, Request
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import PlainTextResponse, JSONResponse from fastapi.responses import JSONResponse
from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse
import uvicorn import uvicorn
from faster_whisper import WhisperModel from faster_whisper import WhisperModel
import torch import torch
from enum import Enum from enum import Enum
from whisper_live import metrics as wl_metrics
from typing import List, Optional from typing import List, Optional
import numpy as np import numpy as np
from websockets.sync.server import serve from websockets.sync.server import serve
@@ -39,6 +43,7 @@ class ClientManager:
self.start_times = {} self.start_times = {}
self.max_clients = max_clients self.max_clients = max_clients
self.max_connection_time = max_connection_time self.max_connection_time = max_connection_time
self.lock = threading.Lock()
def add_client(self, websocket, client): def add_client(self, websocket, client):
""" """
@@ -48,8 +53,9 @@ class ClientManager:
websocket: The websocket associated with the client to add. websocket: The websocket associated with the client to add.
client: The client object to be added and tracked. client: The client object to be added and tracked.
""" """
self.clients[websocket] = client with self.lock:
self.start_times[websocket] = time.time() self.clients[websocket] = client
self.start_times[websocket] = time.time()
def get_client(self, websocket): def get_client(self, websocket):
""" """
@@ -61,9 +67,10 @@ class ClientManager:
Returns: Returns:
The client object if found, False otherwise. The client object if found, False otherwise.
""" """
if websocket in self.clients: with self.lock:
return self.clients[websocket] if websocket in self.clients:
return False return self.clients[websocket]
return False
def remove_client(self, websocket): def remove_client(self, websocket):
""" """
@@ -73,10 +80,11 @@ class ClientManager:
Args: Args:
websocket: The websocket associated with the client to be removed. websocket: The websocket associated with the client to be removed.
""" """
client = self.clients.pop(websocket, None) with self.lock:
client = self.clients.pop(websocket, None)
self.start_times.pop(websocket, None)
if client: if client:
client.cleanup() client.cleanup()
self.start_times.pop(websocket, None)
def get_wait_time(self): def get_wait_time(self):
""" """
@@ -85,11 +93,12 @@ class ClientManager:
Returns: Returns:
The estimated wait time in minutes for new clients to connect. Returns 0 if there are available slots. The estimated wait time in minutes for new clients to connect. Returns 0 if there are available slots.
""" """
wait_time = None with self.lock:
for start_time in self.start_times.values(): wait_time = None
current_client_time_remaining = self.max_connection_time - (time.time() - start_time) for start_time in self.start_times.values():
if wait_time is None or current_client_time_remaining < wait_time: current_client_time_remaining = self.max_connection_time - (time.time() - start_time)
wait_time = current_client_time_remaining if wait_time is None or current_client_time_remaining < wait_time:
wait_time = current_client_time_remaining
return wait_time / 60 if wait_time is not None else 0 return wait_time / 60 if wait_time is not None else 0
def is_server_full(self, websocket, options): def is_server_full(self, websocket, options):
@@ -103,12 +112,18 @@ class ClientManager:
Returns: Returns:
True if the server is full, False otherwise. True if the server is full, False otherwise.
""" """
if len(self.clients) >= self.max_clients: with self.lock:
wait_time = self.get_wait_time() if len(self.clients) >= self.max_clients:
response = {"uid": options["uid"], "status": "WAIT", "message": wait_time} wait_time = None
websocket.send(json.dumps(response)) for start_time in self.start_times.values():
return True remaining = self.max_connection_time - (time.time() - start_time)
return False if wait_time is None or remaining < wait_time:
wait_time = remaining
wait_time_minutes = wait_time / 60 if wait_time is not None else 0
response = {"uid": options["uid"], "status": "WAIT", "message": wait_time_minutes}
websocket.send(json.dumps(response))
return True
return False
def is_client_timeout(self, websocket): def is_client_timeout(self, websocket):
""" """
@@ -120,10 +135,12 @@ class ClientManager:
Returns: Returns:
True if the client's connection time has exceeded the maximum limit, False otherwise. True if the client's connection time has exceeded the maximum limit, False otherwise.
""" """
elapsed_time = time.time() - self.start_times[websocket] with self.lock:
if elapsed_time >= self.max_connection_time: elapsed_time = time.time() - self.start_times[websocket]
self.clients[websocket].disconnect() client = self.clients.get(websocket)
logging.warning(f"Client with uid '{self.clients[websocket].client_uid}' disconnected due to overtime.") if elapsed_time >= self.max_connection_time and client:
client.disconnect()
logging.warning(f"Client with uid '{client.client_uid}' disconnected due to overtime.")
return True return True
return False return False
@@ -160,6 +177,8 @@ class TranscriptionServer:
self.use_vad = True self.use_vad = True
self.single_model = False self.single_model = False
self.batch_config = None self.batch_config = None
self.raw_pcm_input = False
self.segment_post_processor = None
def initialize_client( def initialize_client(
self, websocket, options, faster_whisper_custom_model_path, self, websocket, options, faster_whisper_custom_model_path,
@@ -177,7 +196,7 @@ class TranscriptionServer:
if enable_translation: if enable_translation:
target_language = options.get("target_language", "fr") target_language = options.get("target_language", "fr")
translation_queue = queue.Queue() translation_queue = queue.Queue(maxsize=ServeClientBase.MAX_TRANSLATION_QUEUE_SIZE)
from whisper_live.backend.translation_backend import ServeClientTranslation from whisper_live.backend.translation_backend import ServeClientTranslation
translation_client = ServeClientTranslation( translation_client = ServeClientTranslation(
client_uid=options["uid"], client_uid=options["uid"],
@@ -274,7 +293,10 @@ class TranscriptionServer:
clip_audio=options.get("clip_audio", False), clip_audio=options.get("clip_audio", False),
same_output_threshold=options.get("same_output_threshold", 10), same_output_threshold=options.get("same_output_threshold", 10),
cache_path=self.cache_path, cache_path=self.cache_path,
translation_queue=translation_queue translation_queue=translation_queue,
hotwords=options.get("hotwords"),
diarization=self._create_diarizer(options),
word_timestamps=options.get("word_timestamps", False),
) )
logging.info("Running faster_whisper backend.") logging.info("Running faster_whisper backend.")
@@ -297,12 +319,35 @@ class TranscriptionServer:
if client is None: if client is None:
raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.") raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
# Attach segment post-processor if configured
if self.segment_post_processor is not None:
client.segment_post_processor = self.segment_post_processor
if translation_client: if translation_client:
client.translation_client = translation_client client.translation_client = translation_client
client.translation_thread = translation_thread client.translation_thread = translation_thread
self.client_manager.add_client(websocket, client) self.client_manager.add_client(websocket, client)
def _create_diarizer(self, options):
"""Create a SpeakerDiarizer if the client requested diarization.
Returns:
SpeakerDiarizer or None
"""
if not options.get("enable_diarization", False):
return None
try:
from whisper_live.diarization import SpeakerDiarizer
return SpeakerDiarizer(
similarity_threshold=options.get("diarization_threshold", 0.55),
max_speakers=options.get("max_speakers", 10),
hf_token=options.get("hf_token"),
)
except ImportError:
logging.warning("pyannote.audio not installed; diarization disabled")
return None
def get_audio_from_websocket(self, websocket): def get_audio_from_websocket(self, websocket):
""" """
Receives audio buffer from websocket and creates a numpy array out of it. Receives audio buffer from websocket and creates a numpy array out of it.
@@ -316,6 +361,9 @@ class TranscriptionServer:
frame_data = websocket.recv() frame_data = websocket.recv()
if frame_data == b"END_OF_AUDIO": if frame_data == b"END_OF_AUDIO":
return False return False
if self.raw_pcm_input:
audio_np = np.frombuffer(frame_data, dtype=np.int16)
return audio_np.astype(np.float32) / 32768.0
return np.frombuffer(frame_data, dtype=np.float32) return np.frombuffer(frame_data, dtype=np.float32)
def handle_new_connection(self, websocket, faster_whisper_custom_model_path, def handle_new_connection(self, websocket, faster_whisper_custom_model_path,
@@ -327,6 +375,7 @@ class TranscriptionServer:
self.use_vad = options.get('use_vad') self.use_vad = options.get('use_vad')
if self.client_manager.is_server_full(websocket, options): if self.client_manager.is_server_full(websocket, options):
wl_metrics.track_connection_rejected(reason="full")
websocket.close() websocket.close()
return False # Indicates that the connection should not continue return False # Indicates that the connection should not continue
@@ -334,6 +383,7 @@ class TranscriptionServer:
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE) self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
self.initialize_client(websocket, options, faster_whisper_custom_model_path, self.initialize_client(websocket, options, faster_whisper_custom_model_path,
whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session) whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session)
wl_metrics.track_connection_opened()
return True return True
except json.JSONDecodeError: except json.JSONDecodeError:
logging.error("Failed to decode JSON from client") logging.error("Failed to decode JSON from client")
@@ -412,8 +462,58 @@ class TranscriptionServer:
if self.client_manager.get_client(websocket): if self.client_manager.get_client(websocket):
self.cleanup(websocket) self.cleanup(websocket)
websocket.close() websocket.close()
wl_metrics.track_connection_closed()
del websocket del websocket
def _stream_transcription(self, file, language, prompt, temperature,
timestamp_granularities,
faster_whisper_custom_model_path):
"""Return a StreamingResponse that yields SSE events per segment."""
async def _sse_generator():
tmp_path = None
try:
suffix = os.path.splitext(file.filename)[1] or ".wav"
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
shutil.copyfileobj(file.file, tmp)
tmp_path = tmp.name
device = "cuda" if torch.cuda.is_available() else "cpu"
compute_type = "float16" if device == "cuda" else "int8"
model_name = faster_whisper_custom_model_path or "small"
transcriber = WhisperModel(model_name, device=device, compute_type=compute_type)
segments, info = transcriber.transcribe(
tmp_path,
language=language,
initial_prompt=prompt,
temperature=temperature,
vad_filter=False,
word_timestamps=(timestamp_granularities and "word" in timestamp_granularities),
)
for seg in segments:
seg_dict = {
"id": seg.id,
"start": seg.start,
"end": seg.end,
"text": seg.text.strip(),
}
if timestamp_granularities and "word" in timestamp_granularities:
seg_dict["words"] = [
{"word": w.word, "start": w.start, "end": w.end, "probability": w.probability}
for w in seg.words
]
yield f"data: {json.dumps(seg_dict)}\n\n"
yield "data: [DONE]\n\n"
except Exception as e:
yield f"data: {json.dumps({'error': str(e)})}\n\n"
finally:
if tmp_path and os.path.exists(tmp_path):
os.unlink(tmp_path)
return StreamingResponse(_sse_generator(), media_type="text/event-stream")
def run(self, def run(self,
host, host,
port=9090, port=9090,
@@ -431,7 +531,12 @@ class TranscriptionServer:
cors_origins: Optional[str] = None, cors_origins: Optional[str] = None,
batch_enabled=False, batch_enabled=False,
batch_max_size=8, batch_max_size=8,
batch_window_ms=50): batch_window_ms=50,
raw_pcm_input=False,
metrics_port: int = 0,
api_key: Optional[str] = None,
rate_limit_rpm: int = 0,
segment_post_processor=None):
""" """
Run the transcription server. Run the transcription server.
@@ -447,8 +552,25 @@ class TranscriptionServer:
batch_window_ms (int): Maximum time in milliseconds to wait for batch_window_ms (int): Maximum time in milliseconds to wait for
the batch to fill after the first request arrives. Defaults the batch to fill after the first request arrives. Defaults
to 50. to 50.
segment_post_processor (callable, optional): A callable that receives
a transcription segment dict and returns a modified segment dict.
Applied to every segment before sending to the client. Useful for
plugging in custom post-processing (e.g. formatting, redaction).
Defaults to None.
""" """
self.cache_path = cache_path self.cache_path = cache_path
self.raw_pcm_input = raw_pcm_input
if max_clients < 1:
raise ValueError(f"max_clients must be >= 1, got {max_clients}")
if max_connection_time <= 0:
raise ValueError(f"max_connection_time must be > 0, got {max_connection_time}")
if batch_enabled and batch_max_size < 1:
raise ValueError(f"batch_max_size must be >= 1, got {batch_max_size}")
if batch_enabled and batch_window_ms < 0:
raise ValueError(f"batch_window_ms must be >= 0, got {batch_window_ms}")
self.segment_post_processor = segment_post_processor
self.client_manager = ClientManager(max_clients, max_connection_time) self.client_manager = ClientManager(max_clients, max_connection_time)
if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path): if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
if "/" not in faster_whisper_custom_model_path: if "/" not in faster_whisper_custom_model_path:
@@ -477,6 +599,10 @@ class TranscriptionServer:
if not BackendType.is_valid(backend): if not BackendType.is_valid(backend):
raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}") raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
# Start Prometheus metrics endpoint if port is specified
if metrics_port > 0:
wl_metrics.start_metrics_server(metrics_port)
# New OpenAI-compatible REST API (toggleable via enable_rest boolean) # New OpenAI-compatible REST API (toggleable via enable_rest boolean)
if enable_rest: if enable_rest:
app = FastAPI(title="WhisperLive OpenAI-Compatible API") app = FastAPI(title="WhisperLive OpenAI-Compatible API")
@@ -489,6 +615,34 @@ class TranscriptionServer:
allow_headers=["*"], # Allows all headers allow_headers=["*"], # Allows all headers
) )
# Optional API key authentication
if api_key:
@app.middleware("http")
async def _check_api_key(request: Request, call_next):
auth = request.headers.get("Authorization", "")
if auth != f"Bearer {api_key}":
return JSONResponse({"error": "Invalid or missing API key"}, status_code=401)
return await call_next(request)
# Optional rate limiting (requests per minute per client IP)
if rate_limit_rpm > 0:
_rate_lock = threading.Lock()
_rate_buckets: dict = {} # ip -> deque of timestamps
@app.middleware("http")
async def _rate_limit(request: Request, call_next):
client_ip = request.client.host if request.client else "unknown"
now = time.time()
with _rate_lock:
bucket = _rate_buckets.setdefault(client_ip, collections.deque())
# Discard entries older than 60s
while bucket and bucket[0] < now - 60:
bucket.popleft()
if len(bucket) >= rate_limit_rpm:
return JSONResponse({"error": "Rate limit exceeded"}, status_code=429)
bucket.append(now)
return await call_next(request)
@app.post("/v1/audio/transcriptions") @app.post("/v1/audio/transcriptions")
async def transcribe( async def transcribe(
@@ -504,15 +658,31 @@ class TranscriptionServer:
include: Optional[List[str]] = Form(default=None), include: Optional[List[str]] = Form(default=None),
known_speaker_names: Optional[List[str]] = Form(default=None), known_speaker_names: Optional[List[str]] = Form(default=None),
known_speaker_references: Optional[List[str]] = Form(default=None), known_speaker_references: Optional[List[str]] = Form(default=None),
stream: bool = Form(default=False) stream: bool = Form(default=False),
hotwords: Optional[str] = Form(default=None),
): ):
if stream: if stream:
return JSONResponse({"error": "Streaming not supported in this backend."}, status_code=400) return self._stream_transcription(
if chunking_strategy or known_speaker_names or known_speaker_references: file, language, prompt, temperature,
logging.warning("Diarization/chunking params ignored; not supported.") timestamp_granularities,
faster_whisper_custom_model_path,
)
ignored_params = []
if chunking_strategy:
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
if known_speaker_names:
ignored_params.append("known_speaker_names")
if known_speaker_references:
ignored_params.append("known_speaker_references")
if include:
ignored_params.append(f"include={include}")
if ignored_params:
logging.warning(f"Unsupported OpenAI params ignored: {', '.join(ignored_params)}")
supported_formats = ["json", "text", "srt", "verbose_json", "vtt"] supported_formats = ["json", "text", "srt", "verbose_json", "vtt"]
if response_format not in supported_formats: if response_format not in supported_formats:
wl_metrics.track_rest_request(endpoint="transcriptions", status=400)
return JSONResponse({"error": f"Unsupported response_format. Supported: {supported_formats}"}, status_code=400) return JSONResponse({"error": f"Unsupported response_format. Supported: {supported_formats}"}, status_code=400)
if model != "whisper-1": if model != "whisper-1":
@@ -535,15 +705,18 @@ class TranscriptionServer:
initial_prompt=prompt, initial_prompt=prompt,
temperature=temperature, temperature=temperature,
vad_filter=False, vad_filter=False,
word_timestamps=(timestamp_granularities and "word" in timestamp_granularities) word_timestamps=(timestamp_granularities and "word" in timestamp_granularities),
hotwords=hotwords,
) )
text = " ".join([s.text.strip() for s in segments]) text = " ".join([s.text.strip() for s in segments])
os.unlink(tmp_path) os.unlink(tmp_path)
if response_format == "text": if response_format == "text":
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
return PlainTextResponse(text) return PlainTextResponse(text)
elif response_format == "json": elif response_format == "json":
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
return {"text": text} return {"text": text}
elif response_format == "verbose_json": elif response_format == "verbose_json":
verbose = { verbose = {
@@ -569,6 +742,7 @@ class TranscriptionServer:
if timestamp_granularities and "word" in timestamp_granularities: if timestamp_granularities and "word" in timestamp_granularities:
seg_dict["words"] = [{"word": w.word, "start": w.start, "end": w.end, "probability": w.probability} for w in seg.words] seg_dict["words"] = [{"word": w.word, "start": w.start, "end": w.end, "probability": w.probability} for w in seg.words]
verbose["segments"].append(seg_dict) verbose["segments"].append(seg_dict)
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
return verbose return verbose
elif response_format in ["srt", "vtt"]: elif response_format in ["srt", "vtt"]:
output = [] output = []
@@ -579,8 +753,11 @@ class TranscriptionServer:
output.append(f"{i}\n{start.replace('.', ',')} --> {end.replace('.', ',')}\n{seg.text.strip()}\n") output.append(f"{i}\n{start.replace('.', ',')} --> {end.replace('.', ',')}\n{seg.text.strip()}\n")
else: # vtt else: # vtt
output.append(f"{start} --> {end}\n{seg.text.strip()}\n") output.append(f"{start} --> {end}\n{seg.text.strip()}\n")
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
return PlainTextResponse("\n".join(output)) return PlainTextResponse("\n".join(output))
except Exception as e: except Exception as e:
wl_metrics.track_rest_request(endpoint="transcriptions", status=500)
wl_metrics.track_error("rest_transcription")
return JSONResponse({"error": str(e)}, status_code=500) return JSONResponse({"error": str(e)}, status_code=500)
threading.Thread( threading.Thread(
@@ -592,6 +769,21 @@ class TranscriptionServer:
logging.info(f"✅ OpenAI-Compatible API started on http://0.0.0.0:{rest_port}") logging.info(f"✅ OpenAI-Compatible API started on http://0.0.0.0:{rest_port}")
# Original WebSocket server (always supported) # Original WebSocket server (always supported)
extra_ws_kwargs = {}
if api_key:
def _ws_auth(path, request_headers):
auth = request_headers.get("Authorization", "")
token_param = None
# Check query string for token parameter
if "?" in path:
from urllib.parse import urlparse, parse_qs
parsed = urlparse(path)
token_param = parse_qs(parsed.query).get("token", [None])[0]
if auth == f"Bearer {api_key}" or token_param == api_key:
return None # Allow connection
return (401, [("Content-Type", "text/plain")], b"Unauthorized\n")
extra_ws_kwargs["process_request"] = _ws_auth
with serve( with serve(
functools.partial( functools.partial(
self.recv_audio, self.recv_audio,
@@ -602,7 +794,8 @@ class TranscriptionServer:
trt_py_session=trt_py_session, trt_py_session=trt_py_session,
), ),
host, host,
port port,
**extra_ws_kwargs,
) as server: ) as server:
server.serve_forever() server.serve_forever()
+1 -2
View File
@@ -1,4 +1,3 @@
import os
import textwrap import textwrap
import scipy import scipy
import numpy as np import numpy as np
@@ -8,7 +7,7 @@ from pathlib import Path
def clear_screen(): def clear_screen():
"""Clears the console screen.""" """Clears the console screen."""
os.system("cls" if os.name == "nt" else "clear") print("\033[H\033[2J", end="", flush=True)
def print_transcript(text, translated=False, timestamps=False): def print_transcript(text, translated=False, timestamps=False):