Merge pull request #284 from makaveli10/expose_client_manager_args

Expose client manager args.
This commit is contained in:
Marcus Edel
2024-10-31 15:40:11 -04:00
committed by GitHub
6 changed files with 45 additions and 21 deletions
+6 -1
View File
@@ -77,6 +77,9 @@ If you don't want this, set `--no_single_model`.
- `use_vad`: Whether to use `Voice Activity Detection` on the server.
- `save_output_recording`: Set to True to save the microphone input as a `.wav` file during live transcription. This option is helpful for recording sessions for later playback or analysis. Defaults to `False`.
- `output_recording_filename`: Specifies the `.wav` file path where the microphone input will be saved if `save_output_recording` is set to `True`.
- `max_clients`: Specifies the maximum number of clients the server should allow. Defaults to 4.
- `max_connection_time`: Maximum connection time for each client in seconds. Defaults to 600.
```python
from whisper_live.client import TranscriptionClient
client = TranscriptionClient(
@@ -87,7 +90,9 @@ client = TranscriptionClient(
model="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
output_recording_filename="./output_recording.wav", # Only used for microphone input
max_clients=4,
max_connection_time=600
)
```
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.
+1 -1
View File
@@ -10,4 +10,4 @@ jiwer
evaluate
numpy<2
tiktoken==0.3.3
openai-whisper
openai-whisper==20231117
+3 -1
View File
@@ -48,7 +48,9 @@ class TestClientCallbacks(BaseTestCase):
"language": self.client.language,
"task": self.client.task,
"model": self.client.model,
"use_vad": True
"use_vad": True,
"max_clients": 4,
"max_connection_time": 600,
})
self.client.on_open(self.mock_ws_app)
self.mock_ws_app.send.assert_called_with(expected_message)
+12 -14
View File
@@ -5,10 +5,10 @@ import unittest
from unittest import mock
import numpy as np
import evaluate
import jiwer
from websockets.exceptions import ConnectionClosed
from whisper_live.server import TranscriptionServer
from whisper_live.server import TranscriptionServer, BackendType, ClientManager
from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
from whisper.normalizers import EnglishTextNormalizer
@@ -16,6 +16,7 @@ from whisper.normalizers import EnglishTextNormalizer
class TestTranscriptionServerInitialization(unittest.TestCase):
def test_initialization(self):
server = TranscriptionServer()
server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
self.assertEqual(server.client_manager.max_clients, 4)
self.assertEqual(server.client_manager.max_connection_time, 600)
self.assertDictEqual(server.client_manager.clients, {})
@@ -25,6 +26,7 @@ class TestTranscriptionServerInitialization(unittest.TestCase):
class TestGetWaitTime(unittest.TestCase):
def setUp(self):
self.server = TranscriptionServer()
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
self.server.client_manager.start_times = {
'client1': time.time() - 120,
'client2': time.time() - 300
@@ -49,7 +51,7 @@ class TestServerConnection(unittest.TestCase):
'task': 'transcribe',
'model': 'tiny.en'
})
self.server.recv_audio(mock_websocket, "faster_whisper")
self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
@mock.patch('websockets.WebSocketCommonProtocol')
def test_recv_audio_exception_handling(self, mock_websocket):
@@ -61,7 +63,7 @@ class TestServerConnection(unittest.TestCase):
}), np.array([1, 2, 3]).tobytes()]
with self.assertLogs(level="ERROR"):
self.server.recv_audio(mock_websocket, "faster_whisper")
self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertNotIn(mock_websocket, self.server.client_manager.clients)
@@ -82,7 +84,6 @@ class TestServerInferenceAccuracy(unittest.TestCase):
cls.server_process.wait()
def setUp(self):
self.metric = evaluate.load("wer")
self.normalizer = EnglishTextNormalizer()
def check_prediction(self, srt_path):
@@ -94,11 +95,8 @@ class TestServerInferenceAccuracy(unittest.TestCase):
gt_normalized = self.normalizer(gt)
# calculate WER
wer = self.metric.compute(
predictions=[prediction_normalized],
references=[gt_normalized]
)
self.assertLess(wer, 0.05)
wer_score = jiwer.wer(gt_normalized, prediction_normalized)
self.assertLess(wer_score, 0.05)
def test_inference(self):
client = TranscriptionClient(
@@ -124,10 +122,10 @@ class TestExceptionHandling(unittest.TestCase):
@mock.patch('websockets.WebSocketCommonProtocol')
def test_connection_closed_exception(self, mock_websocket):
mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed")
mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed", rcvd_then_sent=mock.Mock())
with self.assertLogs(level="INFO") as log:
self.server.recv_audio(mock_websocket, "faster_whisper")
self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertTrue(any("Connection closed by client" in message for message in log.output))
@mock.patch('websockets.WebSocketCommonProtocol')
@@ -135,7 +133,7 @@ class TestExceptionHandling(unittest.TestCase):
mock_websocket.recv.return_value = "invalid json"
with self.assertLogs(level="ERROR") as log:
self.server.recv_audio(mock_websocket, "faster_whisper")
self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output))
@mock.patch('websockets.WebSocketCommonProtocol')
@@ -143,7 +141,7 @@ class TestExceptionHandling(unittest.TestCase):
mock_websocket.recv.side_effect = RuntimeError("Unexpected error")
with self.assertLogs(level="ERROR") as log:
self.server.recv_audio(mock_websocket, "faster_whisper")
self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
for message in log.output:
print(message)
print()
+16 -3
View File
@@ -30,7 +30,9 @@ class Client:
model="small",
srt_file_path="output.srt",
use_vad=True,
log_transcription=True
log_transcription=True,
max_clients=4,
max_connection_time=600,
):
"""
Initializes a Client instance for audio recording and streaming to a server.
@@ -59,6 +61,8 @@ class Client:
self.last_segment = None
self.last_received_segment = None
self.log_transcription = log_transcription
self.max_clients = max_clients
self.max_connection_time = max_connection_time
if translate:
self.task = "translate"
@@ -199,7 +203,9 @@ class Client:
"language": self.language,
"task": self.task,
"model": self.model,
"use_vad": self.use_vad
"use_vad": self.use_vad,
"max_clients": self.max_clients,
"max_connection_time": self.max_connection_time,
}
)
)
@@ -681,8 +687,15 @@ class TranscriptionClient(TranscriptionTeeClient):
output_recording_filename="./output_recording.wav",
output_transcription_path="./output.srt",
log_transcription=True,
max_clients=4,
max_connection_time=600,
):
self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad, log_transcription=log_transcription)
self.client = Client(
host, port, lang, translate, model, srt_file_path=output_transcription_path,
use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
max_connection_time=max_connection_time
)
if save_output_recording and not output_recording_filename.endswith(".wav"):
raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
if not output_transcription_path.endswith(".srt"):
+7 -1
View File
@@ -147,7 +147,7 @@ class TranscriptionServer:
RATE = 16000
def __init__(self):
self.client_manager = ClientManager()
self.client_manager = None
self.no_voice_activity_chunks = 0
self.use_vad = True
self.single_model = False
@@ -224,6 +224,12 @@ class TranscriptionServer:
logging.info("New client connected")
options = websocket.recv()
options = json.loads(options)
if self.client_manager is None:
max_clients = options.get('max_clients', 4)
max_connection_time = options.get('max_connection_time', 600)
self.client_manager = ClientManager(max_clients, max_connection_time)
self.use_vad = options.get('use_vad')
if self.client_manager.is_server_full(websocket, options):
websocket.close()