Compare commits
70 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 955665401a | |||
| 17b1639a9b | |||
| 06ec02445d | |||
| daf3da8633 | |||
| 2e466b765a | |||
| dfef376086 | |||
| 26845cabe9 | |||
| c0f101f0d1 | |||
| ec1dc7c6aa | |||
| a71c578570 | |||
| f4f1b1d8be | |||
| ad4d314b79 | |||
| 8c0caf1be0 | |||
| d9459ebf2d | |||
| 5b577b34e4 | |||
| 2debc0ee80 | |||
| 9f7a043d8b | |||
| ee64458194 | |||
| 056774ea50 | |||
| e4160d2d06 | |||
| 471c3fd6b4 | |||
| ac7a9f849c | |||
| ecb052c873 | |||
| ee3113a507 | |||
| 81ff199a70 | |||
| d0ad362b20 | |||
| 44940e2834 | |||
| 32c1b18c9f | |||
| 582d5426d6 | |||
| 1814dd9bfa | |||
| 0747529910 | |||
| cef3352c5b | |||
| 3516f663e5 | |||
| d006d4abb3 | |||
| 6242ad7f81 | |||
| 5334ea0f7a | |||
| b648bcb2a4 | |||
| e86d98dd80 | |||
| c5ec7f4a99 | |||
| 52005b94cb | |||
| f2f769532b | |||
| cdc661ce28 | |||
| 8396763444 | |||
| 8ac98dceec | |||
| 8bde966c1e | |||
| dc4a707f9a | |||
| c028c4b584 | |||
| 4e31f8c61b | |||
| 18de3eacc7 | |||
| 18b897277f | |||
| 3d63e82571 | |||
| ced4bdb737 | |||
| 4210697ca6 | |||
| 445bf26e85 | |||
| b534f9d249 | |||
| 9a71a95ca8 | |||
| 485b211072 | |||
| 19847784ae | |||
| 52d94bf1fb | |||
| 1c663d0bba | |||
| 298a01f1b0 | |||
| a6147a6745 | |||
| 18bce1864a | |||
| 9e5e4a9970 | |||
| 81cdbbca95 | |||
| f5340ddf1e | |||
| b1cd51ac8a | |||
| e41324bf03 | |||
| 31efff9330 | |||
| 68a8b57e66 |
@@ -74,8 +74,34 @@ jobs:
|
||||
# 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
|
||||
|
||||
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:
|
||||
needs: [run-tests, check-code-format]
|
||||
needs: [run-tests, check-code-format, venv-install-smoke-test]
|
||||
runs-on: ubuntu-22.04
|
||||
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
|
||||
steps:
|
||||
@@ -158,7 +184,7 @@ jobs:
|
||||
tags: ghcr.io/collabora/whisperlive-openvino:latest
|
||||
|
||||
publish-to-pypi:
|
||||
needs: [run-tests, check-code-format]
|
||||
needs: [run-tests, check-code-format, venv-install-smoke-test]
|
||||
runs-on: ubuntu-22.04
|
||||
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
|
||||
steps:
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
__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/
|
||||
Audio-Transcription-Chrome/node_modules/
|
||||
Audio-Transcription-Chrome/package-lock.json
|
||||
@@ -0,0 +1,242 @@
|
||||
'use strict';
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Chrome API mock — defined before any extension script is loaded
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const storageData = {};
|
||||
|
||||
global.chrome = {
|
||||
storage: {
|
||||
local: {
|
||||
get: jest.fn((keys, cb) => {
|
||||
if (typeof keys === 'string') {
|
||||
cb({ [keys]: storageData[keys] });
|
||||
} else if (Array.isArray(keys)) {
|
||||
const result = {};
|
||||
keys.forEach(k => { result[k] = storageData[k]; });
|
||||
cb(result);
|
||||
} else {
|
||||
const result = {};
|
||||
Object.keys(keys).forEach(k => {
|
||||
result[k] = k in storageData ? storageData[k] : keys[k];
|
||||
});
|
||||
cb(result);
|
||||
}
|
||||
}),
|
||||
set: jest.fn((obj, cb) => {
|
||||
Object.assign(storageData, obj);
|
||||
if (cb) cb();
|
||||
}),
|
||||
},
|
||||
},
|
||||
runtime: {
|
||||
sendMessage: jest.fn(),
|
||||
onMessage: { addListener: jest.fn() },
|
||||
getURL: jest.fn(path => `chrome-extension://fake-id/${path}`),
|
||||
id: 'fake-extension-id',
|
||||
},
|
||||
tabs: {
|
||||
query: jest.fn(),
|
||||
get: jest.fn(),
|
||||
create: jest.fn(),
|
||||
remove: jest.fn(),
|
||||
sendMessage: jest.fn(),
|
||||
},
|
||||
tabCapture: { capture: jest.fn() },
|
||||
scripting: { executeScript: jest.fn() },
|
||||
};
|
||||
|
||||
// Flush all pending microtasks and one macrotask round so async click
|
||||
// handlers (which await at least one Promise inside) complete fully.
|
||||
const flushPromises = () => new Promise(resolve => setTimeout(resolve, 0));
|
||||
|
||||
function resetStorage() {
|
||||
Object.keys(storageData).forEach(k => delete storageData[k]);
|
||||
}
|
||||
|
||||
function buildPopupDOM() {
|
||||
document.body.innerHTML = `
|
||||
<div id="startCapture"></div>
|
||||
<div id="stopCapture"></div>
|
||||
<input type="checkbox" id="useServerCheckbox">
|
||||
<input type="checkbox" id="useVadCheckbox">
|
||||
<input type="checkbox" id="saveCaptionsCheckbox">
|
||||
<select id="languageDropdown">
|
||||
<option value="" selected></option>
|
||||
<option value="en">English</option>
|
||||
</select>
|
||||
<select id="taskDropdown">
|
||||
<option value="transcribe" selected>Transcribe</option>
|
||||
</select>
|
||||
<select id="modelSizeDropdown">
|
||||
<option value="small" selected>Small</option>
|
||||
<option value="large-v3">Large-v3</option>
|
||||
</select>
|
||||
<select id="captionLinesDropdown">
|
||||
<option value="3" selected>3</option>
|
||||
</select>
|
||||
`;
|
||||
}
|
||||
|
||||
function loadPopup() {
|
||||
jest.resetModules();
|
||||
require('../popup.js');
|
||||
document.dispatchEvent(new Event('DOMContentLoaded'));
|
||||
}
|
||||
|
||||
function clickStart(tabId = 1) {
|
||||
chrome.tabs.query.mockImplementation((_, cb) => cb([{ id: tabId }]));
|
||||
document.getElementById('startCapture').click();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 1. WebSocket URL construction — pure logic extracted from options.js
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('WebSocket URL construction', () => {
|
||||
function buildWsUrl(host, port) {
|
||||
return port
|
||||
? `ws://${host}:${port}/`
|
||||
: `wss://${host}/ws`;
|
||||
}
|
||||
|
||||
test('local server: ws:// with port and trailing slash', () => {
|
||||
expect(buildWsUrl('localhost', '9090')).toBe('ws://localhost:9090/');
|
||||
});
|
||||
|
||||
test('Modal service (empty port): wss:// with /ws path', () => {
|
||||
expect(buildWsUrl('boxerab--aavaaz-live-livetranscriber-web.modal.run', '')).toBe(
|
||||
'wss://boxerab--aavaaz-live-livetranscriber-web.modal.run/ws'
|
||||
);
|
||||
});
|
||||
|
||||
test('custom host+port stays on ws://', () => {
|
||||
expect(buildWsUrl('my-server.example.com', '7090')).toBe(
|
||||
'ws://my-server.example.com:7090/'
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 2. popup.js — host/port selection based on checkbox
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('popup.js host/port selection', () => {
|
||||
beforeEach(() => {
|
||||
resetStorage();
|
||||
buildPopupDOM();
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
test('checkbox unchecked → localhost:9090', async () => {
|
||||
document.getElementById('useServerCheckbox').checked = false;
|
||||
loadPopup();
|
||||
clickStart();
|
||||
await flushPromises();
|
||||
|
||||
const call = chrome.runtime.sendMessage.mock.calls.find(
|
||||
c => c[0] && c[0].action === 'startCapture'
|
||||
);
|
||||
expect(call).toBeDefined();
|
||||
expect(call[0].host).toBe('localhost');
|
||||
expect(call[0].port).toBe('9090');
|
||||
});
|
||||
|
||||
test('checkbox checked → Modal host with empty port', async () => {
|
||||
document.getElementById('useServerCheckbox').checked = true;
|
||||
loadPopup();
|
||||
clickStart();
|
||||
await flushPromises();
|
||||
|
||||
const call = chrome.runtime.sendMessage.mock.calls.find(
|
||||
c => c[0] && c[0].action === 'startCapture'
|
||||
);
|
||||
expect(call).toBeDefined();
|
||||
expect(call[0].host).toBe('boxerab--aavaaz-live-livetranscriber-web.modal.run');
|
||||
expect(call[0].port).toBe('');
|
||||
});
|
||||
|
||||
test('startCapture message carries language, task, and modelSize', async () => {
|
||||
document.getElementById('languageDropdown').value = 'en';
|
||||
document.getElementById('taskDropdown').value = 'transcribe';
|
||||
document.getElementById('modelSizeDropdown').value = 'large-v3';
|
||||
loadPopup();
|
||||
clickStart();
|
||||
await flushPromises();
|
||||
|
||||
const call = chrome.runtime.sendMessage.mock.calls.find(
|
||||
c => c[0] && c[0].action === 'startCapture'
|
||||
);
|
||||
expect(call).toBeDefined();
|
||||
expect(call[0].task).toBe('transcribe');
|
||||
expect(call[0].modelSize).toBe('large-v3');
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 3. popup.js — button state management
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('popup.js button state', () => {
|
||||
beforeEach(() => {
|
||||
resetStorage();
|
||||
buildPopupDOM();
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
test('start button enabled and stop button disabled on load (not capturing)', () => {
|
||||
storageData.capturingState = { isCapturing: false };
|
||||
loadPopup();
|
||||
expect(document.getElementById('startCapture').disabled).toBeFalsy();
|
||||
expect(document.getElementById('stopCapture').disabled).toBe(true);
|
||||
});
|
||||
|
||||
test('stop button enabled on load when already capturing', () => {
|
||||
storageData.capturingState = { isCapturing: true };
|
||||
loadPopup();
|
||||
expect(document.getElementById('stopCapture').disabled).toBe(false);
|
||||
});
|
||||
|
||||
test('clicking stop sends stopCapture action to runtime', () => {
|
||||
storageData.capturingState = { isCapturing: true };
|
||||
loadPopup();
|
||||
// toggleCaptureButtons(true) disables startCapture and enables stopCapture
|
||||
document.getElementById('stopCapture').click();
|
||||
|
||||
expect(chrome.runtime.sendMessage).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ action: 'stopCapture' }),
|
||||
expect.any(Function)
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 4. popup.js — storage state restoration on open
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('popup.js storage restoration', () => {
|
||||
beforeEach(() => {
|
||||
resetStorage();
|
||||
buildPopupDOM();
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
test('restores useServerCheckbox from storage', () => {
|
||||
storageData.useServerState = true;
|
||||
loadPopup();
|
||||
expect(document.getElementById('useServerCheckbox').checked).toBe(true);
|
||||
});
|
||||
|
||||
test('restores language selection from storage', () => {
|
||||
storageData.selectedLanguage = 'en';
|
||||
loadPopup();
|
||||
expect(document.getElementById('languageDropdown').value).toBe('en');
|
||||
});
|
||||
|
||||
test('restores model size from storage', () => {
|
||||
storageData.selectedModelSize = 'large-v3';
|
||||
loadPopup();
|
||||
expect(document.getElementById('modelSizeDropdown').value).toBe('large-v3');
|
||||
});
|
||||
});
|
||||
@@ -3,6 +3,7 @@ var elem_text = null;
|
||||
|
||||
var segments = [];
|
||||
var text_segments = [];
|
||||
var captionLineCount = 3;
|
||||
var allSegments = [];
|
||||
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')) {
|
||||
return;
|
||||
}
|
||||
|
||||
elem_container = document.createElement('div');
|
||||
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:0 24px;font-family:Arial,Helvetica,sans-serif;font-size:22px;font-weight:600;line-height:30px;position:fixed;top:85%;left:50%;transform:translate(-50%,-50%);width:min(80vw,900px);min-height:' + (captionLineCount * 30) + 'px;z-index:2147483647;color:white;text-align:center;letter-spacing:0.01em;text-shadow:0 0 2px #000,0 2px 4px rgba(0,0,0,0.95);cursor:move;';
|
||||
|
||||
for (var i = 0; i < 4; i++) {
|
||||
for (var i = 0; i <= captionLineCount; i++) {
|
||||
elem_text = document.createElement('span');
|
||||
elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;';
|
||||
elem_text.style.cssText = 'position:absolute;left:50%;transform:translateX(-50%);max-width:100%;padding:2px 14px;background:rgba(0,0,0,0.72);border-radius:6px;box-decoration-break:clone;-webkit-box-decoration-break:clone;';
|
||||
elem_text.id = "t" + i;
|
||||
elem_container.appendChild(elem_text);
|
||||
|
||||
if (i == 3) {
|
||||
if (i == captionLineCount) {
|
||||
elem_text.style.top = "-1000px"
|
||||
}
|
||||
}
|
||||
@@ -164,7 +166,7 @@ function get_lines(elem, line_height) {
|
||||
var divHeight = elem.offsetHeight;
|
||||
var lines = divHeight / line_height;
|
||||
|
||||
var original_text = elem.innerHTML;
|
||||
var original_text = elem.textContent;
|
||||
|
||||
var words = original_text.split(' ');
|
||||
var segments = [];
|
||||
@@ -174,7 +176,7 @@ function get_lines(elem, line_height) {
|
||||
for (var i = 0; i < words.length; i++)
|
||||
{
|
||||
segment += words[i] + ' ';
|
||||
elem.innerHTML = segment;
|
||||
elem.textContent = segment;
|
||||
divHeight = elem.offsetHeight;
|
||||
|
||||
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)
|
||||
segments.push(line_segment);
|
||||
|
||||
elem.innerHTML = original_text;
|
||||
elem.textContent = original_text;
|
||||
|
||||
return segments;
|
||||
|
||||
@@ -196,7 +198,7 @@ function get_lines(elem, line_height) {
|
||||
|
||||
function remove_element() {
|
||||
var elem = document.getElementById('transcription')
|
||||
for (var i = 0; i < 4; i++) {
|
||||
for (var i = 0; i <= captionLineCount; i++) {
|
||||
document.getElementById("t" + i).remove();
|
||||
}
|
||||
elem.remove()
|
||||
@@ -205,6 +207,7 @@ function remove_element() {
|
||||
chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
const { type, data } = request;
|
||||
const saveCaptions = data.saveCaptions;
|
||||
const captionLines = data.captionLines || captionLineCount;
|
||||
|
||||
if (type === "STOP") {
|
||||
if (saveCaptions === true) {
|
||||
@@ -234,7 +237,7 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
return true;
|
||||
}
|
||||
|
||||
init_element();
|
||||
init_element(captionLines);
|
||||
|
||||
try {
|
||||
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, "");
|
||||
|
||||
var elem = document.getElementById('t3');
|
||||
var elem = document.getElementById('t' + captionLineCount);
|
||||
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 divHeight = elem.offsetHeight;
|
||||
var lines = divHeight / line_height;
|
||||
@@ -274,29 +277,29 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
text_segments = [];
|
||||
text_segments = get_lines(elem, line_height);
|
||||
|
||||
elem.innerHTML = '';
|
||||
elem.textContent = '';
|
||||
|
||||
if (text_segments.length > 2) {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i];
|
||||
if (text_segments.length > captionLineCount - 1) {
|
||||
for (var i = 0; i < captionLineCount; i++) {
|
||||
document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
|
||||
}
|
||||
} else {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = '';
|
||||
for (var i = 0; i < captionLineCount; i++) {
|
||||
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++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[i];
|
||||
document.getElementById('t' + i).textContent = text_segments[i];
|
||||
}
|
||||
} else {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i];
|
||||
for (var i = 0; i < captionLineCount; 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 elem = document.getElementById('t' + i);
|
||||
|
||||
@@ -129,7 +129,10 @@ async function startRecord(option) {
|
||||
return;
|
||||
}
|
||||
|
||||
socket = new WebSocket(`ws://${option.host}:${option.port}/`);
|
||||
const wsUrl = option.port
|
||||
? `ws://${option.host}:${option.port}/`
|
||||
: `wss://${option.host}/ws`;
|
||||
socket = new WebSocket(wsUrl);
|
||||
isServerReady = false;
|
||||
let language = option.language;
|
||||
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"name": "audio-transcription-chrome",
|
||||
"version": "1.0.0",
|
||||
"description": "Audio Transcription is a Chrome extension that allows users to capture any audio playing on the current tab and transcribe it using OpenAI-whisper in real time. Users will have the option to do voice activity detection as well to not send audio to server when there is no speech.",
|
||||
"main": "audiopreprocessor.js",
|
||||
"scripts": {
|
||||
"test": "jest"
|
||||
},
|
||||
"jest": {
|
||||
"testEnvironment": "jsdom"
|
||||
},
|
||||
"keywords": [],
|
||||
"author": "",
|
||||
"license": "ISC",
|
||||
"devDependencies": {
|
||||
"jest": "^30.4.2",
|
||||
"jest-environment-jsdom": "^30.4.1"
|
||||
}
|
||||
}
|
||||
@@ -23,6 +23,14 @@
|
||||
<input type="checkbox" id="saveCaptionsCheckbox">
|
||||
<label for="saveCaptions">Download SRT file at Stop Capture</label>
|
||||
</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">
|
||||
<label for="languageDropdown">Select Language:</label>
|
||||
<select id="languageDropdown">
|
||||
|
||||
@@ -9,9 +9,11 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
const languageDropdown = document.getElementById('languageDropdown');
|
||||
const taskDropdown = document.getElementById('taskDropdown');
|
||||
const modelSizeDropdown = document.getElementById('modelSizeDropdown');
|
||||
const captionLinesDropdown = document.getElementById('captionLinesDropdown');
|
||||
let selectedLanguage = null;
|
||||
let selectedTask = taskDropdown.value;
|
||||
let selectedModelSize = modelSizeDropdown.value;
|
||||
let selectedCaptionLines = captionLinesDropdown.value;
|
||||
|
||||
// Add click event listeners to the buttons
|
||||
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
|
||||
async function startCapture() {
|
||||
// Ignore click if the button is disabled
|
||||
@@ -81,8 +90,8 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
let port = "9090";
|
||||
const useCollaboraServer = useServerCheckbox.checked;
|
||||
if (useCollaboraServer){
|
||||
host = "transcription.kurg.org"
|
||||
port = "7090"
|
||||
host = "boxerab--aavaaz-live-livetranscriber-web.modal.run"
|
||||
port = ""
|
||||
}
|
||||
|
||||
chrome.runtime.sendMessage(
|
||||
@@ -96,6 +105,7 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
modelSize: selectedModelSize,
|
||||
useVad: useVadCheckbox.checked,
|
||||
saveCaptions: saveCaptionsCheckbox.checked,
|
||||
captionLines: Number(selectedCaptionLines),
|
||||
}, () => {
|
||||
// Update capturing state in storage and toggle the buttons
|
||||
chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => {
|
||||
@@ -144,6 +154,7 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
modelSizeDropdown.disabled = isCapturing;
|
||||
languageDropdown.disabled = isCapturing;
|
||||
taskDropdown.disabled = isCapturing;
|
||||
captionLinesDropdown.disabled = isCapturing;
|
||||
startButton.classList.toggle("disabled", isCapturing);
|
||||
stopButton.classList.toggle("disabled", !isCapturing);
|
||||
}
|
||||
@@ -183,6 +194,11 @@ document.addEventListener("DOMContentLoaded", function () {
|
||||
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) => {
|
||||
if (request.action === "updateSelectedLanguage") {
|
||||
const detectedLanguage = request.detectedLanguage;
|
||||
|
||||
@@ -162,6 +162,7 @@ var elem_text = null;
|
||||
|
||||
var segments = [];
|
||||
var text_segments = [];
|
||||
var captionLineCount = 3;
|
||||
|
||||
function initPopupElement() {
|
||||
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')) {
|
||||
return;
|
||||
}
|
||||
|
||||
elem_container = document.createElement('div');
|
||||
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:0 24px;font-family:Arial,Helvetica,sans-serif;font-size:22px;font-weight:600;line-height:30px;position:fixed;top:85%;left:50%;transform:translate(-50%,-50%);width:min(80vw,900px);min-height:' + (captionLineCount * 30) + 'px;z-index:2147483647;color:white;text-align:center;letter-spacing:0.01em;text-shadow:0 0 2px #000,0 2px 4px rgba(0,0,0,0.95);cursor:move;';
|
||||
|
||||
for (var i = 0; i < 4; i++) {
|
||||
for (var i = 0; i <= captionLineCount; i++) {
|
||||
elem_text = document.createElement('span');
|
||||
elem_text.style.cssText = 'position: absolute;padding-left:16px;padding-right:16px;';
|
||||
elem_text.style.cssText = 'position:absolute;left:50%;transform:translateX(-50%);max-width:100%;padding:2px 14px;background:rgba(0,0,0,0.72);border-radius:6px;box-decoration-break:clone;-webkit-box-decoration-break:clone;';
|
||||
elem_text.id = "t" + i;
|
||||
elem_container.appendChild(elem_text);
|
||||
|
||||
if (i == 3) {
|
||||
if (i == captionLineCount) {
|
||||
elem_text.style.top = "-1000px"
|
||||
}
|
||||
}
|
||||
@@ -286,7 +288,7 @@ function get_lines(elem, line_height) {
|
||||
var divHeight = elem.offsetHeight;
|
||||
var lines = divHeight / line_height;
|
||||
|
||||
var original_text = elem.innerHTML;
|
||||
var original_text = elem.textContent;
|
||||
|
||||
var words = original_text.split(' ');
|
||||
var segments = [];
|
||||
@@ -296,7 +298,7 @@ function get_lines(elem, line_height) {
|
||||
for (var i = 0; i < words.length; i++)
|
||||
{
|
||||
segment += words[i] + ' ';
|
||||
elem.innerHTML = segment;
|
||||
elem.textContent = segment;
|
||||
divHeight = elem.offsetHeight;
|
||||
|
||||
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)
|
||||
segments.push(line_segment);
|
||||
|
||||
elem.innerHTML = original_text;
|
||||
elem.textContent = original_text;
|
||||
|
||||
return segments;
|
||||
|
||||
@@ -318,7 +320,7 @@ function get_lines(elem, line_height) {
|
||||
|
||||
function remove_element() {
|
||||
var elem = document.getElementById('transcription')
|
||||
for (var i = 0; i < 4; i++) {
|
||||
for (var i = 0; i <= captionLineCount; i++) {
|
||||
document.getElementById("t" + i).remove();
|
||||
}
|
||||
elem.remove()
|
||||
@@ -327,6 +329,7 @@ function remove_element() {
|
||||
browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
const { action, data } = request;
|
||||
const saveCaption = data.saveCaption || false;
|
||||
const captionLines = data.captionLines || captionLineCount;
|
||||
|
||||
if (action === "startCapture") {
|
||||
isCapturing = true;
|
||||
@@ -364,7 +367,7 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
|
||||
} else if (action === "show_transcript"){
|
||||
if (!isCapturing) return;
|
||||
init_element();
|
||||
init_element(captionLines);
|
||||
message = JSON.parse(data.data);
|
||||
message = message["segments"];
|
||||
|
||||
@@ -391,10 +394,10 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
}
|
||||
text = text.replace(/(\r\n|\n|\r)/gm, "");
|
||||
|
||||
var elem = document.getElementById('t3');
|
||||
elem.innerHTML = text;
|
||||
var elem = document.getElementById('t' + captionLineCount);
|
||||
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 divHeight = elem.offsetHeight;
|
||||
var lines = divHeight / line_height;
|
||||
@@ -402,29 +405,29 @@ browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
text_segments = [];
|
||||
text_segments = get_lines(elem, line_height);
|
||||
|
||||
elem.innerHTML = '';
|
||||
elem.textContent = '';
|
||||
|
||||
if (text_segments.length > 2) {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i];
|
||||
if (text_segments.length > captionLineCount - 1) {
|
||||
for (var i = 0; i < captionLineCount; i++) {
|
||||
document.getElementById('t' + i).textContent = text_segments[text_segments.length - captionLineCount + i];
|
||||
}
|
||||
} else {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = '';
|
||||
for (var i = 0; i < captionLineCount; i++) {
|
||||
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++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[i];
|
||||
document.getElementById('t' + i).textContent = text_segments[i];
|
||||
}
|
||||
} else {
|
||||
for (var i = 0; i < 3; i++) {
|
||||
document.getElementById('t' + i).innerHTML = text_segments[text_segments.length - 3 + i];
|
||||
for (var i = 0; i < captionLineCount; 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 elem = document.getElementById('t' + i);
|
||||
|
||||
@@ -24,6 +24,14 @@
|
||||
<label for="saveCaption">Download SRT file at Stop Capture</label>
|
||||
</div>
|
||||
<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">
|
||||
<label for="languageDropdown">Select Language:</label>
|
||||
<select id="languageDropdown">
|
||||
|
||||
@@ -8,9 +8,11 @@ document.addEventListener("DOMContentLoaded", function() {
|
||||
const languageDropdown = document.getElementById('languageDropdown');
|
||||
const taskDropdown = document.getElementById('taskDropdown');
|
||||
const modelSizeDropdown = document.getElementById('modelSizeDropdown');
|
||||
const captionLinesDropdown = document.getElementById('captionLinesDropdown');
|
||||
let selectedLanguage = null;
|
||||
let selectedTask = taskDropdown.value;
|
||||
let selectedModelSize = modelSizeDropdown.value;
|
||||
let selectedCaptionLines = captionLinesDropdown.value;
|
||||
|
||||
|
||||
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() {
|
||||
let host = "localhost";
|
||||
let port = "9090";
|
||||
@@ -93,6 +102,7 @@ document.addEventListener("DOMContentLoaded", function() {
|
||||
modelSize: selectedModelSize,
|
||||
useVad: useVadCheckbox.checked,
|
||||
saveCaption: saveCaptionCheckbox.checked,
|
||||
captionLines: Number(selectedCaptionLines),
|
||||
}
|
||||
});
|
||||
toggleCaptureButtons(true);
|
||||
@@ -136,6 +146,7 @@ document.addEventListener("DOMContentLoaded", function() {
|
||||
modelSizeDropdown.disabled = isCapturing;
|
||||
languageDropdown.disabled = isCapturing;
|
||||
taskDropdown.disabled = isCapturing;
|
||||
captionLinesDropdown.disabled = isCapturing;
|
||||
startButton.classList.toggle("disabled", isCapturing);
|
||||
stopButton.classList.toggle("disabled", !isCapturing);
|
||||
}
|
||||
@@ -175,6 +186,11 @@ document.addEventListener("DOMContentLoaded", function() {
|
||||
browser.storage.local.set({ selectedModelSize });
|
||||
});
|
||||
|
||||
captionLinesDropdown.addEventListener('change', function() {
|
||||
selectedCaptionLines = captionLinesDropdown.value;
|
||||
browser.storage.local.set({ selectedCaptionLines });
|
||||
});
|
||||
|
||||
browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
|
||||
if (request.action === "updateSelectedLanguage") {
|
||||
const detectedLanguage = request.data;
|
||||
|
||||
@@ -23,6 +23,8 @@ The app streams microphone audio to a WhisperLive server via WebSocket and displ
|
||||
|
||||
## Getting Started
|
||||
|
||||
This directory contains the Swift source files for the iOS client, but it does not include a generated `.xcodeproj` or `.xcodeworkspace`. Create a new Xcode project and add these files to it.
|
||||
|
||||
1. Clone the repository (your fork):
|
||||
|
||||
```bash
|
||||
@@ -30,16 +32,24 @@ The app streams microphone audio to a WhisperLive server via WebSocket and displ
|
||||
cd whisperlive/Audio-Transcription-iOS
|
||||
```
|
||||
|
||||
2. Open the `.xcodeproj` or `.xcodeworkspace` in Xcode
|
||||
2. In Xcode, choose **File ▸ New ▸ Project…**, then create an iOS **App** project with SwiftUI.
|
||||
|
||||
3. Add the following to your `Info.plist`:
|
||||
3. Add the Swift files from this directory to the new app target:
|
||||
|
||||
- `AudioStream.swift`
|
||||
- `AudioWebSocket.swift`
|
||||
- `ContentView.swift`
|
||||
- `RecordingViewModel.swift`
|
||||
- `WhisperLive_iOS_ClientApp.swift`
|
||||
|
||||
4. Use `WhisperLive-iOS-Client-Info.plist` as a reference for your app's `Info.plist`, or add the microphone usage description manually:
|
||||
|
||||
```xml
|
||||
<key>NSMicrophoneUsageDescription</key>
|
||||
<string>This app requires microphone access for transcription.</string>
|
||||
```
|
||||
|
||||
4. Run the app on a physical device (recommended)
|
||||
5. Run the app on a physical device (recommended)
|
||||
|
||||
## Running on a Physical Device (with Free Apple ID)
|
||||
|
||||
@@ -92,12 +102,12 @@ Now you can run and debug the app on your real device!
|
||||
## Folder Structure
|
||||
```
|
||||
Audio-Transcription-iOS/
|
||||
├── AudioViewModel.swift
|
||||
├── AudioStreamer.swift
|
||||
├── AudioStream.swift
|
||||
├── AudioWebSocket.swift
|
||||
├── RecordingView.swift
|
||||
├── ContentView.swift
|
||||
├── RecordingViewModel.swift
|
||||
├── WhisperLive-iOS-Client-Info.plist
|
||||
├── WhisperLive_iOS_ClientApp.swift
|
||||
├── Info.plist
|
||||
├── README.md
|
||||
```
|
||||
|
||||
|
||||
@@ -17,34 +17,40 @@ input from microphone and pre-recorded audio files.
|
||||
- [Getting Started](#getting-started)
|
||||
- [Running the Server](#running-the-server)
|
||||
- [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)
|
||||
- [Streaming Client (Manual Audio Chunking)](#streaming-client-manual-audio-chunking)
|
||||
- [Browser Extensions](#browser-extensions)
|
||||
- [Whisper Live Server in Docker](#whisper-live-server-in-docker)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Future Work](#future-work)
|
||||
- [Blog Posts](#blog-posts)
|
||||
- [Contact](#contact)
|
||||
- [Citations](#citations)
|
||||
|
||||
## Installation
|
||||
- Install PortAudio
|
||||
- Install PortAudio (required system dependency for microphone input via PyAudio)
|
||||
```bash
|
||||
bash scripts/setup.sh
|
||||
```
|
||||
On Debian/Ubuntu this installs `portaudio19-dev`, on Fedora `portaudio-devel`, on macOS it uses Homebrew (`portaudio`).
|
||||
|
||||
- Install 3.12 venv (on Fedora `sudo dnf install -y python3.12 python3.12-pip`)
|
||||
|
||||
```bash
|
||||
python3.12 -m venv whisper_env
|
||||
source whisper_env/bin/activate
|
||||
```
|
||||
|
||||
- Install whisper-live from pip
|
||||
```bash
|
||||
pip install whisper-live
|
||||
```
|
||||
|
||||
|
||||
- Install 3.12 venv on Fedora
|
||||
|
||||
```bash
|
||||
sudo dnf install -y python3.12 python3.12-pip
|
||||
python3.12 -m venv whisper_env
|
||||
source whisper_env/bin/activate
|
||||
```
|
||||
|
||||
|
||||
### OpenAI REST interface
|
||||
|
||||
#### Server
|
||||
@@ -101,6 +107,7 @@ python3 run_server.py -p 9090 \
|
||||
--max_clients 4 \
|
||||
--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_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).
|
||||
@@ -111,6 +118,9 @@ python3 run_server.py -p 9090 \
|
||||
python3 run_server.py -p 9090 -b openvino
|
||||
```
|
||||
|
||||
### Setting up AMD ROCm for faster_whisper backend
|
||||
- Please follow [ROCm_whisper readme](https://github.com/collabora/WhisperLive/blob/main/ROCm_whisper.md) for setup of AMD ROCm GPU support with the CTranslate2 ROCm wheel.
|
||||
|
||||
|
||||
#### Controlling OpenMP Threads
|
||||
To control the number of threads used by OpenMP, you can set the `OMP_NUM_THREADS` environment variable. This is useful for managing CPU resources and ensuring consistent performance. If not specified, `OMP_NUM_THREADS` is set to `1` by default. You can change this by using the `--omp_num_threads` argument:
|
||||
@@ -162,6 +172,7 @@ client = TranscriptionClient(
|
||||
mute_audio_playback=False, # Only used for file input, False by Default
|
||||
enable_translation=True,
|
||||
target_language="hi",
|
||||
initial_prompt=None, # Add context for the model, e.g. 'Jane Doe context'
|
||||
)
|
||||
```
|
||||
It connects to the server running on localhost at port 9090. Using a multilingual model, language for the transcription will be automatically detected. You can also use the language option to specify the target language for the transcription, in this case, English ("en"). The translate option should be set to `True` if we want to translate from the source language to English and `False` if we want to transcribe in the source language.
|
||||
@@ -186,6 +197,130 @@ 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")
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
||||
The OpenAI-compatible REST endpoint also accepts `known_speaker_names` and uploaded `known_speaker_references` multipart fields. When speaker fields are supplied with `response_format="verbose_json"`, segments include a `speaker` field.
|
||||
|
||||
#### 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]. Clients can also set `audio_format` in the initial websocket options to `float32` (default), `int16`, or `uint8`.
|
||||
|
||||
## Streaming Client (manual audio streaming from any source)
|
||||
|
||||
`StreamingTranscriptionClient` lets you push raw PCM audio bytes from any source — a live microphone capture loop, a network stream, an audio pipeline — and receive transcripts via callbacks as speech is detected. Unlike `TranscriptionClient`, it does not manage audio capture internally; you control when and how audio is fed.
|
||||
|
||||
A runnable example that reads from an audio file and streams the chunks is at [`examples/manual_audio_chunking.py`](examples/manual_audio_chunking.py):
|
||||
|
||||
```bash
|
||||
python examples/manual_audio_chunking.py --file assets/jfk.flac
|
||||
```
|
||||
|
||||
Example usage:
|
||||
|
||||
```python
|
||||
from whisper_live.client import StreamingTranscriptionClient
|
||||
|
||||
client = StreamingTranscriptionClient(
|
||||
"localhost", 9090,
|
||||
lang="en",
|
||||
model="small",
|
||||
on_session_started=lambda: print("Server ready"),
|
||||
on_partial_transcript=lambda text, segs: print(f"… {text}", end="\r"),
|
||||
on_committed_transcript=lambda text, segs: print(f"✓ {text}"),
|
||||
on_error=lambda e: print(f"Error: {e}"),
|
||||
on_close=lambda: print("Closed"),
|
||||
)
|
||||
|
||||
with client:
|
||||
for chunk in my_audio_source: # any cadence, any chunk size
|
||||
client.send(chunk, pcm_format="int16")
|
||||
```
|
||||
|
||||
Audio must be **mono, 16 kHz PCM**. Two formats are accepted:
|
||||
|
||||
| `pcm_format` | Description |
|
||||
|---|---|
|
||||
| `"int16"` (default for raw microphone data) | 16-bit signed integers, normalized internally |
|
||||
| `"float32"` | 32-bit floats in `[-1, 1]`, passed through directly |
|
||||
|
||||
NumPy arrays can be sent with `send_array()`:
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
samples = np.frombuffer(raw_bytes, dtype=np.int16)
|
||||
client.send_array(samples)
|
||||
```
|
||||
|
||||
**Callbacks**
|
||||
|
||||
| Callback | Signature | When fired |
|
||||
|---|---|---|
|
||||
| `on_session_started` | `() -> None` | Server handshake complete, ready to receive audio |
|
||||
| `on_partial_transcript` | `(text, segments) -> None` | In-progress segment updated |
|
||||
| `on_committed_transcript` | `(text, segments) -> None` | Segment finalized |
|
||||
| `on_translation` | `(text, segments) -> None` | Translated segment ready (requires `enable_translation=True`) |
|
||||
| `on_error` | `(error) -> None` | WebSocket error |
|
||||
| `on_close` | `() -> None` | Connection closed |
|
||||
|
||||
## Browser Extensions
|
||||
- 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
|
||||
@@ -206,19 +341,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.
|
||||
```bash
|
||||
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
|
||||
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 int4 # int4 weight only quantization
|
||||
|
||||
# Run server with small.en
|
||||
# Run server with small.en (pick one engine)
|
||||
python3 run_server.py --port 9090 \
|
||||
--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_int8"
|
||||
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
|
||||
# or int8 / 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
|
||||
@@ -226,12 +362,36 @@ Refer to [`ios-client`](https://github.com/collabora/WhisperLive/tree/main/Audio
|
||||
docker run -it --device=/dev/dri -p 9090:9090 ghcr.io/collabora/whisperlive-openvino
|
||||
```
|
||||
|
||||
- AMD ROCm (faster-whisper on AMD GPU via CTranslate2 ROCm wheel)
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.rocm -t whisperlive-rocm .
|
||||
docker run --rm -it --device=/dev/kfd --device=/dev/dri \
|
||||
--group-add "$(getent group video | cut -d: -f3)" \
|
||||
--group-add "$(getent group render | cut -d: -f3)" \
|
||||
-p 9090:9090 whisperlive-rocm
|
||||
```
|
||||
|
||||
- CPU
|
||||
- Faster-whisper
|
||||
```bash
|
||||
docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu:latest
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
#### macOS OpenMP runtime conflict
|
||||
On macOS, especially on Intel Macs, `faster_whisper`/`ctranslate2` can conflict with OpenMP runtimes loaded by other Python packages. If the server aborts with a duplicate OpenMP runtime error, run the server with `KMP_DUPLICATE_LIB_OK=TRUE`:
|
||||
|
||||
```bash
|
||||
KMP_DUPLICATE_LIB_OK=TRUE python3 run_server.py --port 9090 \
|
||||
--backend faster_whisper \
|
||||
--max_clients 4 \
|
||||
--max_connection_time 600 \
|
||||
--no_single_model
|
||||
```
|
||||
|
||||
This workaround is intended for local development and testing. For production deployments, prefer using a clean environment that loads only one OpenMP runtime.
|
||||
|
||||
## Future Work
|
||||
- [x] Add translation to other languages on top of transcription.
|
||||
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
# WhisperLive-ROCm
|
||||
Run WhisperLive's `faster_whisper` backend on AMD GPUs using the official [CTranslate2 ROCm wheel](https://github.com/OpenNMT/CTranslate2/releases). Tested on Radeon AI PRO R9700 (gfx1201/RDNA4) and Ryzen AI Max+ 395 / Radeon 8060S (gfx1151/Strix Halo).
|
||||
|
||||
## Docker Installation (recommended)
|
||||
- Install [docker](https://docs.docker.com/engine/install/)
|
||||
|
||||
- Build and run the WhisperLive ROCm image:
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.rocm -t whisperlive-rocm .
|
||||
docker run --rm -it \
|
||||
--device=/dev/kfd --device=/dev/dri \
|
||||
--group-add "$(getent group video | cut -d: -f3)" \
|
||||
--group-add "$(getent group render | cut -d: -f3)" \
|
||||
-p 9090:9090 whisperlive-rocm
|
||||
```
|
||||
|
||||
## Native Installation
|
||||
|
||||
### Prerequisites
|
||||
- AMD GPU with ROCm support (see [supported GPUs](https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html))
|
||||
- ROCm 7.2+ installed ([installation guide](https://rocm.docs.amd.com/en/latest/deploy/linux/quick_start.html))
|
||||
- User in `video` and `render` groups (`sudo usermod -aG video,render $USER`, re-login)
|
||||
- Python 3.12
|
||||
|
||||
### Verify ROCm is working
|
||||
```bash
|
||||
rocminfo | grep -E 'Name:|gfx'
|
||||
# Should show your GPU, e.g. "Name: gfx1151" or "Name: gfx1201"
|
||||
```
|
||||
|
||||
### Install CTranslate2 ROCm wheel
|
||||
The default `pip install ctranslate2` installs a CUDA-only wheel. Replace it with the official ROCm wheel from the [CTranslate2 releases page](https://github.com/OpenNMT/CTranslate2/releases):
|
||||
|
||||
```bash
|
||||
# Download the ROCm wheels archive (v4.8.0)
|
||||
curl -LO https://github.com/OpenNMT/CTranslate2/releases/download/v4.8.0/rocm-python-wheels-Linux.zip
|
||||
|
||||
# Extract the Python 3.12 wheel
|
||||
unzip -j rocm-python-wheels-Linux.zip 'temp-linux/ctranslate2-*-cp312-*manylinux*x86_64.whl'
|
||||
|
||||
# Install (replaces any existing ctranslate2)
|
||||
pip install ctranslate2-*-cp312-*.whl
|
||||
```
|
||||
|
||||
### Install WhisperLive server requirements
|
||||
```bash
|
||||
pip install -r requirements/server.txt
|
||||
```
|
||||
|
||||
### Verify GPU is visible to CTranslate2
|
||||
```bash
|
||||
python -c "import ctranslate2; print('devices:', ctranslate2.get_cuda_device_count())"
|
||||
```
|
||||
Expected output: `devices: 1` (CTranslate2 uses the name "cuda" even on ROCm).
|
||||
|
||||
If you see `devices: 0`, check:
|
||||
- Your user is in `video` and `render` groups (re-login after adding)
|
||||
- `/dev/kfd` exists and is accessible
|
||||
- The ROCm wheel was installed (not the default PyPI CUDA-only one)
|
||||
|
||||
## Run WhisperLive Server with ROCm
|
||||
```bash
|
||||
python3 run_server.py --port 9090 --backend faster_whisper
|
||||
```
|
||||
|
||||
The server automatically uses the AMD GPU when the CTranslate2 ROCm wheel is installed. For multi-GPU systems, use `HIP_VISIBLE_DEVICES=N` to select a specific GPU.
|
||||
@@ -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)
|
||||
@@ -0,0 +1,44 @@
|
||||
# docker/Dockerfile.rocm
|
||||
#
|
||||
# WhisperLive faster_whisper backend on AMD ROCm GPUs.
|
||||
# Uses the official CTranslate2 ROCm wheel (ships kernels for gfx803 through
|
||||
# gfx1201 including Strix Halo gfx1151 and RDNA4 gfx1200/1201).
|
||||
#
|
||||
# Build:
|
||||
# docker build -f docker/Dockerfile.rocm -t whisperlive-rocm .
|
||||
#
|
||||
# Run (expose the WebSocket port; add --enable_rest --rest_port 8000 -p 8000:8000 for REST):
|
||||
# docker run --rm -it \
|
||||
# --device=/dev/kfd --device=/dev/dri \
|
||||
# --group-add "$(getent group video | cut -d: -f3)" \
|
||||
# --group-add "$(getent group render | cut -d: -f3)" \
|
||||
# -p 9090:9090 whisperlive-rocm
|
||||
|
||||
FROM rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0
|
||||
|
||||
ARG DEBIAN_FRONTEND=noninteractive
|
||||
ARG CT2_WHEEL_URL=https://github.com/OpenNMT/CTranslate2/releases/download/v4.8.0/rocm-python-wheels-Linux.zip
|
||||
|
||||
RUN apt-get update -qq && \
|
||||
apt-get install -y --no-install-recommends curl unzip portaudio19-dev && \
|
||||
apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install the CTranslate2 ROCm wheel (official release artifact).
|
||||
# This replaces any CUDA-only ctranslate2 and enables GPU on AMD.
|
||||
RUN curl -sL "${CT2_WHEEL_URL}" -o /tmp/ct2-rocm.zip && \
|
||||
unzip -j /tmp/ct2-rocm.zip 'temp-linux/ctranslate2-*-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl' -d /tmp && \
|
||||
pip install --no-cache-dir --force-reinstall /tmp/ctranslate2-*-cp312-*.whl && \
|
||||
rm -f /tmp/ct2-rocm.zip /tmp/ctranslate2-*.whl
|
||||
|
||||
# Install server requirements
|
||||
COPY requirements/server.txt /app/
|
||||
RUN pip install --no-cache-dir -r server.txt && rm server.txt
|
||||
|
||||
COPY whisper_live /app/whisper_live
|
||||
COPY run_server.py /app
|
||||
|
||||
EXPOSE 9090
|
||||
|
||||
CMD ["python", "run_server.py", "--port", "9090", "--backend", "faster_whisper"]
|
||||
@@ -0,0 +1,78 @@
|
||||
"""
|
||||
Manual audio chunking example for WhisperLive.
|
||||
|
||||
Streams an audio file to a running WhisperLive server in real-time sized chunks,
|
||||
printing partial transcripts when speech is detected and committed transcripts
|
||||
when each segment is finalized.
|
||||
|
||||
Usage:
|
||||
python examples/manual_audio_chunking.py --file assets/jfk.flac
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import wave
|
||||
|
||||
try:
|
||||
from whisper_live.client import StreamingTranscriptionClient
|
||||
from whisper_live.utils import resample
|
||||
except ImportError: # just in case whisper_live isn't installed.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
print("[INFO] whisper_live not installed or the current version does not have StreamingTranscriptionClient. Will attempt to import from local source.")
|
||||
from whisper_live.client import StreamingTranscriptionClient
|
||||
from whisper_live.utils import resample
|
||||
|
||||
SAMPLE_RATE = 16000
|
||||
|
||||
|
||||
def stream_audio_file(path: str, client: StreamingTranscriptionClient, chunk_ms: int = 50) -> None:
|
||||
"""Read an audio file, resample to 16 kHz mono if needed, and pace chunks in real time."""
|
||||
resampled_path = resample(path)
|
||||
try:
|
||||
with wave.open(resampled_path, "rb") as wf:
|
||||
frames_per_chunk = SAMPLE_RATE * chunk_ms // 1000
|
||||
chunk_duration = frames_per_chunk / SAMPLE_RATE
|
||||
while chunk := wf.readframes(frames_per_chunk):
|
||||
client.send(chunk, pcm_format="int16")
|
||||
time.sleep(chunk_duration)
|
||||
finally:
|
||||
os.remove(resampled_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Stream an audio file to WhisperLive.")
|
||||
parser.add_argument("--file", "-f", required=True, help="Audio file to transcribe (any format supported by ffmpeg).")
|
||||
parser.add_argument("--server", "-s", default="localhost")
|
||||
parser.add_argument("--port", "-p", type=int, default=9090)
|
||||
parser.add_argument("--model", "-m", default="small")
|
||||
parser.add_argument("--lang", "-l", default="en")
|
||||
parser.add_argument("--chunk_ms", type=int, default=50, help="Chunk size in ms.")
|
||||
args = parser.parse_args()
|
||||
|
||||
client = StreamingTranscriptionClient(
|
||||
args.server, args.port,
|
||||
lang=args.lang,
|
||||
model=args.model,
|
||||
on_session_started=lambda: print("[INFO] Server ready.\n"),
|
||||
on_partial_transcript=lambda text, _: print(f"\r… {text:<80}", end="", flush=True),
|
||||
on_committed_transcript=lambda text, _: print(f"\r✓ {text:<80}"),
|
||||
on_error=lambda e: print(f"\n[ERROR] {e}"),
|
||||
on_close=lambda: print("\n[INFO] Connection closed."),
|
||||
)
|
||||
|
||||
with client:
|
||||
print(f"[INFO] Streaming {args.file} in {args.chunk_ms} ms chunks.")
|
||||
stream_audio_file(args.file, client, chunk_ms=args.chunk_ms)
|
||||
|
||||
print("\n[INFO] Final transcript:")
|
||||
for seg in client.transcript:
|
||||
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()}")
|
||||
seg = client.last_partial
|
||||
if seg:
|
||||
print(f" [{float(seg['start']):.2f}s → {float(seg['end']):.2f}s] {seg['text'].strip()} (partial)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
python_files = test_*.py
|
||||
python_classes = Test*
|
||||
python_functions = test_*
|
||||
@@ -1,6 +1,7 @@
|
||||
faster-whisper==1.2.0
|
||||
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
|
||||
kaldialign
|
||||
soundfile
|
||||
@@ -8,8 +9,9 @@ scipy
|
||||
av
|
||||
jiwer
|
||||
evaluate
|
||||
numpy<2
|
||||
numpy>=1.26.4,<2.5
|
||||
openai-whisper==20250625
|
||||
pyannote.audio
|
||||
tokenizers==0.20.3
|
||||
transformers[torch]
|
||||
sentencepiece
|
||||
|
||||
+9
-2
@@ -33,7 +33,8 @@ if __name__ == '__main__':
|
||||
help='Language code for transcription, e.g., "en" for English.')
|
||||
parser.add_argument('--translate', '-t',
|
||||
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',
|
||||
action='store_true',
|
||||
help='Mute audio playback during transcription.')
|
||||
@@ -42,7 +43,7 @@ if __name__ == '__main__':
|
||||
help='Save the output recording, only used for microphone input.')
|
||||
parser.add_argument('--enable_translation',
|
||||
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',
|
||||
type=str,
|
||||
default='fr',
|
||||
@@ -57,6 +58,12 @@ if __name__ == '__main__':
|
||||
|
||||
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(
|
||||
args.server,
|
||||
args.port,
|
||||
|
||||
@@ -84,6 +84,31 @@ if __name__ == "__main__":
|
||||
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()
|
||||
|
||||
if args.backend == "tensorrt":
|
||||
@@ -113,4 +138,8 @@ if __name__ == "__main__":
|
||||
batch_enabled=args.batch_inference,
|
||||
batch_max_size=args.batch_max_size,
|
||||
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,
|
||||
)
|
||||
+1
-1
@@ -19,7 +19,7 @@ elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
|
||||
source /etc/os-release
|
||||
fi
|
||||
|
||||
if [[ "${ID:-}" == "fedora" ]]; then
|
||||
if [[ "$(command -v dnf)" ]]; then
|
||||
echo "Detected Fedora, using dnf for installation"
|
||||
dnf install -y portaudio-devel wget
|
||||
else
|
||||
|
||||
@@ -28,8 +28,11 @@ setup(
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3 :: Only",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"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",
|
||||
],
|
||||
packages=find_packages(
|
||||
@@ -43,11 +46,13 @@ setup(
|
||||
),
|
||||
install_requires=[
|
||||
"PyAudio",
|
||||
"av",
|
||||
"faster-whisper==1.2.0",
|
||||
"torch",
|
||||
"torchaudio",
|
||||
"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",
|
||||
"websocket-client",
|
||||
"numba",
|
||||
@@ -56,12 +61,24 @@ setup(
|
||||
"soundfile",
|
||||
"tokenizers==0.20.3",
|
||||
"librosa",
|
||||
"numpy==1.26.4",
|
||||
"numpy>=1.26.4,<2.5",
|
||||
"openvino",
|
||||
"openvino-genai",
|
||||
"openvino-tokenizers",
|
||||
"optimum",
|
||||
"optimum-intel",
|
||||
"fastapi",
|
||||
"uvicorn",
|
||||
"python-multipart",
|
||||
# CTranslate2 (faster-whisper's backend) is hard-linked against
|
||||
# libcublas.so.12 / libcudnn.so.9 but doesn't declare the matching
|
||||
# wheels as runtime deps. torch >=2.12 also dropped the cu12
|
||||
# wheels in favor of cu13, so users no longer get cu12 transitively.
|
||||
# Without these wheels GPU inference dies at first transcription:
|
||||
# ERROR: Library libcublas.so.12 is not found or cannot be loaded
|
||||
# Skip only for CPU-only inference.
|
||||
"nvidia-cublas-cu12; sys_platform == 'linux'",
|
||||
"nvidia-cudnn-cu12; sys_platform == 'linux'",
|
||||
],
|
||||
python_requires=">=3.9"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,644 @@
|
||||
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 WaitTrackingEvent:
|
||||
"""Threading event that records when wait() is entered."""
|
||||
|
||||
def __init__(self):
|
||||
self._event = threading.Event()
|
||||
self.wait_started = threading.Event()
|
||||
|
||||
def wait(self, timeout=None):
|
||||
self.wait_started.set()
|
||||
return self._event.wait(timeout)
|
||||
|
||||
def set(self):
|
||||
self._event.set()
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._event, name)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
def test_exception_releases_lock_without_signaling_frames_ready(self):
|
||||
ws = MagicMock()
|
||||
client = ConcreteServeClient(client_uid="test", websocket=ws)
|
||||
client.frames_np = np.array([0.1], dtype=np.float32)
|
||||
|
||||
with patch("whisper_live.backend.base.np.concatenate", side_effect=RuntimeError("boom")):
|
||||
with self.assertRaisesRegex(RuntimeError, "boom"):
|
||||
client.add_frames(np.array([0.2], dtype=np.float32))
|
||||
|
||||
self.assertFalse(client.lock.locked())
|
||||
self.assertFalse(client.frames_ready.is_set())
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def _supports_thread_time():
|
||||
thread_time = getattr(time, "thread_time", None)
|
||||
if thread_time is None:
|
||||
return False
|
||||
try:
|
||||
thread_time()
|
||||
except NotImplementedError:
|
||||
return False
|
||||
return True
|
||||
|
||||
class TestSpeechToTextWaitingBehavior(unittest.TestCase):
|
||||
"""Tests the first-frame wait behavior in speech_to_text()."""
|
||||
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
||||
self.client.frames_ready = WaitTrackingEvent()
|
||||
self.transcribe_called = threading.Event()
|
||||
self.thread_started = threading.Event()
|
||||
self.cpu_used = None
|
||||
self.thread = None
|
||||
|
||||
def tearDown(self):
|
||||
if self.thread is not None and self.thread.is_alive():
|
||||
self.client.exit = True
|
||||
# release wait() directly so a broken cleanup() cannot hang the test process
|
||||
self.client.frames_ready.set()
|
||||
self.thread.join(timeout=1.0)
|
||||
|
||||
def _start_speech_thread(self, target=None):
|
||||
self.thread = threading.Thread(target=target or self.client.speech_to_text)
|
||||
self.thread.start()
|
||||
return self.thread
|
||||
|
||||
def _join_speech_thread(self):
|
||||
self.thread.join(timeout=1.0)
|
||||
return not self.thread.is_alive()
|
||||
|
||||
def _transcribe_once(self, input_sample):
|
||||
# mark the first processing step after wait and stop the loop
|
||||
self.transcribe_called.set()
|
||||
self.client.exit = True
|
||||
return []
|
||||
|
||||
def _measure_waiting_cpu(self):
|
||||
# measure CPU consumed by speech_to_text loop while it waits for the first frame
|
||||
self.thread_started.set()
|
||||
start_cpu = time.thread_time()
|
||||
self.client.speech_to_text()
|
||||
self.cpu_used = time.thread_time() - start_cpu
|
||||
|
||||
def test_waits_for_first_frame_before_transcribing(self):
|
||||
self.client.transcribe_audio = MagicMock(side_effect=self._transcribe_once)
|
||||
self._start_speech_thread()
|
||||
self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0))
|
||||
self.assertFalse(self.transcribe_called.is_set())
|
||||
|
||||
self.client.add_frames(np.zeros(self.client.RATE, dtype=np.float32))
|
||||
|
||||
self.assertTrue(self.transcribe_called.wait(timeout=1.0))
|
||||
self.assertTrue(self._join_speech_thread())
|
||||
|
||||
def test_cleanup_unblocks_waiting_thread_without_audio(self):
|
||||
self.client.transcribe_audio = MagicMock()
|
||||
self._start_speech_thread()
|
||||
self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0))
|
||||
|
||||
self.client.cleanup()
|
||||
|
||||
self.assertTrue(self._join_speech_thread())
|
||||
self.assertTrue(self.client.exit)
|
||||
self.client.transcribe_audio.assert_not_called()
|
||||
|
||||
def test_exit_flag_unblocks_waiting_thread_without_signal(self):
|
||||
self.client.transcribe_audio = MagicMock()
|
||||
self._start_speech_thread()
|
||||
self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0))
|
||||
|
||||
self.client.exit = True
|
||||
|
||||
self.assertTrue(self._join_speech_thread())
|
||||
self.client.transcribe_audio.assert_not_called()
|
||||
|
||||
@unittest.skipUnless(_supports_thread_time(), "time.thread_time() not supported")
|
||||
def test_waiting_for_first_frame_uses_negligible_thread_cpu(self):
|
||||
self._start_speech_thread(target=self._measure_waiting_cpu)
|
||||
self.assertTrue(self.thread_started.wait(timeout=1.0))
|
||||
|
||||
time.sleep(0.25)
|
||||
self.client.cleanup()
|
||||
|
||||
self.assertTrue(self._join_speech_thread())
|
||||
self.assertIsNotNone(self.cpu_used)
|
||||
self.assertLess(self.cpu_used, 0.1)
|
||||
|
||||
|
||||
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()
|
||||
+47
-15
@@ -4,9 +4,10 @@ import scipy
|
||||
import websocket
|
||||
import copy
|
||||
import unittest
|
||||
from io import StringIO
|
||||
from unittest.mock import patch, MagicMock
|
||||
from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
|
||||
from whisper_live.utils import resample
|
||||
from whisper_live.utils import print_transcript, resample
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@@ -43,21 +44,25 @@ class TestClientWebSocketCommunication(BaseTestCase):
|
||||
|
||||
class TestClientCallbacks(BaseTestCase):
|
||||
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.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):
|
||||
message = json.dumps(
|
||||
@@ -112,6 +117,33 @@ class TestAudioResampling(unittest.TestCase):
|
||||
os.remove(resampled_audio)
|
||||
|
||||
|
||||
class TestPrintTranscript(unittest.TestCase):
|
||||
@patch("whisper_live.utils.shutil.get_terminal_size")
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_transcript_respects_narrow_terminals(self, mock_stdout, mock_terminal_size):
|
||||
mock_terminal_size.return_value = os.terminal_size((20, 20))
|
||||
|
||||
print_transcript(["This transcript should still wrap cleanly on a narrow terminal."])
|
||||
|
||||
output_lines = [line for line in mock_stdout.getvalue().splitlines() if line.strip()]
|
||||
self.assertGreater(len(output_lines), 1)
|
||||
self.assertTrue(all(len(line) <= 20 for line in output_lines))
|
||||
|
||||
@patch("whisper_live.utils.shutil.get_terminal_size")
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_transcript_indents_timestamp_continuations(self, mock_stdout, mock_terminal_size):
|
||||
mock_terminal_size.return_value = os.terminal_size((32, 20))
|
||||
|
||||
print_transcript(
|
||||
[{"start": "00:00", "end": "00:05", "text": "This line should wrap and keep its timestamp indentation."}],
|
||||
timestamps=True,
|
||||
)
|
||||
|
||||
output_lines = [line.rstrip() for line in mock_stdout.getvalue().splitlines() if line.strip()]
|
||||
self.assertGreater(len(output_lines), 1)
|
||||
self.assertTrue(output_lines[1].startswith(" " * len("[00:00 -> 00:05] ")))
|
||||
|
||||
|
||||
class TestSendingAudioPacket(BaseTestCase):
|
||||
def test_send_packet(self):
|
||||
self.client.send_packet_to_server(self.mock_audio_packet)
|
||||
|
||||
@@ -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()
|
||||
@@ -0,0 +1,201 @@
|
||||
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_enroll_speaker_uses_known_name(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)
|
||||
self.assertTrue(d.enroll_speaker("Alice", audio))
|
||||
|
||||
self._set_embedding(d, [0.99, 0.01, 0.0])
|
||||
speaker = d.identify_speaker(audio)
|
||||
self.assertEqual(speaker, "Alice")
|
||||
|
||||
def test_speaker_names_label_new_speakers(self):
|
||||
d = self._make_diarizer(speaker_names=["Alice"])
|
||||
self._set_embedding(d, [1.0, 0.0, 0.0])
|
||||
audio = np.zeros(16000, dtype=np.float32)
|
||||
speaker = d.identify_speaker(audio)
|
||||
self.assertEqual(speaker, "Alice")
|
||||
|
||||
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()
|
||||
|
||||
|
||||
class TestRestDiarizationHelpers(unittest.TestCase):
|
||||
def test_normalize_form_list_accepts_repeated_or_comma_separated_values(self):
|
||||
from whisper_live.server import TranscriptionServer
|
||||
|
||||
self.assertEqual(
|
||||
TranscriptionServer._normalize_form_list(["Alice,Bob", "Carol"]),
|
||||
["Alice", "Bob", "Carol"],
|
||||
)
|
||||
|
||||
def test_speaker_labels_for_segments(self):
|
||||
from whisper_live.server import TranscriptionServer
|
||||
|
||||
segment = MagicMock()
|
||||
segment.start = 0.0
|
||||
segment.end = 1.0
|
||||
diarizer = MagicMock()
|
||||
diarizer.identify_speaker.return_value = "Alice"
|
||||
|
||||
speakers = TranscriptionServer._speaker_labels_for_segments(
|
||||
[segment],
|
||||
np.zeros(16000, dtype=np.float32),
|
||||
diarizer,
|
||||
)
|
||||
|
||||
self.assertEqual(speakers, {0: "Alice"})
|
||||
diarizer.identify_speaker.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
@@ -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_uint8_audio_format_normalizes_unsigned_pcm(self):
|
||||
import numpy as np
|
||||
ws = MagicMock()
|
||||
self.server.audio_formats[ws] = "uint8"
|
||||
pcm = np.array([0, 128, 255], dtype=np.uint8)
|
||||
ws.recv.return_value = pcm.tobytes()
|
||||
result = self.server.get_audio_from_websocket(ws)
|
||||
expected = (pcm.astype(np.float32) - 128.0) / 128.0
|
||||
np.testing.assert_array_almost_equal(result, expected)
|
||||
|
||||
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
|
||||
|
||||
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, File
|
||||
from fastapi.testclient import TestClient
|
||||
from typing import Optional, List
|
||||
|
||||
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[UploadFile]] = File(default=None),
|
||||
stream: bool = Form(default=False),
|
||||
):
|
||||
ignored_params = []
|
||||
if chunking_strategy:
|
||||
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
|
||||
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_supported(self):
|
||||
resp = self._post(known_speaker_names="alice")
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
ignored = resp.json()["ignored"]
|
||||
self.assertFalse(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.assertEqual(len(ignored), 1)
|
||||
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,210 @@
|
||||
import json
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import numpy as np
|
||||
|
||||
from whisper_live.client import Client, StreamingTranscriptionClient
|
||||
|
||||
|
||||
class StreamingClientTestCase(unittest.TestCase):
|
||||
@patch('whisper_live.client.websocket.WebSocketApp')
|
||||
def setUp(self, mock_websocket):
|
||||
self.mock_websocket = mock_websocket
|
||||
self.mock_ws_app = mock_websocket.return_value
|
||||
self.mock_ws_app.send = MagicMock()
|
||||
|
||||
self.committed = []
|
||||
self.partials = []
|
||||
self.session_started = []
|
||||
|
||||
self.client = StreamingTranscriptionClient(
|
||||
host='localhost',
|
||||
port=9090,
|
||||
lang="en",
|
||||
on_session_started=lambda: self.session_started.append(True),
|
||||
on_committed_transcript=lambda text, segs: self.committed.append((text, segs)),
|
||||
on_partial_transcript=lambda text, segs: self.partials.append((text, segs)),
|
||||
)
|
||||
self._inner = self.client._client
|
||||
|
||||
def tearDown(self):
|
||||
self._inner.close_websocket()
|
||||
self.mock_websocket.stop()
|
||||
|
||||
def _server_ready(self, backend="faster_whisper"):
|
||||
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||
"uid": self._inner.uid,
|
||||
"message": "SERVER_READY",
|
||||
"backend": backend,
|
||||
}))
|
||||
|
||||
def _send_segments(self, segments):
|
||||
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||
"uid": self._inner.uid,
|
||||
"segments": segments,
|
||||
}))
|
||||
|
||||
|
||||
class TestPcmFormatConversion(StreamingClientTestCase):
|
||||
def test_int16_is_normalized_to_float32(self):
|
||||
self._server_ready()
|
||||
raw = np.array([0, 16384, -32768], dtype=np.int16).tobytes()
|
||||
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||
self.client.send(raw, pcm_format="int16")
|
||||
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||
np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4)
|
||||
|
||||
def test_float32_passes_through(self):
|
||||
self._server_ready()
|
||||
raw = np.array([0.1, -0.2], dtype=np.float32).tobytes()
|
||||
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||
self.client.send(raw, pcm_format="float32")
|
||||
self.assertEqual(mock_send.call_args[0][0], raw)
|
||||
|
||||
def test_default_format_is_int16(self):
|
||||
self._server_ready()
|
||||
raw = np.array([32767], dtype=np.int16).tobytes()
|
||||
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||
self.client.send(raw)
|
||||
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||
self.assertAlmostEqual(float(sent[0]), 32767 / 32768.0, places=4)
|
||||
|
||||
def test_unsupported_format_raises(self):
|
||||
self._server_ready()
|
||||
with self.assertRaises(ValueError):
|
||||
self.client.send(b"\x00\x00", pcm_format="int8")
|
||||
|
||||
def test_send_after_close_raises(self):
|
||||
self.client._closed = True
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.client.send(b"\x00\x00", pcm_format="int16")
|
||||
|
||||
def test_send_array_normalizes_integers(self):
|
||||
with patch.object(self._inner, 'send_packet_to_server') as mock_send:
|
||||
self.client.send_array(np.array([0, 16384, -32768], dtype=np.int16))
|
||||
sent = np.frombuffer(mock_send.call_args[0][0], dtype=np.float32)
|
||||
np.testing.assert_allclose(sent, [0.0, 0.5, -1.0], atol=1e-4)
|
||||
|
||||
|
||||
class TestTranscriptDispatch(StreamingClientTestCase):
|
||||
def test_partial_then_committed(self):
|
||||
self._server_ready()
|
||||
self._send_segments([{"start": 0, "end": 1, "text": "hello", "completed": False}])
|
||||
self.assertEqual(len(self.partials), 1)
|
||||
self.assertEqual(self.partials[0][0], "hello")
|
||||
self.assertEqual(len(self.committed), 0)
|
||||
|
||||
self._send_segments([{"start": 0, "end": 1, "text": "hello world", "completed": True}])
|
||||
self.assertEqual(len(self.committed), 1)
|
||||
self.assertEqual(self.committed[0][0], "hello world")
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
|
||||
def test_committed_deduplicated(self):
|
||||
self._server_ready()
|
||||
seg = {"start": 0, "end": 1, "text": "hi", "completed": True}
|
||||
self._send_segments([seg])
|
||||
self._send_segments([seg])
|
||||
self.assertEqual(len(self.committed), 1)
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
|
||||
def test_committed_backend_agnostic(self):
|
||||
"""Committed dispatch must work for non-faster_whisper backends."""
|
||||
self._server_ready(backend="tensorrt")
|
||||
self._send_segments([{"start": 0, "end": 1, "text": "trt seg", "completed": True}])
|
||||
self.assertEqual(len(self.committed), 1)
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
|
||||
def test_last_partial_alias(self):
|
||||
self._server_ready()
|
||||
self._send_segments([{"start": 0, "end": 1, "text": "pending", "completed": False}])
|
||||
self.assertIsNotNone(self.client.last_partial)
|
||||
self.assertIs(self.client.last_partial, self.client.last_segment)
|
||||
|
||||
|
||||
class TestConnectLifecycle(StreamingClientTestCase):
|
||||
def test_connect_returns_after_ready(self):
|
||||
self._server_ready()
|
||||
self.assertIs(self.client.connect(), self.client)
|
||||
self.assertEqual(len(self.session_started), 1)
|
||||
|
||||
def test_connect_times_out(self):
|
||||
self.client._ready_timeout = 0.1
|
||||
with self.assertRaises(TimeoutError):
|
||||
self.client.connect()
|
||||
|
||||
def test_connect_raises_on_server_error(self):
|
||||
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||
"uid": self._inner.uid,
|
||||
"status": "ERROR",
|
||||
"message": "boom",
|
||||
}))
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.client.connect()
|
||||
|
||||
def test_connect_raises_when_server_full(self):
|
||||
self._inner.on_message(self.mock_ws_app, json.dumps({
|
||||
"uid": self._inner.uid,
|
||||
"status": "WAIT",
|
||||
"message": 5,
|
||||
}))
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.client.connect()
|
||||
|
||||
def test_close_sends_end_of_audio(self):
|
||||
self._server_ready()
|
||||
self._inner.recording = False # pretend server already closed
|
||||
with patch.object(self._inner, 'send_packet_to_server') as mock_send, \
|
||||
patch.object(self._inner, 'close_websocket') as mock_close:
|
||||
self.client.close()
|
||||
mock_send.assert_called_once_with(Client.END_OF_AUDIO.encode("utf-8"))
|
||||
mock_close.assert_called_once()
|
||||
|
||||
def test_close_waits_for_server_then_times_out(self):
|
||||
self._server_ready()
|
||||
self.assertTrue(self._inner.recording) # server still "processing"
|
||||
start = time.time()
|
||||
with patch.object(self._inner, 'send_packet_to_server'), \
|
||||
patch.object(self._inner, 'close_websocket') as mock_close:
|
||||
self.client.close(timeout=0.2)
|
||||
self.assertGreaterEqual(time.time() - start, 0.2)
|
||||
mock_close.assert_called_once()
|
||||
|
||||
def test_close_returns_early_when_server_closes(self):
|
||||
self._server_ready()
|
||||
|
||||
def close_soon(_msg):
|
||||
self._inner.recording = False
|
||||
|
||||
with patch.object(self._inner, 'send_packet_to_server', side_effect=close_soon), \
|
||||
patch.object(self._inner, 'close_websocket') as mock_close:
|
||||
start = time.time()
|
||||
self.client.close(timeout=10.0)
|
||||
self.assertLess(time.time() - start, 1.0)
|
||||
mock_close.assert_called_once()
|
||||
|
||||
|
||||
class TestErrorHandling(StreamingClientTestCase):
|
||||
def test_close_frame_not_reported_as_error(self):
|
||||
"""A normal CLOSE control frame (opcode 8) must not fire on_error."""
|
||||
self._server_ready()
|
||||
errors = []
|
||||
self.client._client._on_error_hook = errors.append
|
||||
close_frame = MagicMock()
|
||||
close_frame.opcode = 8
|
||||
self._inner.on_error(self.mock_ws_app, close_frame)
|
||||
self.assertEqual(errors, [])
|
||||
self.assertFalse(self._inner.server_error)
|
||||
|
||||
def test_real_error_still_reported(self):
|
||||
self._server_ready()
|
||||
errors = []
|
||||
self.client._client._on_error_hook = errors.append
|
||||
self._inner.on_error(self.mock_ws_app, RuntimeError("boom"))
|
||||
self.assertEqual(len(errors), 1)
|
||||
self.assertTrue(self._inner.server_error)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -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()
|
||||
@@ -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 @@
|
||||
__version__ = "0.8.0"
|
||||
__version__ = "0.9.0"
|
||||
|
||||
+134
-23
@@ -5,12 +5,25 @@ import time
|
||||
import queue
|
||||
import numpy as np
|
||||
|
||||
from whisper_live import metrics as wl_metrics
|
||||
|
||||
|
||||
class ServeClientBase(object):
|
||||
RATE = 16000
|
||||
SERVER_READY = "SERVER_READY"
|
||||
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."""
|
||||
FIRST_FRAME_WAIT_TIMEOUT_S = 0.1
|
||||
"""Interval in seconds for re-checking exit while waiting for the first audio frame."""
|
||||
|
||||
client_uid: str
|
||||
"""A unique identifier for the client."""
|
||||
websocket: object
|
||||
@@ -24,6 +37,9 @@ class ServeClientBase(object):
|
||||
same_output_threshold: int
|
||||
"""Number of repeated outputs before considering it as a valid segment."""
|
||||
|
||||
MAX_TRANSCRIPT_LENGTH = 500
|
||||
MAX_TRANSLATION_QUEUE_SIZE = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client_uid,
|
||||
@@ -33,6 +49,8 @@ class ServeClientBase(object):
|
||||
clip_audio=False,
|
||||
same_output_threshold=10,
|
||||
translation_queue=None,
|
||||
diarization=None,
|
||||
word_timestamps=False,
|
||||
):
|
||||
self.client_uid = client_uid
|
||||
self.websocket = websocket
|
||||
@@ -40,6 +58,8 @@ class ServeClientBase(object):
|
||||
self.no_speech_thresh = no_speech_thresh
|
||||
self.clip_audio = clip_audio
|
||||
self.same_output_threshold = same_output_threshold
|
||||
self.diarization = diarization
|
||||
self.word_timestamps = word_timestamps
|
||||
|
||||
self.frames = b""
|
||||
self.timestamp_offset = 0.0
|
||||
@@ -54,15 +74,24 @@ class ServeClientBase(object):
|
||||
self.end_time_for_same_output = None
|
||||
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
|
||||
self.lock = threading.Lock()
|
||||
self.frames_ready = threading.Event()
|
||||
|
||||
def speech_to_text(self):
|
||||
"""
|
||||
Process an audio stream in an infinite loop, continuously transcribing the speech.
|
||||
|
||||
This method continuously receives audio frames, performs real-time transcription, and sends
|
||||
transcribed segments to the client via a WebSocket connection.
|
||||
transcribed segments to the client via a WebSocket connection. The loop blocks until the first
|
||||
audio frame arrives when a client is connected but still idle.
|
||||
|
||||
If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
|
||||
It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
|
||||
@@ -78,6 +107,8 @@ class ServeClientBase(object):
|
||||
break
|
||||
|
||||
if self.frames_np is None:
|
||||
while self.frames_np is None and not self.exit:
|
||||
self.frames_ready.wait(timeout=self.FIRST_FRAME_WAIT_TIMEOUT_S)
|
||||
continue
|
||||
|
||||
if self.clip_audio:
|
||||
@@ -89,16 +120,20 @@ class ServeClientBase(object):
|
||||
continue
|
||||
try:
|
||||
input_sample = input_bytes.copy()
|
||||
t0 = time.time()
|
||||
result = self.transcribe_audio(input_sample)
|
||||
|
||||
if result is None or self.language is None:
|
||||
self.timestamp_offset += duration
|
||||
time.sleep(0.25) # wait for voice activity, result is None when no voice activity
|
||||
continue
|
||||
wl_metrics.track_transcription_latency(time.time() - t0)
|
||||
wl_metrics.track_audio_processed(duration)
|
||||
self.handle_transcription_output(result, duration)
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
|
||||
wl_metrics.track_error("transcription")
|
||||
time.sleep(0.01)
|
||||
|
||||
def transcribe_audio(self):
|
||||
@@ -107,7 +142,7 @@ class ServeClientBase(object):
|
||||
def handle_transcription_output(self, result, duration):
|
||||
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.
|
||||
|
||||
@@ -115,18 +150,25 @@ class ServeClientBase(object):
|
||||
start (float): The start 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.
|
||||
speaker (str, optional): Speaker label from diarization.
|
||||
words (list, optional): Word-level timestamps and probabilities.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary representing the formatted transcription segment, including
|
||||
'start' and 'end' times as strings with three decimal places and the 'text'
|
||||
of the transcription.
|
||||
"""
|
||||
return {
|
||||
seg = {
|
||||
'start': "{:.3f}".format(start),
|
||||
'end': "{:.3f}".format(end),
|
||||
'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):
|
||||
"""
|
||||
@@ -134,7 +176,8 @@ class ServeClientBase(object):
|
||||
|
||||
This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
|
||||
of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
|
||||
to prevent excessive memory usage.
|
||||
to prevent excessive memory usage. When the first frame arrives, it also wakes the transcription
|
||||
thread so processing can begin.
|
||||
|
||||
If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
|
||||
of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided
|
||||
@@ -144,20 +187,20 @@ class ServeClientBase(object):
|
||||
frame_np (numpy.ndarray): The audio frame data as a NumPy array.
|
||||
|
||||
"""
|
||||
self.lock.acquire()
|
||||
if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
|
||||
self.frames_offset += 30.0
|
||||
self.frames_np = self.frames_np[int(30*self.RATE):]
|
||||
# check timestamp offset(should be >= self.frame_offset)
|
||||
# this basically means that there is no speech as timestamp offset hasnt updated
|
||||
# and is less than frame_offset
|
||||
if self.timestamp_offset < self.frames_offset:
|
||||
self.timestamp_offset = self.frames_offset
|
||||
if self.frames_np is None:
|
||||
self.frames_np = frame_np.copy()
|
||||
else:
|
||||
self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
|
||||
self.lock.release()
|
||||
with self.lock:
|
||||
if self.frames_np is not None and self.frames_np.shape[0] > self.MAX_BUFFER_DURATION_S*self.RATE:
|
||||
self.frames_offset += float(self.BUFFER_TRIM_DURATION_S)
|
||||
self.frames_np = self.frames_np[int(self.BUFFER_TRIM_DURATION_S*self.RATE):]
|
||||
# check timestamp offset(should be >= self.frame_offset)
|
||||
# this basically means that there is no speech as timestamp offset hasnt updated
|
||||
# and is less than frame_offset
|
||||
if self.timestamp_offset < self.frames_offset:
|
||||
self.timestamp_offset = self.frames_offset
|
||||
if self.frames_np is None:
|
||||
self.frames_np = frame_np.copy()
|
||||
else:
|
||||
self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
|
||||
self.frames_ready.set()
|
||||
|
||||
def clip_audio_if_no_valid_segment(self):
|
||||
"""
|
||||
@@ -166,9 +209,9 @@ class ServeClientBase(object):
|
||||
no valid segment for the last 30 seconds from whisper
|
||||
"""
|
||||
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
|
||||
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):
|
||||
"""
|
||||
@@ -234,9 +277,23 @@ class ServeClientBase(object):
|
||||
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.
|
||||
|
||||
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:
|
||||
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:
|
||||
self.websocket.send(
|
||||
json.dumps({
|
||||
@@ -244,6 +301,8 @@ class ServeClientBase(object):
|
||||
"segments": segments,
|
||||
})
|
||||
)
|
||||
for seg in segments:
|
||||
wl_metrics.track_segment_emitted(completed=seg.get("completed", False))
|
||||
except Exception as e:
|
||||
logging.error(f"[ERROR]: Sending data to client: {e}")
|
||||
|
||||
@@ -271,6 +330,7 @@ class ServeClientBase(object):
|
||||
"""
|
||||
logging.info("Cleaning up.")
|
||||
self.exit = True
|
||||
self.frames_ready.set()
|
||||
|
||||
def get_segment_no_speech_prob(self, segment):
|
||||
return getattr(segment, "no_speech_prob", 0)
|
||||
@@ -281,6 +341,45 @@ class ServeClientBase(object):
|
||||
def get_segment_end(self, segment):
|
||||
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):
|
||||
"""
|
||||
Processes the segments from Whisper and updates the transcript.
|
||||
@@ -310,7 +409,9 @@ class ServeClientBase(object):
|
||||
continue
|
||||
if self.get_segment_no_speech_prob(s) > self.no_speech_thresh:
|
||||
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)
|
||||
|
||||
if self.translation_queue:
|
||||
@@ -323,12 +424,14 @@ class ServeClientBase(object):
|
||||
# Process the last segment if its no_speech_prob is acceptable.
|
||||
if self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh:
|
||||
self.current_out += segments[-1].text
|
||||
words = self._extract_words(segments[-1], self.timestamp_offset)
|
||||
with self.lock:
|
||||
last_segment = self.format_segment(
|
||||
self.timestamp_offset + self.get_segment_start(segments[-1]),
|
||||
self.timestamp_offset + min(duration, self.get_segment_end(segments[-1])),
|
||||
self.current_out,
|
||||
completed=False
|
||||
completed=False,
|
||||
words=words
|
||||
)
|
||||
|
||||
# Handle repeated output logic.
|
||||
@@ -376,4 +479,12 @@ class ServeClientBase(object):
|
||||
with self.lock:
|
||||
self.timestamp_offset += offset
|
||||
|
||||
self._trim_transcript()
|
||||
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:]
|
||||
|
||||
@@ -34,6 +34,9 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
same_output_threshold=7,
|
||||
cache_path="~/.cache/whisper-live/",
|
||||
translation_queue=None,
|
||||
hotwords=None,
|
||||
diarization=None,
|
||||
word_timestamps=False,
|
||||
):
|
||||
"""
|
||||
Initialize a ServeClient instance.
|
||||
@@ -63,7 +66,9 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
no_speech_thresh,
|
||||
clip_audio,
|
||||
same_output_threshold,
|
||||
translation_queue
|
||||
translation_queue,
|
||||
diarization,
|
||||
word_timestamps,
|
||||
)
|
||||
self.cache_path = cache_path
|
||||
self.model_sizes = [
|
||||
@@ -78,6 +83,7 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
self.task = task
|
||||
self.initial_prompt = initial_prompt
|
||||
self.vad_parameters = vad_parameters or {"threshold": 0.5}
|
||||
self.hotwords = hotwords
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if device == "cuda":
|
||||
@@ -213,6 +219,8 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
initial_prompt=self.initial_prompt,
|
||||
use_vad=self.use_vad,
|
||||
vad_parameters=self.vad_parameters if self.use_vad else None,
|
||||
word_timestamps=self.word_timestamps,
|
||||
client_uid=self.client_uid,
|
||||
)
|
||||
ServeClientFasterWhisper.BATCH_WORKER.submit(request)
|
||||
request.future.wait(timeout=30)
|
||||
@@ -231,7 +239,9 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
language=self.language,
|
||||
task=self.task,
|
||||
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:
|
||||
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ class ServeClientOpenVINO(ServeClientBase):
|
||||
no_speech_thresh=0.45,
|
||||
clip_audio=False,
|
||||
same_output_threshold=10,
|
||||
diarization=None,
|
||||
):
|
||||
"""
|
||||
Initialize a ServeClient instance.
|
||||
@@ -56,6 +57,8 @@ class ServeClientOpenVINO(ServeClientBase):
|
||||
no_speech_thresh,
|
||||
clip_audio,
|
||||
same_output_threshold,
|
||||
None, # translation_queue — OpenVINO backend does not support translation
|
||||
diarization, # speaker diarization — passed through to base class
|
||||
)
|
||||
self.language = "en" if language is None else language
|
||||
if not self.language.startswith("<|"):
|
||||
@@ -96,6 +99,20 @@ class ServeClientOpenVINO(ServeClientBase):
|
||||
logging.info(f"Using OpenVINO device: {self.device}")
|
||||
logging.info(f"Running OpenVINO backend with language: {self.language} and task: {self.task}")
|
||||
|
||||
def get_segment_end(self, segment):
|
||||
"""
|
||||
Override base class implementation to handle OpenVINO's end timestamp sentinel value.
|
||||
|
||||
WhisperDecodedResultChunk.end_ts is -1.0 when the model did not predict an ending
|
||||
timestamp (e.g. audio cut off mid-word). A negative end_ts causes a negative array
|
||||
index in _identify_speaker(), producing an empty audio slice and silently disabling
|
||||
diarization. Fall back to start_ts + 1.0 second in that case.
|
||||
"""
|
||||
end = getattr(segment, "end_ts", -1.0)
|
||||
if end < 0:
|
||||
return getattr(segment, "start_ts", 0) + 1.0
|
||||
return end
|
||||
|
||||
def create_model(self, model_id):
|
||||
"""
|
||||
Instantiates a new model, sets it as the transcriber.
|
||||
|
||||
@@ -74,6 +74,8 @@ class BatchRequest:
|
||||
initial_prompt: Optional[str] = None
|
||||
use_vad: bool = True
|
||||
vad_parameters: Optional[Dict] = None
|
||||
word_timestamps: bool = False
|
||||
client_uid: Optional[str] = None
|
||||
# Signaling
|
||||
future: threading.Event = field(default_factory=threading.Event)
|
||||
# Results (filled by batch worker)
|
||||
@@ -307,36 +309,87 @@ class BatchInferenceWorker:
|
||||
tokenizers_list.append(tokenizer)
|
||||
prompts.append(prompt)
|
||||
|
||||
# Step 4: Batch GPU generate
|
||||
# Step 4: Batch GPU generate with per-item temperature fallback.
|
||||
# Mirrors faster_whisper.transcribe()'s fallback loop. Items that
|
||||
# pass quality thresholds at lower temperature keep their result;
|
||||
# only failed items are re-decoded at the next temperature.
|
||||
suppress_tokens = get_suppressed_tokens(tokenizers_list[0], [-1])
|
||||
|
||||
results = self.transcriber.model.generate(
|
||||
encoder_output,
|
||||
prompts,
|
||||
beam_size=5,
|
||||
patience=1,
|
||||
length_penalty=1,
|
||||
max_length=self.transcriber.max_length,
|
||||
suppress_blank=True,
|
||||
suppress_tokens=suppress_tokens,
|
||||
return_scores=True,
|
||||
return_no_speech_prob=True,
|
||||
sampling_temperature=0.0,
|
||||
repetition_penalty=1,
|
||||
no_repeat_ngram_size=0,
|
||||
)
|
||||
temperatures = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
|
||||
comp_thresh = 2.4
|
||||
logprob_thresh = -1.0
|
||||
no_speech_thresh = 0.6
|
||||
|
||||
n = len(preprocessed)
|
||||
final_results = [None] * n # tuples of (gen_result, avg_logprob, used_temp)
|
||||
pending_indices = list(range(n))
|
||||
|
||||
for temp in temperatures:
|
||||
if not pending_indices:
|
||||
break
|
||||
|
||||
if len(pending_indices) == n:
|
||||
sub_encoder = encoder_output
|
||||
else:
|
||||
# Re-encode features for just the pending items to get
|
||||
# an encoder_output of the right batch dimension.
|
||||
sub_feature_batch = np.stack(
|
||||
[preprocessed[i][1] for i in pending_indices]
|
||||
)
|
||||
sub_encoder = self.transcriber.encode(sub_feature_batch)
|
||||
sub_prompts = [prompts[i] for i in pending_indices]
|
||||
|
||||
gen_kwargs = dict(
|
||||
beam_size=5 if temp == 0.0 else 1,
|
||||
patience=1,
|
||||
length_penalty=1,
|
||||
max_length=self.transcriber.max_length,
|
||||
suppress_blank=True,
|
||||
suppress_tokens=suppress_tokens,
|
||||
return_scores=True,
|
||||
return_no_speech_prob=True,
|
||||
sampling_temperature=temp,
|
||||
repetition_penalty=1,
|
||||
no_repeat_ngram_size=0,
|
||||
)
|
||||
batch_results = self.transcriber.model.generate(
|
||||
sub_encoder, sub_prompts, **gen_kwargs
|
||||
)
|
||||
|
||||
next_pending = []
|
||||
for j, idx in enumerate(pending_indices):
|
||||
gen_result = batch_results[j]
|
||||
tokens = gen_result.sequences_ids[0]
|
||||
seq_len = len(tokens)
|
||||
cum_logprob = gen_result.scores[0] * seq_len
|
||||
avg_logprob = cum_logprob / (seq_len + 1) if seq_len > 0 else 0.0
|
||||
raw_text = tokenizers_list[idx].decode(tokens).strip()
|
||||
comp_ratio = get_compression_ratio(raw_text) if raw_text else 0.0
|
||||
|
||||
bad = (
|
||||
comp_ratio > comp_thresh
|
||||
or avg_logprob < logprob_thresh
|
||||
)
|
||||
# High no_speech + low logprob -> treat as silence, accept empty.
|
||||
is_silence = (
|
||||
gen_result.no_speech_prob > no_speech_thresh
|
||||
and avg_logprob < logprob_thresh
|
||||
)
|
||||
|
||||
if not bad or is_silence or temp == temperatures[-1]:
|
||||
final_results[idx] = (gen_result, avg_logprob, temp)
|
||||
else:
|
||||
next_pending.append(idx)
|
||||
|
||||
pending_indices = next_pending
|
||||
|
||||
# Step 5: Per-item segment parsing and result dispatch
|
||||
for i, (req, features, audio, duration, speech_chunks) in enumerate(preprocessed):
|
||||
try:
|
||||
tokenizer = tokenizers_list[i]
|
||||
gen_result = results[i]
|
||||
gen_result, avg_logprob, used_temp = final_results[i]
|
||||
|
||||
tokens = gen_result.sequences_ids[0]
|
||||
seq_len = len(tokens)
|
||||
cum_logprob = gen_result.scores[0] * seq_len
|
||||
avg_logprob = cum_logprob / (seq_len + 1) if seq_len > 0 else 0.0
|
||||
|
||||
segment_size = int(ceil(duration) * self.transcriber.frames_per_second)
|
||||
|
||||
subsegments, _, _ = self.transcriber._split_segments_by_timestamps(
|
||||
@@ -364,7 +417,7 @@ class BatchInferenceWorker:
|
||||
compression_ratio=get_compression_ratio(text),
|
||||
no_speech_prob=gen_result.no_speech_prob,
|
||||
words=None,
|
||||
temperature=0.0,
|
||||
temperature=used_temp,
|
||||
))
|
||||
|
||||
req.result = segments
|
||||
|
||||
+293
-11
@@ -11,6 +11,7 @@ import websocket
|
||||
import uuid
|
||||
import time
|
||||
import av
|
||||
from typing import Callable, Literal, Optional
|
||||
import whisper_live.utils as utils
|
||||
|
||||
|
||||
@@ -43,6 +44,14 @@ class Client:
|
||||
translation_srt_file_path="output_translated.srt",
|
||||
enable_timestamps=False,
|
||||
display_segments=4,
|
||||
hotwords=None,
|
||||
enable_diarization=False,
|
||||
max_speakers=10,
|
||||
word_timestamps=False,
|
||||
max_retries=0,
|
||||
retry_delay=5,
|
||||
initial_prompt=None,
|
||||
vad_parameters=None,
|
||||
):
|
||||
"""
|
||||
Initializes a Client instance for audio recording and streaming to a server.
|
||||
@@ -69,6 +78,8 @@ class Client:
|
||||
target_language (str, optional): Target language for translation. Defaults to 'fr'.
|
||||
translation_callback (callable, optional): A callback function to handle translation results. Default is None.
|
||||
translation_srt_file_path (str, optional): The file path to save the translated output SRT file. Default is "output_translated.srt".
|
||||
initial_prompt (str, optional): Optional text to provide context to the model (e.g. domain vocabulary or names). Default is None.
|
||||
vad_parameters (dict, optional): Optional voice-activity-detection parameters passed to the server backend. Default is None.
|
||||
"""
|
||||
self.recording = False
|
||||
self.task = "transcribe"
|
||||
@@ -97,25 +108,29 @@ class Client:
|
||||
self.translation_callback = translation_callback
|
||||
self.translation_srt_file_path = translation_srt_file_path
|
||||
self.last_translated_segment = None
|
||||
|
||||
self.initial_prompt = initial_prompt
|
||||
self.vad_parameters = vad_parameters
|
||||
|
||||
if translate:
|
||||
self.task = "translate"
|
||||
self.enable_timestamps = enable_timestamps
|
||||
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
|
||||
|
||||
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_url = f"{socket_protocol}://{host}:{port}"
|
||||
self.client_socket = websocket.WebSocketApp(
|
||||
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
|
||||
),
|
||||
)
|
||||
self.socket_url = f"{socket_protocol}://{host}:{port}"
|
||||
self._create_websocket()
|
||||
else:
|
||||
print("[ERROR]: No host or port specified.")
|
||||
return
|
||||
@@ -131,6 +146,18 @@ class Client:
|
||||
self.translated_transcript = []
|
||||
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):
|
||||
"""Handles server status messages."""
|
||||
status = message_data["status"]
|
||||
@@ -273,6 +300,15 @@ class Client:
|
||||
self.recording = 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):
|
||||
"""
|
||||
Callback function called when the WebSocket connection is successfully opened.
|
||||
@@ -299,6 +335,12 @@ class Client:
|
||||
"same_output_threshold": self.same_output_threshold,
|
||||
"enable_translation": self.enable_translation,
|
||||
"target_language": self.target_language,
|
||||
"hotwords": self.hotwords,
|
||||
"enable_diarization": self.enable_diarization,
|
||||
"max_speakers": self.max_speakers,
|
||||
"word_timestamps": self.word_timestamps,
|
||||
"initial_prompt": self.initial_prompt,
|
||||
"vad_parameters": self.vad_parameters,
|
||||
}
|
||||
)
|
||||
)
|
||||
@@ -820,7 +862,14 @@ class TranscriptionClient(TranscriptionTeeClient):
|
||||
translation_srt_file_path="./output_translated.srt",
|
||||
enable_timestamps=False,
|
||||
display_segments=4,
|
||||
hotwords=None,
|
||||
enable_diarization=False,
|
||||
max_speakers=10,
|
||||
word_timestamps=False,
|
||||
initial_prompt=None,
|
||||
vad_parameters=None,
|
||||
):
|
||||
|
||||
self.client = Client(
|
||||
host,
|
||||
port,
|
||||
@@ -842,6 +891,12 @@ class TranscriptionClient(TranscriptionTeeClient):
|
||||
translation_srt_file_path=translation_srt_file_path,
|
||||
enable_timestamps=enable_timestamps,
|
||||
display_segments=display_segments,
|
||||
hotwords=hotwords,
|
||||
enable_diarization=enable_diarization,
|
||||
max_speakers=max_speakers,
|
||||
word_timestamps=word_timestamps,
|
||||
initial_prompt=initial_prompt,
|
||||
vad_parameters=vad_parameters,
|
||||
)
|
||||
|
||||
if save_output_recording and not output_recording_filename.endswith(".wav"):
|
||||
@@ -857,3 +912,230 @@ class TranscriptionClient(TranscriptionTeeClient):
|
||||
output_recording_filename=output_recording_filename,
|
||||
mute_audio_playback=mute_audio_playback
|
||||
)
|
||||
|
||||
|
||||
PcmFormat = Literal["float32", "int16"]
|
||||
|
||||
|
||||
class _HookedClient(Client):
|
||||
"""Client subclass that exposes lifecycle callbacks not available on the base class."""
|
||||
|
||||
def __init__(self, *args, on_session_started=None, on_error_hook=None, on_close_hook=None, **kwargs):
|
||||
self._on_session_started = on_session_started
|
||||
self._on_error_hook = on_error_hook
|
||||
self._on_close_hook = on_close_hook
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def on_message(self, ws, message):
|
||||
was_recording = self.recording
|
||||
super().on_message(ws, message)
|
||||
if not was_recording and self.recording and self._on_session_started:
|
||||
self._on_session_started()
|
||||
|
||||
def on_error(self, ws, error):
|
||||
# websocket-client surfaces the server's CLOSE control frame (opcode 8)
|
||||
# through on_error during shutdown; a normal close is not an error.
|
||||
if getattr(error, "opcode", None) == 8:
|
||||
return
|
||||
if self._on_error_hook:
|
||||
self._on_error_hook(error)
|
||||
super().on_error(ws, error)
|
||||
|
||||
def on_close(self, ws, close_status_code, close_msg):
|
||||
if self._on_close_hook:
|
||||
self._on_close_hook()
|
||||
super().on_close(ws, close_status_code, close_msg)
|
||||
|
||||
|
||||
class StreamingTranscriptionClient:
|
||||
"""Feed raw PCM audio in chunks; receive partial and committed transcripts via callbacks.
|
||||
|
||||
Args:
|
||||
host: WhisperLive server hostname.
|
||||
port: WhisperLive server port.
|
||||
lang: Language code (e.g. ``"en"``). ``None`` enables auto-detection.
|
||||
model: Whisper model size (``"tiny"``, ``"base"``, ``"small"``, ``"medium"``, ``"large"``).
|
||||
use_vad: Enable server-side voice activity detection.
|
||||
use_wss: Use ``wss://`` instead of ``ws://``.
|
||||
send_last_n_segments: How many recent segments the server echoes per update.
|
||||
no_speech_thresh: Segments with no-speech probability above this are discarded.
|
||||
clip_audio: Drop audio with no valid segments.
|
||||
same_output_threshold: Repeated identical outputs before a segment is committed.
|
||||
enable_translation: Enable post-transcription translation.
|
||||
target_language: Target language for translation (e.g. ``"fr"``).
|
||||
ready_timeout: Seconds to wait for ``SERVER_READY`` before raising ``TimeoutError``.
|
||||
on_session_started: Called once when the server is ready to receive audio.
|
||||
on_partial_transcript: Called on each in-progress segment update with ``(text, segments)``.
|
||||
on_committed_transcript: Called for each finalized segment with ``(text, segments)``.
|
||||
on_translation: Called for each translated segment with ``(text, segments)``.
|
||||
on_error: Called on WebSocket errors with the exception.
|
||||
on_close: Called when the connection closes.
|
||||
|
||||
Example::
|
||||
|
||||
client = StreamingTranscriptionClient(
|
||||
"localhost", 9090,
|
||||
lang="en",
|
||||
on_partial_transcript=lambda text, _: print(f"… {text}", end="\\r"),
|
||||
on_committed_transcript=lambda text, _: print(f"✓ {text}"),
|
||||
)
|
||||
with client:
|
||||
for chunk in my_audio_source:
|
||||
client.send(chunk, pcm_format="int16")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
*,
|
||||
lang: Optional[str] = None,
|
||||
model: str = "small",
|
||||
use_vad: bool = True,
|
||||
use_wss: bool = False,
|
||||
send_last_n_segments: int = 10,
|
||||
no_speech_thresh: float = 0.45,
|
||||
clip_audio: bool = False,
|
||||
same_output_threshold: int = 10,
|
||||
enable_translation: bool = False,
|
||||
target_language: str = "fr",
|
||||
ready_timeout: float = 30.0,
|
||||
on_session_started: Optional[Callable[[], None]] = None,
|
||||
on_partial_transcript: Optional[Callable[[str, list], None]] = None,
|
||||
on_committed_transcript: Optional[Callable[[str, list], None]] = None,
|
||||
on_translation: Optional[Callable[[str, list], None]] = None,
|
||||
on_error: Optional[Callable[[Exception], None]] = None,
|
||||
on_close: Optional[Callable[[], None]] = None,
|
||||
):
|
||||
self._on_partial_transcript = on_partial_transcript
|
||||
self._on_committed_transcript = on_committed_transcript
|
||||
self._ready_timeout = ready_timeout
|
||||
self._closed = False
|
||||
self._transcript = []
|
||||
self._committed_keys = set()
|
||||
|
||||
self._client = _HookedClient(
|
||||
host=host,
|
||||
port=port,
|
||||
lang=lang,
|
||||
model=model,
|
||||
use_vad=use_vad,
|
||||
use_wss=use_wss,
|
||||
log_transcription=False,
|
||||
send_last_n_segments=send_last_n_segments,
|
||||
no_speech_thresh=no_speech_thresh,
|
||||
clip_audio=clip_audio,
|
||||
same_output_threshold=same_output_threshold,
|
||||
enable_translation=enable_translation,
|
||||
target_language=target_language,
|
||||
transcription_callback=self._dispatch_transcript,
|
||||
translation_callback=on_translation,
|
||||
on_session_started=on_session_started,
|
||||
on_error_hook=on_error,
|
||||
on_close_hook=on_close,
|
||||
)
|
||||
|
||||
def _dispatch_transcript(self, text: str, segments: list) -> None:
|
||||
for seg in segments:
|
||||
if not seg.get("completed", False):
|
||||
continue
|
||||
key = (seg.get("start"), seg.get("end"), seg.get("text"))
|
||||
if key in self._committed_keys:
|
||||
continue
|
||||
self._committed_keys.add(key)
|
||||
self._transcript.append(seg)
|
||||
if self._on_committed_transcript:
|
||||
self._on_committed_transcript(seg["text"].strip(), [seg])
|
||||
|
||||
last = segments[-1] if segments else None
|
||||
if last and not last.get("completed", False) and self._on_partial_transcript:
|
||||
self._on_partial_transcript(last["text"].strip(), [last])
|
||||
|
||||
def connect(self) -> "StreamingTranscriptionClient":
|
||||
"""Block until the server is ready. Returns self for use as a context manager."""
|
||||
deadline = time.time() + self._ready_timeout
|
||||
while not self._client.recording:
|
||||
if self._client.server_error:
|
||||
raise RuntimeError(getattr(self._client, "error_message", "Server reported an error."))
|
||||
if self._client.waiting:
|
||||
raise RuntimeError("Server is full.")
|
||||
if time.time() > deadline:
|
||||
raise TimeoutError("Timed out waiting for server ready.")
|
||||
time.sleep(0.05)
|
||||
return self
|
||||
|
||||
def send(self, audio_bytes: bytes, pcm_format: PcmFormat = "int16") -> None:
|
||||
"""Send one PCM chunk. Any chunk size is fine; must be mono 16 kHz.
|
||||
|
||||
Args:
|
||||
audio_bytes: Raw PCM payload.
|
||||
pcm_format: ``"int16"`` is normalized to float32; ``"float32"`` passes through.
|
||||
"""
|
||||
if self._closed:
|
||||
raise RuntimeError("Client is already closed.")
|
||||
if not audio_bytes:
|
||||
return
|
||||
if pcm_format == "float32":
|
||||
payload = audio_bytes
|
||||
elif pcm_format == "int16":
|
||||
samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
|
||||
payload = samples.tobytes()
|
||||
else:
|
||||
raise ValueError(f"Unsupported pcm_format: {pcm_format!r}")
|
||||
self._client.send_packet_to_server(payload)
|
||||
|
||||
def send_array(self, samples: np.ndarray) -> None:
|
||||
"""Send a numpy array (any numeric dtype, mono, 16 kHz).
|
||||
|
||||
Args:
|
||||
samples: 1-D numpy array of audio samples.
|
||||
"""
|
||||
if samples.ndim != 1:
|
||||
raise ValueError("Expected mono (1-D) array.")
|
||||
if np.issubdtype(samples.dtype, np.integer):
|
||||
info = np.iinfo(samples.dtype)
|
||||
samples = samples.astype(np.float32) / max(abs(info.min), info.max)
|
||||
elif samples.dtype != np.float32:
|
||||
samples = samples.astype(np.float32)
|
||||
self._client.send_packet_to_server(samples.tobytes())
|
||||
|
||||
@property
|
||||
def transcript(self) -> list:
|
||||
"""All committed segments received so far."""
|
||||
return self._transcript
|
||||
|
||||
@property
|
||||
def last_partial(self) -> Optional[dict]:
|
||||
"""The most recent in-progress segment, or ``None`` if none pending."""
|
||||
return self._client.last_segment
|
||||
|
||||
# Alias for ``last_partial``; kept for readability at call sites.
|
||||
last_segment = last_partial
|
||||
|
||||
def close(self, timeout: float = 15.0) -> None:
|
||||
"""Signal end-of-stream, wait for the server to finish, then close.
|
||||
|
||||
After ``END_OF_AUDIO`` the server transcribes any buffered audio, sends
|
||||
the final committed segment, and closes the connection. Waiting for that
|
||||
server-initiated close keeps the last segment from being dropped.
|
||||
|
||||
Args:
|
||||
timeout: Maximum seconds to wait for the server to close before
|
||||
forcing the connection shut.
|
||||
"""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
try:
|
||||
self._client.send_packet_to_server(Client.END_OF_AUDIO.encode("utf-8"))
|
||||
deadline = time.time() + timeout
|
||||
while self._client.recording and time.time() < deadline:
|
||||
time.sleep(0.05)
|
||||
finally:
|
||||
self._client.close_websocket()
|
||||
|
||||
def __enter__(self) -> "StreamingTranscriptionClient":
|
||||
return self.connect()
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None:
|
||||
self.close()
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
def load_audio(file_path, sample_rate=16000):
|
||||
"""Load an audio file as mono float32 PCM at the requested sample rate."""
|
||||
import av
|
||||
|
||||
container = av.open(file_path)
|
||||
resampler = av.AudioResampler(format="flt", layout="mono", rate=sample_rate)
|
||||
chunks = []
|
||||
|
||||
try:
|
||||
for frame in container.decode(audio=0):
|
||||
for resampled_frame in resampler.resample(frame):
|
||||
chunks.append(
|
||||
resampled_frame.to_ndarray().reshape(-1).astype(np.float32)
|
||||
)
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
if not chunks:
|
||||
return np.array([], dtype=np.float32)
|
||||
return np.concatenate(chunks)
|
||||
|
||||
|
||||
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,
|
||||
speaker_names=None,
|
||||
):
|
||||
self.similarity_threshold = similarity_threshold
|
||||
self.max_speakers = max_speakers
|
||||
self.speaker_names = list(speaker_names or [])
|
||||
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 _next_speaker_id(self):
|
||||
if self._speaker_count < len(self.speaker_names):
|
||||
return self.speaker_names[self._speaker_count]
|
||||
return f"SPEAKER_{self._speaker_count:02d}"
|
||||
|
||||
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 = self._next_speaker_id()
|
||||
self._speaker_count += 1
|
||||
self.speakers[speaker_id] = embedding
|
||||
return speaker_id
|
||||
|
||||
def enroll_speaker(self, speaker_name, audio_np, sample_rate=16000):
|
||||
"""Enroll a known speaker from reference audio."""
|
||||
embedding = self._compute_embedding(audio_np, sample_rate)
|
||||
if embedding is None:
|
||||
return False
|
||||
self.speakers[speaker_name] = embedding
|
||||
return True
|
||||
|
||||
def reset(self):
|
||||
"""Reset all speaker state."""
|
||||
self.speakers.clear()
|
||||
self._speaker_count = 0
|
||||
@@ -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()
|
||||
+316
-37
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
import collections
|
||||
import queue
|
||||
import json
|
||||
import functools
|
||||
@@ -8,16 +9,18 @@ import logging
|
||||
import shutil
|
||||
import tempfile
|
||||
from typing import Optional, List
|
||||
from fastapi import FastAPI, UploadFile, Form
|
||||
from fastapi import FastAPI, UploadFile, Form, Request, File
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from starlette.responses import PlainTextResponse, JSONResponse
|
||||
from fastapi.responses import JSONResponse
|
||||
from starlette.responses import PlainTextResponse, StreamingResponse
|
||||
import uvicorn
|
||||
from faster_whisper import WhisperModel
|
||||
import torch
|
||||
|
||||
from enum import Enum
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
from whisper_live import metrics as wl_metrics
|
||||
from websockets.sync.server import serve
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from whisper_live.vad import VoiceActivityDetector
|
||||
@@ -39,6 +42,7 @@ class ClientManager:
|
||||
self.start_times = {}
|
||||
self.max_clients = max_clients
|
||||
self.max_connection_time = max_connection_time
|
||||
self.lock = threading.Lock()
|
||||
|
||||
def add_client(self, websocket, client):
|
||||
"""
|
||||
@@ -48,8 +52,9 @@ class ClientManager:
|
||||
websocket: The websocket associated with the client to add.
|
||||
client: The client object to be added and tracked.
|
||||
"""
|
||||
self.clients[websocket] = client
|
||||
self.start_times[websocket] = time.time()
|
||||
with self.lock:
|
||||
self.clients[websocket] = client
|
||||
self.start_times[websocket] = time.time()
|
||||
|
||||
def get_client(self, websocket):
|
||||
"""
|
||||
@@ -61,9 +66,10 @@ class ClientManager:
|
||||
Returns:
|
||||
The client object if found, False otherwise.
|
||||
"""
|
||||
if websocket in self.clients:
|
||||
return self.clients[websocket]
|
||||
return False
|
||||
with self.lock:
|
||||
if websocket in self.clients:
|
||||
return self.clients[websocket]
|
||||
return False
|
||||
|
||||
def remove_client(self, websocket):
|
||||
"""
|
||||
@@ -73,10 +79,11 @@ class ClientManager:
|
||||
Args:
|
||||
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:
|
||||
client.cleanup()
|
||||
self.start_times.pop(websocket, None)
|
||||
|
||||
def get_wait_time(self):
|
||||
"""
|
||||
@@ -85,11 +92,12 @@ class ClientManager:
|
||||
Returns:
|
||||
The estimated wait time in minutes for new clients to connect. Returns 0 if there are available slots.
|
||||
"""
|
||||
wait_time = None
|
||||
for start_time in self.start_times.values():
|
||||
current_client_time_remaining = self.max_connection_time - (time.time() - start_time)
|
||||
if wait_time is None or current_client_time_remaining < wait_time:
|
||||
wait_time = current_client_time_remaining
|
||||
with self.lock:
|
||||
wait_time = None
|
||||
for start_time in self.start_times.values():
|
||||
current_client_time_remaining = self.max_connection_time - (time.time() - start_time)
|
||||
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
|
||||
|
||||
def is_server_full(self, websocket, options):
|
||||
@@ -103,12 +111,18 @@ class ClientManager:
|
||||
Returns:
|
||||
True if the server is full, False otherwise.
|
||||
"""
|
||||
if len(self.clients) >= self.max_clients:
|
||||
wait_time = self.get_wait_time()
|
||||
response = {"uid": options["uid"], "status": "WAIT", "message": wait_time}
|
||||
websocket.send(json.dumps(response))
|
||||
return True
|
||||
return False
|
||||
with self.lock:
|
||||
if len(self.clients) >= self.max_clients:
|
||||
wait_time = None
|
||||
for start_time in self.start_times.values():
|
||||
remaining = self.max_connection_time - (time.time() - start_time)
|
||||
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):
|
||||
"""
|
||||
@@ -120,10 +134,12 @@ class ClientManager:
|
||||
Returns:
|
||||
True if the client's connection time has exceeded the maximum limit, False otherwise.
|
||||
"""
|
||||
elapsed_time = time.time() - self.start_times[websocket]
|
||||
if elapsed_time >= self.max_connection_time:
|
||||
self.clients[websocket].disconnect()
|
||||
logging.warning(f"Client with uid '{self.clients[websocket].client_uid}' disconnected due to overtime.")
|
||||
with self.lock:
|
||||
elapsed_time = time.time() - self.start_times[websocket]
|
||||
client = self.clients.get(websocket)
|
||||
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 False
|
||||
|
||||
@@ -160,6 +176,9 @@ class TranscriptionServer:
|
||||
self.use_vad = True
|
||||
self.single_model = False
|
||||
self.batch_config = None
|
||||
self.raw_pcm_input = False
|
||||
self.audio_formats = {}
|
||||
self.segment_post_processor = None
|
||||
|
||||
def initialize_client(
|
||||
self, websocket, options, faster_whisper_custom_model_path,
|
||||
@@ -177,7 +196,7 @@ class TranscriptionServer:
|
||||
|
||||
if enable_translation:
|
||||
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
|
||||
translation_client = ServeClientTranslation(
|
||||
client_uid=options["uid"],
|
||||
@@ -239,6 +258,7 @@ class TranscriptionServer:
|
||||
no_speech_thresh=options.get("no_speech_thresh", 0.45),
|
||||
clip_audio=options.get("clip_audio", False),
|
||||
same_output_threshold=options.get("same_output_threshold", 10),
|
||||
diarization=self._create_diarizer(options),
|
||||
)
|
||||
logging.info("Running OpenVINO backend.")
|
||||
except Exception as e:
|
||||
@@ -274,7 +294,10 @@ class TranscriptionServer:
|
||||
clip_audio=options.get("clip_audio", False),
|
||||
same_output_threshold=options.get("same_output_threshold", 10),
|
||||
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.")
|
||||
@@ -297,12 +320,35 @@ class TranscriptionServer:
|
||||
if client is None:
|
||||
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:
|
||||
client.translation_client = translation_client
|
||||
client.translation_thread = translation_thread
|
||||
|
||||
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):
|
||||
"""
|
||||
Receives audio buffer from websocket and creates a numpy array out of it.
|
||||
@@ -316,6 +362,13 @@ class TranscriptionServer:
|
||||
frame_data = websocket.recv()
|
||||
if frame_data == b"END_OF_AUDIO":
|
||||
return False
|
||||
audio_format = self.audio_formats.get(websocket)
|
||||
if audio_format == "uint8":
|
||||
audio_np = np.frombuffer(frame_data, dtype=np.uint8)
|
||||
return (audio_np.astype(np.float32) - 128.0) / 128.0
|
||||
if self.raw_pcm_input or audio_format == "int16":
|
||||
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)
|
||||
|
||||
def handle_new_connection(self, websocket, faster_whisper_custom_model_path,
|
||||
@@ -327,13 +380,19 @@ class TranscriptionServer:
|
||||
|
||||
self.use_vad = options.get('use_vad')
|
||||
if self.client_manager.is_server_full(websocket, options):
|
||||
wl_metrics.track_connection_rejected(reason="full")
|
||||
websocket.close()
|
||||
return False # Indicates that the connection should not continue
|
||||
audio_format = options.get("audio_format", "float32")
|
||||
if audio_format not in {"float32", "int16", "uint8"}:
|
||||
raise ValueError(f"Unsupported audio_format: {audio_format}")
|
||||
self.audio_formats[websocket] = audio_format
|
||||
|
||||
if self.backend.is_tensorrt():
|
||||
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
|
||||
self.initialize_client(websocket, options, faster_whisper_custom_model_path,
|
||||
whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session)
|
||||
wl_metrics.track_connection_opened()
|
||||
return True
|
||||
except json.JSONDecodeError:
|
||||
logging.error("Failed to decode JSON from client")
|
||||
@@ -412,8 +471,119 @@ class TranscriptionServer:
|
||||
if self.client_manager.get_client(websocket):
|
||||
self.cleanup(websocket)
|
||||
websocket.close()
|
||||
wl_metrics.track_connection_closed()
|
||||
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")
|
||||
|
||||
@staticmethod
|
||||
def _normalize_form_list(values):
|
||||
"""Normalize repeated or comma-separated multipart form fields."""
|
||||
if not values:
|
||||
return []
|
||||
normalized = []
|
||||
for value in values:
|
||||
if isinstance(value, str):
|
||||
normalized.extend(item.strip() for item in value.split(",") if item.strip())
|
||||
return normalized
|
||||
|
||||
async def _create_rest_diarizer(self, known_speaker_names, known_speaker_references):
|
||||
"""Create a diarizer from OpenAI-compatible known speaker fields."""
|
||||
speaker_names = self._normalize_form_list(known_speaker_names)
|
||||
speaker_references = known_speaker_references or []
|
||||
|
||||
if speaker_references and not speaker_names:
|
||||
raise ValueError("known_speaker_references requires matching known_speaker_names")
|
||||
if speaker_names and speaker_references and len(speaker_names) != len(speaker_references):
|
||||
raise ValueError("known_speaker_names and known_speaker_references must have the same length")
|
||||
if not speaker_names and not speaker_references:
|
||||
return None
|
||||
|
||||
from whisper_live.diarization import SpeakerDiarizer, load_audio
|
||||
|
||||
diarizer = SpeakerDiarizer(
|
||||
max_speakers=max(10, len(speaker_names)),
|
||||
speaker_names=speaker_names,
|
||||
)
|
||||
|
||||
for speaker_name, reference in zip(speaker_names, speaker_references):
|
||||
suffix = os.path.splitext(reference.filename or "")[1] or ".wav"
|
||||
reference_path = None
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(await reference.read())
|
||||
reference_path = tmp.name
|
||||
audio_np = load_audio(reference_path)
|
||||
if not diarizer.enroll_speaker(speaker_name, audio_np):
|
||||
raise ValueError(f"known_speaker_references for '{speaker_name}' is too short")
|
||||
finally:
|
||||
if reference_path and os.path.exists(reference_path):
|
||||
os.unlink(reference_path)
|
||||
|
||||
return diarizer
|
||||
|
||||
@staticmethod
|
||||
def _speaker_labels_for_segments(segments, audio_np, diarizer, sample_rate=16000):
|
||||
if diarizer is None or audio_np is None:
|
||||
return {}
|
||||
labels = {}
|
||||
for index, segment in enumerate(segments):
|
||||
start = max(0, int(segment.start * sample_rate))
|
||||
end = min(len(audio_np), int(segment.end * sample_rate))
|
||||
if end <= start:
|
||||
continue
|
||||
speaker = diarizer.identify_speaker(audio_np[start:end], sample_rate)
|
||||
if speaker:
|
||||
labels[index] = speaker
|
||||
return labels
|
||||
|
||||
def run(self,
|
||||
host,
|
||||
port=9090,
|
||||
@@ -431,7 +601,12 @@ class TranscriptionServer:
|
||||
cors_origins: Optional[str] = None,
|
||||
batch_enabled=False,
|
||||
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.
|
||||
|
||||
@@ -447,8 +622,25 @@ class TranscriptionServer:
|
||||
batch_window_ms (int): Maximum time in milliseconds to wait for
|
||||
the batch to fill after the first request arrives. Defaults
|
||||
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.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)
|
||||
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:
|
||||
@@ -472,11 +664,18 @@ class TranscriptionServer:
|
||||
logging.info("Custom model option was provided. Switching to single model mode.")
|
||||
self.single_model = True
|
||||
# TODO: load model initially
|
||||
elif batch_enabled:
|
||||
logging.info("Batch inference enabled. Switching to single model mode for stock model.")
|
||||
self.single_model = True
|
||||
else:
|
||||
logging.info("Single model mode currently only works with custom models.")
|
||||
if not BackendType.is_valid(backend):
|
||||
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)
|
||||
if enable_rest:
|
||||
app = FastAPI(title="WhisperLive OpenAI-Compatible API")
|
||||
@@ -489,6 +688,34 @@ class TranscriptionServer:
|
||||
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")
|
||||
async def transcribe(
|
||||
@@ -503,22 +730,35 @@ class TranscriptionServer:
|
||||
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)
|
||||
known_speaker_references: Optional[List[UploadFile]] = File(default=None),
|
||||
stream: bool = Form(default=False),
|
||||
hotwords: Optional[str] = Form(default=None),
|
||||
):
|
||||
if stream:
|
||||
return JSONResponse({"error": "Streaming not supported in this backend."}, status_code=400)
|
||||
if chunking_strategy or known_speaker_names or known_speaker_references:
|
||||
logging.warning("Diarization/chunking params ignored; not supported.")
|
||||
return self._stream_transcription(
|
||||
file, language, prompt, temperature,
|
||||
timestamp_granularities,
|
||||
faster_whisper_custom_model_path,
|
||||
)
|
||||
|
||||
ignored_params = []
|
||||
if chunking_strategy:
|
||||
ignored_params.append(f"chunking_strategy='{chunking_strategy}'")
|
||||
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"]
|
||||
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)
|
||||
|
||||
if model != "whisper-1":
|
||||
logging.warning(f"Model '{model}' requested; using 'small' as fallback.")
|
||||
model_name = faster_whisper_custom_model_path or "small"
|
||||
|
||||
tmp_path = None
|
||||
try:
|
||||
suffix = os.path.splitext(file.filename)[1] or ".wav"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
@@ -535,15 +775,18 @@ class TranscriptionServer:
|
||||
initial_prompt=prompt,
|
||||
temperature=temperature,
|
||||
vad_filter=False,
|
||||
word_timestamps=(timestamp_granularities and "word" in timestamp_granularities)
|
||||
word_timestamps=(timestamp_granularities and "word" in timestamp_granularities),
|
||||
hotwords=hotwords,
|
||||
)
|
||||
segments = list(segments)
|
||||
|
||||
text = " ".join([s.text.strip() for s in segments])
|
||||
os.unlink(tmp_path)
|
||||
|
||||
if response_format == "text":
|
||||
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
|
||||
return PlainTextResponse(text)
|
||||
elif response_format == "json":
|
||||
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
|
||||
return {"text": text}
|
||||
elif response_format == "verbose_json":
|
||||
verbose = {
|
||||
@@ -553,7 +796,17 @@ class TranscriptionServer:
|
||||
"text": text,
|
||||
"segments": []
|
||||
}
|
||||
for seg in segments:
|
||||
speaker_labels = {}
|
||||
try:
|
||||
rest_diarizer = await self._create_rest_diarizer(known_speaker_names, known_speaker_references)
|
||||
except ValueError as e:
|
||||
wl_metrics.track_rest_request(endpoint="transcriptions", status=400)
|
||||
return JSONResponse({"error": str(e)}, status_code=400)
|
||||
if rest_diarizer is not None:
|
||||
from whisper_live.diarization import load_audio
|
||||
audio_np = load_audio(tmp_path)
|
||||
speaker_labels = self._speaker_labels_for_segments(segments, audio_np, rest_diarizer)
|
||||
for index, seg in enumerate(segments):
|
||||
seg_dict = {
|
||||
"id": seg.id,
|
||||
"seek": seg.seek,
|
||||
@@ -566,9 +819,12 @@ class TranscriptionServer:
|
||||
"compression_ratio": seg.compression_ratio,
|
||||
"no_speech_prob": seg.no_speech_prob
|
||||
}
|
||||
if index in speaker_labels:
|
||||
seg_dict["speaker"] = speaker_labels[index]
|
||||
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]
|
||||
verbose["segments"].append(seg_dict)
|
||||
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
|
||||
return verbose
|
||||
elif response_format in ["srt", "vtt"]:
|
||||
output = []
|
||||
@@ -579,9 +835,15 @@ class TranscriptionServer:
|
||||
output.append(f"{i}\n{start.replace('.', ',')} --> {end.replace('.', ',')}\n{seg.text.strip()}\n")
|
||||
else: # vtt
|
||||
output.append(f"{start} --> {end}\n{seg.text.strip()}\n")
|
||||
wl_metrics.track_rest_request(endpoint="transcriptions", status=200)
|
||||
return PlainTextResponse("\n".join(output))
|
||||
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)
|
||||
finally:
|
||||
if tmp_path and os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
|
||||
threading.Thread(
|
||||
target=uvicorn.run,
|
||||
@@ -592,6 +854,21 @@ class TranscriptionServer:
|
||||
logging.info(f"✅ OpenAI-Compatible API started on http://0.0.0.0:{rest_port}")
|
||||
|
||||
# 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(
|
||||
functools.partial(
|
||||
self.recv_audio,
|
||||
@@ -602,7 +879,8 @@ class TranscriptionServer:
|
||||
trt_py_session=trt_py_session,
|
||||
),
|
||||
host,
|
||||
port
|
||||
port,
|
||||
**extra_ws_kwargs,
|
||||
) as server:
|
||||
server.serve_forever()
|
||||
|
||||
@@ -652,3 +930,4 @@ class TranscriptionServer:
|
||||
if hasattr(client, 'translation_thread') and client.translation_thread:
|
||||
client.translation_thread.join(timeout=2.0)
|
||||
self.client_manager.remove_client(websocket)
|
||||
self.audio_formats.pop(websocket, None)
|
||||
|
||||
+19
-7
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import shutil
|
||||
import textwrap
|
||||
import scipy
|
||||
import numpy as np
|
||||
@@ -8,19 +9,30 @@ from pathlib import Path
|
||||
|
||||
def clear_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):
|
||||
"""Prints formatted transcript text."""
|
||||
"""Prints formatted transcript text in a subtitle-like block."""
|
||||
terminal_width = shutil.get_terminal_size((80, 20)).columns
|
||||
wrap_width = max(10, min(80, terminal_width - 8))
|
||||
|
||||
if timestamps:
|
||||
lines = []
|
||||
for t in text:
|
||||
print(f'[{t["start"]} -> {t["end"]}] {t["text"]}')
|
||||
prefix = f'[{t["start"]} -> {t["end"]}] '
|
||||
wrapper = textwrap.TextWrapper(
|
||||
width=wrap_width,
|
||||
subsequent_indent=" " * len(prefix),
|
||||
)
|
||||
lines.extend(wrapper.wrap(f'{prefix}{t["text"]}'))
|
||||
else:
|
||||
wrapper = textwrap.TextWrapper(width=60)
|
||||
text=" ".join(text) if translated else "".join(text)
|
||||
for line in wrapper.wrap(text=text):
|
||||
print(line)
|
||||
wrapper = textwrap.TextWrapper(width=wrap_width)
|
||||
transcript = " ".join(text) if translated else "".join(text)
|
||||
lines = wrapper.wrap(text=transcript)
|
||||
|
||||
for line in lines[-3:]:
|
||||
print(line.center(terminal_width))
|
||||
|
||||
|
||||
def format_time(s):
|
||||
|
||||
Reference in New Issue
Block a user