diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 29010bf..5d6a2bb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,4 +1,4 @@ -name: CI +name: Test & Build CI/CD on: push: @@ -7,46 +7,83 @@ on: tags: - v* pull_request: - branches: - - main + branches: [ main ] + types: [opened, synchronize, reopened] jobs: - build-and-push-package: + test: runs-on: ubuntu-latest - + timeout-minutes: 60 + strategy: + matrix: + python-version: [3.8, 3.9, '3.10', '3.11'] steps: - - name: Check Out Repository - uses: actions/checkout@v2 + - uses: actions/checkout@v2 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-version }} - - name: Set up Python - uses: actions/setup-python@v2 - with: - python-version: 3.8 - - - name: Set up FFmpeg - uses: FedericoCarboni/setup-ffmpeg@v2 - - - name: Install Additional requirements - run: | - sudo apt-get -y install portaudio19-dev wget - shell: bash + - name: Cache Python dependencies + uses: actions/cache@v2 + with: + path: | + ~/.cache/pip + !~/.cache/pip/log + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }} + restore-keys: | + ${{ runner.os }}-pip-${{ matrix.python-version }}- + + - name: Install system dependencies + run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements/server.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install -r requirements/client.txt + + - name: Run tests + run: | + echo "Running tests with Python ${{ matrix.python-version }}" + python -m unittest discover -s tests - - name: Install Client Requirements - run: pip install -r requirements/client.txt + build-and-push: + needs: test + runs-on: ubuntu-latest + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags') + steps: + - uses: actions/checkout@v2 - - name: Install Server Requirements - run: pip install -r requirements/server.txt + - name: Set up Python 3.8 + uses: actions/setup-python@v2 + with: + python-version: 3.8 - - name: Install Wheel for build - run: pip install wheel twine - - - name: Build wheel - run: | - python setup.py sdist bdist_wheel - - - name: Push package on Test PyPI - if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags') - uses: pypa/gh-action-pypi-publish@release/v1 - with: - user: __token__ - password: ${{ secrets.PYPI_API_TOKEN }} + - name: Cache Python dependencies + uses: actions/cache@v2 + with: + path: | + ~/.cache/pip + !~/.cache/pip/log + key: ubuntu-latest-pip-3.8-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }} + restore-keys: | + ubuntu-latest-pip-3.8- + + - name: Install system dependencies + run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev + + - name: Install Python dependencies + run: | + pip install -r requirements/server.txt + pip install -r requirements/client.txt + + - name: Build package + run: python setup.py sdist bdist_wheel + + - name: Publish package to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/assets/jfk.flac b/assets/jfk.flac new file mode 100644 index 0000000..e44b7c1 Binary files /dev/null and b/assets/jfk.flac differ diff --git a/requirements/server.txt b/requirements/server.txt index f7cf7f6..f9fe92f 100644 --- a/requirements/server.txt +++ b/requirements/server.txt @@ -8,3 +8,5 @@ kaldialign soundfile ffmpeg-python scipy +jiwer +evaluate \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 0000000..2afaba1 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,109 @@ +import json +import os +import scipy +import websocket +import unittest +from unittest.mock import patch, MagicMock +from whisper_live.client import TranscriptionClient, resample + + +class BaseTestCase(unittest.TestCase): + @patch('whisper_live.client.websocket.WebSocketApp') + @patch('whisper_live.client.pyaudio.PyAudio') + def setUp(self, mock_pyaudio, mock_websocket): + self.mock_pyaudio_instance = MagicMock() + mock_pyaudio.return_value = self.mock_pyaudio_instance + self.mock_stream = MagicMock() + self.mock_pyaudio_instance.open.return_value = self.mock_stream + + self.mock_ws_app = mock_websocket.return_value + self.mock_ws_app.send = MagicMock() + + self.client = TranscriptionClient(host='localhost', port=9090, lang="en").client + + self.mock_pyaudio = mock_pyaudio + self.mock_websocket = mock_websocket + + def tearDown(self): + self.client.close_websocket() + self.mock_pyaudio.stop() + self.mock_websocket.stop() + del self.client + + +class TestClientWebSocketCommunication(BaseTestCase): + def test_websocket_communication(self): + expected_url = 'ws://localhost:9090' + self.mock_websocket.assert_called() + self.assertEqual(self.mock_websocket.call_args[0][0], expected_url) + + +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, + }) + self.client.on_open(self.mock_ws_app) + self.mock_ws_app.send.assert_called_with(expected_message) + + def test_on_message(self): + message = json.dumps( + { + "uid": self.client.uid, + "message": "SERVER_READY", + "backend": "faster_whisper" + } + ) + self.client.on_message(self.mock_ws_app, message) + + message = json.dumps({ + "uid": self.client.uid, + "segments": [ + {"start": 0, "end": 1, "text": "Test transcript"}, + {"start": 1, "end": 2, "text": "Test transcript 2"}, + {"start": 2, "end": 3, "text": "Test transcript 3"} + ] + }) + self.client.on_message(self.mock_ws_app, message) + + # Assert that the transcript was updated correctly + self.assertEqual(len(self.client.transcript), 2) + self.assertEqual(self.client.transcript[1]['text'], "Test transcript 2") + + def test_on_close(self): + close_status_code = 1000 + close_msg = "Normal closure" + self.client.on_close(self.mock_ws_app, close_status_code, close_msg) + + self.assertFalse(self.client.recording) + self.assertFalse(self.client.server_error) + self.assertFalse(self.client.waiting) + + def test_on_error(self): + error_message = "Test Error" + self.client.on_error(self.mock_ws_app, error_message) + + self.assertTrue(self.client.server_error) + self.assertEqual(self.client.error_message, error_message) + + +class TestAudioResampling(unittest.TestCase): + def test_resample_audio(self): + original_audio = "assets/jfk.flac" + expected_sr = 16000 + resampled_audio = resample(original_audio, expected_sr) + + sr, _ = scipy.io.wavfile.read(resampled_audio) + self.assertEqual(sr, expected_sr) + + os.remove(resampled_audio) + + +class TestSendingAudioPacket(BaseTestCase): + def test_send_packet(self): + mock_audio_packet = b'\x00\x01\x02\x03' + self.client.send_packet_to_server(mock_audio_packet) + self.client.client_socket.send.assert_called_with(mock_audio_packet, websocket.ABNF.OPCODE_BINARY) diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..34c1288 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,105 @@ +import subprocess +import time +import json +import unittest +from unittest import mock + +import numpy as np +import evaluate +from whisper_live.server import TranscriptionServer +from whisper_live.client import TranscriptionClient +from whisper.normalizers import EnglishTextNormalizer + + +class TestTranscriptionServerInitialization(unittest.TestCase): + def test_initialization(self): + server = TranscriptionServer() + self.assertEqual(server.max_clients, 4) + self.assertEqual(server.max_connection_time, 600) + self.assertDictEqual(server.clients, {}) + self.assertDictEqual(server.websockets, {}) + self.assertDictEqual(server.clients_start_time, {}) + + +class TestGetWaitTime(unittest.TestCase): + def setUp(self): + self.server = TranscriptionServer() + self.server.clients_start_time = { + 'client1': time.time() - 120, + 'client2': time.time() - 300 + } + self.server.max_connection_time = 600 + + def test_get_wait_time(self): + expected_wait_time = (600 - (time.time() - self.server.clients_start_time['client2'])) / 60 + print(self.server.get_wait_time(), expected_wait_time) + self.assertAlmostEqual(self.server.get_wait_time(), expected_wait_time, places=2) + + +class TestServerConnection(unittest.TestCase): + def setUp(self): + self.server = TranscriptionServer() + + @mock.patch('websockets.WebSocketCommonProtocol') + def test_connection(self, mock_websocket): + mock_websocket.recv.return_value = json.dumps({ + 'uid': 'test_client', + 'language': 'en', + 'task': 'transcribe', + 'model': 'tiny.en' + }) + self.server.recv_audio(mock_websocket, "faster_whisper") + + + @mock.patch('websockets.WebSocketCommonProtocol') + def test_recv_audio_exception_handling(self, mock_websocket): + mock_websocket.recv.side_effect = [json.dumps({ + 'uid': 'test_client', + 'language': 'en', + 'task': 'transcribe', + 'model': 'tiny.en' + }), np.array([1, 2, 3]).tobytes()] + + with self.assertLogs(level="ERROR"): + self.server.recv_audio(mock_websocket, "faster_whisper") + + self.assertNotIn(mock_websocket, self.server.clients) + + +class TestServerInferenceAccuracy(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.server_process = subprocess.Popen(["python", "run_server.py"]) # Adjust the command as needed + time.sleep(2) + + @classmethod + def tearDownClass(cls): + cls.server_process.terminate() + cls.server_process.wait() + + @mock.patch('pyaudio.PyAudio') + def setUp(self, mock_pyaudio): + self.mock_pyaudio = mock_pyaudio.return_value + self.mock_stream = mock.MagicMock() + self.mock_pyaudio.open.return_value = self.mock_stream + self.metric = evaluate.load("wer") + self.normalizer = EnglishTextNormalizer() + self.client = TranscriptionClient( + "localhost", "9090", model="base.en", lang="en", + ) + + def test_inference(self): + gt = "And so my fellow Americans, ask not, what your country can do for you. Ask what you can do for your country!" + self.client("assets/jfk.flac") + with open("output.srt", "r") as f: + lines = f.readlines() + prediction = " ".join([l.strip() for l in lines[2::4]]) + prediction_normalized = self.normalizer(prediction) + gt_normalized = self.normalizer(gt) + + # calculate WER + wer = self.metric.compute( + predictions=[prediction_normalized], + references=[gt_normalized] + ) + self.assertLess(wer, 0.05) diff --git a/tests/test_vad.py b/tests/test_vad.py new file mode 100644 index 0000000..8ca7ee3 --- /dev/null +++ b/tests/test_vad.py @@ -0,0 +1,28 @@ +import unittest +import numpy as np +import torch +import scipy.io as sio +from whisper_live.tensorrt_utils import load_audio +from whisper_live.vad import VoiceActivityDetection + + +class TestVoiceActivityDetection(unittest.TestCase): + def setUp(self): + self.vad = VoiceActivityDetection() + self.sample_rate = 16000 + + def generate_silence(self, duration_seconds): + return np.zeros(int(self.sample_rate * duration_seconds), dtype=np.float32) + + def load_speech_segment(self, filepath): + return load_audio(filepath) + + def test_vad_silence_detection(self): + silence = self.generate_silence(3) + speech_prob = self.vad(torch.from_numpy(silence.copy()), self.sample_rate).item() + self.assertLess(speech_prob, 0.5, "VAD incorrectly identified silence as speech.") + + def test_vad_speech_detection(self): + audio_tensor = torch.from_numpy(load_audio("assets/jfk.flac")) + speech_prob = self.vad(audio_tensor, self.sample_rate).item() + self.assertGreater(speech_prob, 0.5, "VAD failed to identify speech segment.") \ No newline at end of file diff --git a/whisper_live/client.py b/whisper_live/client.py index f8df387..15850a0 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -233,10 +233,15 @@ class Client: print(element) def on_error(self, ws, error): - print(error) + print(f"[ERROR] WebSocket Error: {error}") + self.server_error = True + self.error_message = error def on_close(self, ws, close_status_code, close_msg): print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}") + self.recording = False + self.server_error = False + self.waiting = False def on_open(self, ws): """ diff --git a/whisper_live/server.py b/whisper_live/server.py index 4a8a1d8..e9e8493 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -1,5 +1,4 @@ import os -import websockets import time import threading import json @@ -12,10 +11,8 @@ from websockets.sync.server import serve import torch import numpy as np -import queue from whisper_live.vad import VoiceActivityDetection -from scipy.io.wavfile import write import functools from whisper_live.vad import VoiceActivityDetection @@ -23,7 +20,7 @@ from whisper_live.transcriber import WhisperModel try: from whisper_live.transcriber_tensorrt import WhisperTRTLLM except Exception as e: - logging.warn("cannot import WhisperTRTLLM") + pass class TranscriptionServer: diff --git a/whisper_live/vad.py b/whisper_live/vad.py index 53170f9..31edac8 100644 --- a/whisper_live/vad.py +++ b/whisper_live/vad.py @@ -10,9 +10,7 @@ import onnxruntime class VoiceActivityDetection(): def __init__(self, force_onnx_cpu=True): - print("downloading ONNX model...") path = self.download() - print("loading session") opts = onnxruntime.SessionOptions() opts.log_severity_level = 3 @@ -20,13 +18,11 @@ class VoiceActivityDetection(): opts.inter_op_num_threads = 1 opts.intra_op_num_threads = 1 - print("loading onnx model") if force_onnx_cpu and 'CPUExecutionProvider' in onnxruntime.get_available_providers(): self.session = onnxruntime.InferenceSession(path, providers=['CPUExecutionProvider'], sess_options=opts) else: self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts) - print("reset states") self.reset_states() self.sample_rates = [8000, 16000] @@ -110,7 +106,6 @@ class VoiceActivityDetection(): # Check if the model file already exists if not os.path.exists(model_filename): # If it doesn't exist, download the model using wget - print("Downloading VAD ONNX model...") try: subprocess.run(["wget", "-O", model_filename, model_url], check=True) except subprocess.CalledProcessError: