bump version 0.2.0

Docker ghcr ci

.github/workflows/ci.yml

@@ -1,4 +1,4 @@
-name: CI
+name: Test & Build CI/CD
 
 on:
   push:
@@ -7,45 +7,161 @@ on:
     tags:
       - v*
   pull_request:
-    branches:
+    branches: [ main ]
-      - main
+    types: [opened, synchronize, reopened]
 
 jobs:
-  build-and-push-package:
+  run-tests:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: [3.8, 3.9, '3.10', 3.11]
+    steps:
+      - uses: actions/checkout@v2
+      
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - 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
+  
+  check-code-format:
+    runs-on: ubuntu-22.04
+    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
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install flake8
+
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          # 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
+
+  build-and-push-docker-cpu:
+    needs: [run-tests, check-code-format]
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    steps:
+      - uses: actions/checkout@v2
+      
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GHCR_TOKEN }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.cpu
+          push: true
+          tags: ghcr.io/collabora/whisperlive-cpu:latest
+  
+  build-and-push-docker-gpu:
+    needs: [run-tests, check-code-format, build-and-push-docker-cpu]
+    timeout-minutes: 20
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GHCR_TOKEN }}
+
+      - name: Docker Prune
+        run: docker system prune -af
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker GPU image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.gpu
+          push: true
+          tags: ghcr.io/collabora/whisperlive-gpu:latest
+
+  publish-to-pypi:
+    needs: [run-tests, check-code-format]
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set up Python 3.8
         uses: actions/setup-python@v2
         with:
           python-version: 3.8
 
-    - name: Set up FFmpeg
+      - name: Cache Python dependencies
-      uses: FedericoCarboni/setup-ffmpeg@v2
+        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 Additional requirements
+      - name: Install system dependencies
+        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+
+      - name: Install Python dependencies
         run: |
-        sudo apt-get -y install portaudio19-dev wget
+          pip install -r requirements/server.txt
-      shell: bash
+          pip install -r requirements/client.txt
 
-    - name: Install Client Requirements
+      - name: Build package
-      run: pip install -r requirements/client.txt
+        run: python setup.py sdist bdist_wheel
 
-    - name: Install Server Requirements
+      - name: Publish package to PyPI
-      run: pip install -r requirements/server.txt
-
-    - 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__

Audio-Transcription-Chrome/background.js

@@ -157,7 +157,8 @@ async function startCapture(options) {
           multilingual: options.useMultilingual,
           language: options.language,
           task: options.task,
-          modelSize: options.modelSize
+          modelSize: options.modelSize,
+          useVad: options.useVad,
         },
       });
     } else {

Audio-Transcription-Chrome/options.js

@@ -99,7 +99,8 @@ async function startRecord(option) {
           uid: uuid,
           language: option.language,
           task: option.task,
-          model: option.modelSize
+          model: option.modelSize,
+          use_vad: option.useVad
         })
       );
     };

Audio-Transcription-Chrome/popup.html

@@ -15,6 +15,10 @@
     <input type="checkbox" id="useServerCheckbox">
     <label for="useServerCheckbox">Use Collabora Whisper-Live Server</label>
   </div>
+  <div class="checkbox-container">
+    <input type="checkbox" id="useVadCheckbox">
+    <label for="useVadCheckbox">Use Voice Activity Detection</label>
+  </div>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>
     <select id="languageDropdown">

Audio-Transcription-Chrome/popup.js

@@ -4,6 +4,7 @@ document.addEventListener("DOMContentLoaded", function () {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
+  const useVadCheckbox = document.getElementById("useVadCheckbox");
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
   const modelSizeDropdown = document.getElementById('modelSizeDropdown');
@@ -31,6 +32,12 @@ document.addEventListener("DOMContentLoaded", function () {
     }
   });
 
+  chrome.storage.local.get("useVadState", ({ useVadState }) => {
+    if (useVadState !== undefined) {
+      useVadCheckbox.checked = useVadState;
+    }
+  });
+
   chrome.storage.local.get("selectedLanguage", ({ selectedLanguage: storedLanguage }) => {
     if (storedLanguage !== undefined) {
       languageDropdown.value = storedLanguage;
@@ -79,7 +86,8 @@ document.addEventListener("DOMContentLoaded", function () {
         port: port,
         language: selectedLanguage,
         task: selectedTask,
-        modelSize: selectedModelSize
+        modelSize: selectedModelSize,
+        useVad: useVadCheckbox.checked,
       }, () => {
         // Update capturing state in storage and toggle the buttons
         chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => {
@@ -119,6 +127,7 @@ document.addEventListener("DOMContentLoaded", function () {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing;
+    useVadCheckbox.disabled = isCapturing;
     modelSizeDropdown.disabled = isCapturing;
     languageDropdown.disabled = isCapturing;
     taskDropdown.disabled = isCapturing; 
@@ -132,6 +141,11 @@ document.addEventListener("DOMContentLoaded", function () {
     chrome.storage.local.set({ useServerState });
   });
 
+  useVadCheckbox.addEventListener("change", () => {
+    const useVadState = useVadCheckbox.checked;
+    chrome.storage.local.set({ useVadState });
+  });
+
   languageDropdown.addEventListener('change', function() {
     if (languageDropdown.value === "") {
       selectedLanguage = null;

Audio-Transcription-Firefox/content.js

@@ -74,7 +74,8 @@ function startRecording(data) {
             uid: uuid,
             language: data.language,
             task: data.task,
-            model: data.modelSize
+            model: data.modelSize,
+            use_vad: data.useVad
         })
       );
     };

Audio-Transcription-Firefox/popup.html

@@ -15,6 +15,10 @@
     <input type="checkbox" id="useServerCheckbox">
     <label for="useServerCheckbox">Use Collabora Whisper-Live Server</label>
   </div>
+  <div class="checkbox-container">
+    <input type="checkbox" id="useVadCheckbox">
+    <label for="useVadCheckbox">Use Voice Activity Detection</label>
+  </div>
   <textarea id="waitTextBox" style="display: none;"></textarea>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>

Audio-Transcription-Firefox/popup.js

@@ -3,6 +3,7 @@ document.addEventListener("DOMContentLoaded", function() {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
+  const useVadCheckbox = document.getElementById("useVadCheckbox");
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
   const modelSizeDropdown = document.getElementById('modelSizeDropdown');
@@ -34,6 +35,12 @@ document.addEventListener("DOMContentLoaded", function() {
     }
   });
 
+  browser.storage.local.get("useVadState", ({ useVadState }) => {
+    if (useVadState !== undefined) {
+      useVadCheckbox.checked = useVadState;
+    }
+  });
+
   browser.storage.local.get("selectedLanguage", ({ selectedLanguage: storedLanguage }) => {
     if (storedLanguage !== undefined) {
       languageDropdown.value = storedLanguage;
@@ -76,7 +83,8 @@ document.addEventListener("DOMContentLoaded", function() {
               port: port,
               language: selectedLanguage,
               task: selectedTask,
-              modelSize: selectedModelSize
+              modelSize: selectedModelSize,
+              useVad: useVadCheckbox.checked,
             } 
           });
         toggleCaptureButtons(true);
@@ -115,6 +123,7 @@ document.addEventListener("DOMContentLoaded", function() {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing;
+    useVadCheckbox.disabled = isCapturing;
     modelSizeDropdown.disabled = isCapturing;
     languageDropdown.disabled = isCapturing;
     taskDropdown.disabled = isCapturing; 
@@ -128,6 +137,11 @@ document.addEventListener("DOMContentLoaded", function() {
     browser.storage.local.set({ useServerState });
   });
 
+  useVadCheckbox.addEventListener("change", () => {
+    const useVadState = useVadCheckbox.checked;
+    browser.storage.local.set({ useVadState });
+  });
+
   languageDropdown.addEventListener('change', function() {
     if (languageDropdown.value === "") {
       selectedLanguage = null;

README.md

@@ -1,9 +1,15 @@
-# whisper-live
+# WhisperLive
-A nearly-live implementation of OpenAI's Whisper.
 
-This project is a real-time transcription application that uses the OpenAI Whisper model to convert speech input into text output. It can be used to transcribe both live audio input from microphone and pre-recorded audio files.
+<h2 align="center">
+  <a href="https://www.youtube.com/watch?v=0PHWCApIcCI"><img
+src="https://img.youtube.com/vi/0PHWCApIcCI/0.jpg" style="background-color:rgba(0,0,0,0);" height=300 alt="WhisperLive"></a>
+  <br><br>A nearly-live implementation of OpenAI's Whisper.
+<br><br>
+</h2>
 
-Unlike traditional speech recognition systems that rely on continuous audio streaming, we use [voice activity detection (VAD)](https://github.com/snakers4/silero-vad) to detect the presence of speech and only send the audio data to whisper when speech is detected. This helps to reduce the amount of data sent to the whisper model and improves the accuracy of the transcription output.
+This project is a real-time transcription application that uses the OpenAI Whisper model
+to convert speech input into text output. It can be used to transcribe both live audio
+input from microphone and pre-recorded audio files.
 
 ## Installation
 - Install PyAudio and ffmpeg
@@ -50,7 +56,7 @@ python3 run_server.py -p 9090 \
 
 
 ### Running the Client
-- To transcribe an audio file:
+- Initializing the client:
 ```python
 from whisper_live.client import TranscriptionClient
 client = TranscriptionClient(
@@ -58,58 +64,43 @@ client = TranscriptionClient(
   9090,
   lang="en",
   translate=False,
-  model="small"
+  model="small",
+  use_vad=False,
 )
+```
+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.
 
+- Trancribe an audio file:
+```python
 client("tests/jfk.wav")
 ```
-This command transcribes the specified audio file (audio.wav) using the Whisper model. 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.
 
 - To transcribe from microphone:
 ```python
-from whisper_live.client import TranscriptionClient
-client = TranscriptionClient(
-  "localhost",
-  9090,
-  lang="hi",
-  translate=True,
-  model="small"
-)
 client()
 ```
-This command captures audio from the microphone and sends it to the server for transcription. It uses the multilingual model with `hi` as the selected language. We use whisper `small` by default but can be changed to any other option based on the requirements and the hardware running the server.
 
 - To transcribe from a HLS stream:
 ```python
-from whisper_live.client import TranscriptionClient
-client = TranscriptionClient(host, port, lang="en", translate=False) 
 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") 
 ```
-This command streams audio into the server from a HLS stream. It uses the same options as the previous command, using the multilingual model and specifying the target language and task.
 
-## Transcribe audio from browser
+## Browser Extensions
-- Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server)
+- Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server).
-
+- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) for setup instructions.
-### Chrome Extension
-- Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) to use Chrome extension.
-
-### Firefox Extension
-- Refer to [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) to use Mozilla Firefox extension.
 
 ## Whisper Live Server in Docker
 - GPU
   - Faster-Whisper
   ```bash
-  docker build . -t whisper-live -f docker/Dockerfile.gpu
+  docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest
-  docker run -it --gpus all -p 9090:9090 whisper-live:latest
   ```
 
   - TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup docker and use TensorRT backend. We provide a pre-built docker image which has TensorRT-LLM built and ready to use.
 
 - CPU
 ```bash
-docker build . -t whisper-live -f docker/Dockerfile.cpu
+docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu:latest
-docker run -it -p 9090:9090 whisper-live:latest
 ```
 **Note**: By default we use "small" model size. To build docker image for a different model size, change the size in server.py and then build the docker image.
 
@@ -140,6 +131,5 @@ We are available to help you with both Open Source and proprietary AI projects.
   publisher = {GitHub},
   journal = {GitHub repository},
   howpublished = {\url{https://github.com/snakers4/silero-vad}},
-  commit = {insert_some_commit_here},
   email = {hello@silero.ai}
 }

TensorRT_whisper.md

@@ -21,7 +21,7 @@ docker pull ghcr.io/collabora/whisperbot-base:latest
 ```bash
 docker run -it --gpus all --shm-size=8g \
        --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
-       -v /path/to/WhisperLive:/home/WhisperLive \
+       -p 9090:9090 -v /path/to/WhisperLive:/home/WhisperLive \
        ghcr.io/collabora/whisperbot-base:latest
 ```
 
@@ -48,7 +48,7 @@ bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small
 cd /home/WhisperLive
 
 # Install requirements
-bash scripts/setup.sh
+apt update && bash scripts/setup.sh
 pip install -r requirements/server.txt
 
 # Required to create mel spectogram

docker/Dockerfile.cpu

@@ -1,45 +1,24 @@
-FROM ubuntu:focal
+FROM python:3.8-slim-buster
 
 ARG DEBIAN_FRONTEND=noninteractive
 
-# Remove any third-party apt sources to avoid issues with expiring keys.
+RUN apt-get update && apt-get install -y --no-install-recommends \
-RUN rm -f /etc/apt/sources.list.d/*.list
-
-# Install some basic utilities.
-RUN apt-get update && apt-get install -y \
     curl \
     ca-certificates \
     sudo \
     git \
     bzip2 \
     libx11-6 \
+ && apt-get clean \
  && rm -rf /var/lib/apt/lists/*
 
-RUN apt update
-
-# install python
-RUN apt install software-properties-common -y && \
-    add-apt-repository ppa:deadsnakes/ppa && \
-    apt update
-
-RUN apt install python3-dev -y && \
-    apt install python-is-python3
-    
-
-# install pip
-RUN apt install python3-pip -y
-
-# Create a working directory.
-RUN mkdir /app
 WORKDIR /app
 
-COPY scripts/setup.sh /app
+COPY scripts/setup.sh requirements/server.txt /app/
-COPY requirements/ /app
 
-RUN bash setup.sh
+RUN apt update && bash setup.sh && pip install -r server.txt
-RUN pip install -r server.txt
 
 COPY whisper_live /app/whisper_live
-
 COPY run_server.py /app
+
 CMD ["python", "run_server.py"]

docker/Dockerfile.gpu

@@ -1,44 +1,30 @@
-FROM nvidia/cuda:11.2.2-cudnn8-runtime-ubuntu20.04
+FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
-
 ARG DEBIAN_FRONTEND=noninteractive
 
 # Remove any third-party apt sources to avoid issues with expiring keys.
 RUN rm -f /etc/apt/sources.list.d/*.list
 
 # Install some basic utilities.
-RUN apt-get update && apt-get install -y \
+RUN apt-get update && apt-get install -y --no-install-recommends \
     curl \
     ca-certificates \
     sudo \
     git \
     bzip2 \
     libx11-6 \
+    python3-dev \
+    python3-pip \
+    && python3 -m pip install --upgrade pip \
     && rm -rf /var/lib/apt/lists/*
 
-RUN apt update
-
-# install python
-RUN apt install software-properties-common -y && \
-    add-apt-repository ppa:deadsnakes/ppa && \
-    apt update
-
-RUN apt install python3-dev -y && \
-    apt install python-is-python3
-    
-
-# install pip
-RUN apt install python3-pip -y
-
 # Create a working directory.
 RUN mkdir /app
 WORKDIR /app
 
-COPY scripts/setup.sh /app
+COPY scripts/setup.sh requirements/server.txt /app
-COPY requirements/ /app
 
-RUN apt update --fix-missing
+RUN apt update && bash setup.sh && rm setup.sh
-RUN bash setup.sh
+RUN pip install -r server.txt && rm server.txt
-RUN pip install -r server.txt
 
 COPY whisper_live /app/whisper_live
 

requirements/server.txt

@@ -8,3 +8,5 @@ kaldialign
 soundfile
 ffmpeg-python
 scipy
+jiwer
+evaluate

setup.py

@@ -10,7 +10,8 @@ HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
 
 # This call to setup() does all the work
-setup(name="whisper-live",
+setup(
+    name="whisper-live",
     version=__version__,
     description="A nearly-live implementation of OpenAI's Whisper.",
     long_description=README,
@@ -32,7 +33,8 @@ setup(name="whisper-live",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
     ],
     packages=find_packages(
-        exclude=("examples",
+        exclude=(
+            "examples",
             "Audio-Transcription-Chrome",
             "Audio-Transcription-Firefox",
             "requirements",

tests/test_client.py

@@ -0,0 +1,111 @@
+import json
+import os
+import scipy
+import websocket
+import unittest
+from unittest.mock import patch, MagicMock
+from whisper_live.client import TranscriptionClient
+from whisper_live.utils import 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,
+            "use_vad": True
+        })
+        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)

tests/test_server.py

@@ -0,0 +1,137 @@
+import subprocess
+import time
+import json
+import unittest
+from unittest import mock
+
+import numpy as np
+import evaluate
+
+from websockets.exceptions import ConnectionClosed
+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.client_manager.max_clients, 4)
+        self.assertEqual(server.client_manager.max_connection_time, 600)
+        self.assertDictEqual(server.client_manager.clients, {})
+        self.assertDictEqual(server.client_manager.start_times, {})
+
+
+class TestGetWaitTime(unittest.TestCase):
+    def setUp(self):
+        self.server = TranscriptionServer()
+        self.server.client_manager.start_times = {
+            'client1': time.time() - 120,
+            'client2': time.time() - 300
+        }
+        self.server.client_manager.max_connection_time = 600
+
+    def test_get_wait_time(self):
+        expected_wait_time = (600 - (time.time() - self.server.client_manager.start_times['client2'])) / 60
+        print(self.server.client_manager.get_wait_time(), expected_wait_time)
+        self.assertAlmostEqual(self.server.client_manager.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.client_manager.clients)
+
+
+class TestServerInferenceAccuracy(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.server_process = subprocess.Popen(["python", "run_server.py"])
+        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([line.strip() for line 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)
+
+
+class TestExceptionHandling(unittest.TestCase):
+    def setUp(self):
+        self.server = TranscriptionServer()
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_connection_closed_exception(self, mock_websocket):
+        mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed")
+
+        with self.assertLogs(level="INFO") as log:
+            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.assertTrue(any("Connection closed by client" in message for message in log.output))
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_json_decode_exception(self, mock_websocket):
+        mock_websocket.recv.return_value = "invalid json"
+
+        with self.assertLogs(level="ERROR") as log:
+            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output))
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_unexpected_exception_handling(self, mock_websocket):
+        mock_websocket.recv.side_effect = RuntimeError("Unexpected error")
+
+        with self.assertLogs(level="ERROR") as log:
+            self.server.recv_audio(mock_websocket, "faster_whisper")
+            for message in log.output:
+                print(message)
+            print()
+            self.assertTrue(any("Unexpected error" in message for message in log.output))

tests/test_vad.py

@@ -0,0 +1,26 @@
+import unittest
+import numpy as np
+from whisper_live.tensorrt_utils import load_audio
+from whisper_live.vad import VoiceActivityDetector
+
+
+class TestVoiceActivityDetection(unittest.TestCase):
+    def setUp(self):
+        self.vad = VoiceActivityDetector()
+        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)
+        is_speech_present = self.vad(silence.copy())
+        self.assertFalse(is_speech_present, "VAD incorrectly identified silence as speech.")
+
+    def test_vad_speech_detection(self):
+        audio_tensor = load_audio("assets/jfk.flac")
+        is_speech_present = self.vad(audio_tensor)
+        self.assertTrue(is_speech_present, "VAD failed to identify speech segment.")

whisper_live/__version__.py

@@ -1 +1 @@
-__version__="0.1.0"
+__version__ = "0.2.0"

whisper_live/client.py

@@ -2,68 +2,14 @@ import os
 import wave
 
 import numpy as np
-import scipy
-import ffmpeg
 import pyaudio
 import threading
-import textwrap
 import json
 import websocket
 import uuid
 import time
-
+import ffmpeg
-
+import whisper_live.utils as utils
-def format_time(s):
-    """Convert seconds (float) to SRT time format."""
-    hours = int(s // 3600)
-    minutes = int((s % 3600) // 60)
-    seconds = int(s % 60)
-    milliseconds = int((s - int(s)) * 1000)
-    return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
-
-def create_srt_file(segments, output_file):
-    with open(output_file, 'w', encoding='utf-8') as srt_file:
-        segment_number = 1
-        for segment in segments:
-            start_time = format_time(float(segment['start']))
-            end_time = format_time(float(segment['end']))
-            text = segment['text']
-            
-            srt_file.write(f"{segment_number}\n")
-            srt_file.write(f"{start_time} --> {end_time}\n")
-            srt_file.write(f"{text}\n\n")
-            
-            segment_number += 1
-
-
-def resample(file: str, sr: int = 16000):
-    """
-    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
-    Open an audio file and read as mono waveform, resampling as necessary,
-    save the resampled audio
-
-    Args:
-        file (str): The audio file to open
-        sr (int): The sample rate to resample the audio if necessary
-    
-    Returns:
-        resampled_file (str): The resampled audio file
-    """
-    try:
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
-        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
-        out, _ = (
-            ffmpeg.input(file, threads=0)
-            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
-            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
-        )
-    except ffmpeg.Error as e:
-        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
-    np_buffer = np.frombuffer(out, dtype=np.int16)
-
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
-    return resampled_file
 
 
 class Client:
@@ -71,6 +17,7 @@ class Client:
     Handles audio recording, streaming, and communication with a server using WebSocket.
     """
     INSTANCES = {}
+    END_OF_AUDIO = "END_OF_AUDIO"
 
     def __init__(
         self,
@@ -79,7 +26,8 @@ class Client:
         lang=None,
         translate=False,
         model="small",
-        srt_file_path="output.srt"
+        srt_file_path="output.srt",
+        use_vad=True
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -109,6 +57,8 @@ class Client:
         self.model = model
         self.server_error = False
         self.srt_file_path = srt_file_path
+        self.use_vad = use_vad
+        self.last_recieved_segment = None
 
         if translate:
             self.task = "translate"
@@ -150,6 +100,40 @@ class Client:
         self.transcript = []
         print("[INFO]: * recording")
 
+    def handle_status_messages(self, message_data):
+        """Handles server status messages."""
+        status = message_data["status"]
+        if status == "WAIT":
+            self.waiting = True
+            print(f"[INFO]: Server is full. Estimated wait time {round(message_data['message'])} minutes.")
+        elif status == "ERROR":
+            print(f"Message from Server: {message_data['message']}")
+            self.server_error = True
+        elif status == "WARNING":
+            print(f"Message from Server: {message_data['message']}")
+
+    def process_segments(self, segments):
+        """Processes transcript segments."""
+        text = []
+        for i, seg in enumerate(segments):
+            if not text or text[-1] != seg["text"]:
+                text.append(seg["text"])
+                if i == len(segments) - 1:
+                    self.last_segment = seg
+                elif (self.server_backend == "faster_whisper" and
+                      (not self.transcript or
+                        float(seg['start']) >= float(self.transcript[-1]['end']))):
+                    self.transcript.append(seg)
+        # update last received segment and last valild responsne time
+        if self.last_recieved_segment is None or self.last_recieved_segment != segments[-1]["text"]:
+            self.last_response_recieved = time.time()
+            self.last_recieved_segment = segments[-1]["text"]
+
+        # Truncate to last 3 entries for brevity.
+        text = text[-3:]
+        utils.clear_screen()
+        utils.print_transcript(text)
+
     def on_message(self, ws, message):
         """
         Callback function called when a message is received from the server.
@@ -163,7 +147,6 @@ class Client:
             message (str): The received message from the server.
 
         """
-        self.last_response_recieved = time.time()
         message = json.loads(message)
 
         if self.uid != message.get("uid"):
@@ -171,21 +154,15 @@ class Client:
             return
 
         if "status" in message.keys():
-            if message["status"] == "WAIT":
+            self.handle_status_messages(message)
-                self.waiting = True
-                print(
-                    f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes."
-                )
-            elif message["status"] == "ERROR":
-                print(f"Message from Server: {message['message']}")
-                self.server_error = True
             return
 
         if "message" in message.keys() and message["message"] == "DISCONNECT":
-            print("[INFO]: Server overtime disconnected.")
+            print("[INFO]: Server disconnected due to overtime.")
             self.recording = False
 
         if "message" in message.keys() and message["message"] == "SERVER_READY":
+            self.last_response_recieved = time.time()
             self.recording = True
             self.server_backend = message["backend"]
             print(f"[INFO]: Server Running with backend {self.server_backend}")
@@ -199,44 +176,19 @@ class Client:
             )
             return
 
-        if "segments" not in message.keys():
+        if "segments" in message.keys():
-            return
+            self.process_segments(message["segments"])
-
-        message = message["segments"]
-        text = []
-        n_segments = len(message)
-
-        if n_segments:
-            for i, seg in enumerate(message):
-                if text and text[-1] == seg["text"]:
-                    # already got it
-                    continue
-                text.append(seg["text"])
-
-                if i == n_segments-1: 
-                    self.last_segment = seg
-                elif self.server_backend == "faster_whisper":
-                    if not len(self.transcript) or float(seg['start']) >= float(self.transcript[-1]['end']):
-                        self.transcript.append(seg)
-
-        # keep only last 3
-        if len(text) > 3:
-            text = text[-3:]
-        wrapper = textwrap.TextWrapper(width=60)
-        word_list = wrapper.wrap(text="".join(text))
-        # Print each line.
-        if os.name == "nt":
-            os.system("cls")
-        else:
-            os.system("clear")
-        for element in word_list:
-            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):
         """
@@ -257,6 +209,7 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
+                    "use_vad": self.use_vad
                 }
             )
         )
@@ -331,7 +284,7 @@ class Client:
                 assert self.last_response_recieved
                 while time.time() - self.last_response_recieved < self.disconnect_if_no_response_for:
                     continue
-
+                self.send_packet_to_server(Client.END_OF_AUDIO.encode('utf-8'))
                 if self.server_backend == "faster_whisper":
                     self.write_srt_file(self.srt_file_path)
                 self.stream.close()
@@ -428,7 +381,6 @@ class Client:
 
         print("[INFO]: HLS stream processing finished.")
 
-
     def record(self, out_file="output_recording.wav"):
         """
         Record audio data from the input stream and save it to a WAV file.
@@ -443,7 +395,8 @@ class Client:
         the method combines all the saved audio chunks into the specified `out_file`.
 
         Args:
-            out_file (str, optional): The name of the output WAV file to save the entire recording. Default is "output_recording.wav".
+            out_file (str, optional): The name of the output WAV file to save the entire recording.
+                                      Default is "output_recording.wav".
 
         """
         n_audio_file = 0
@@ -453,7 +406,7 @@ class Client:
             for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
                 if not self.recording:
                     break
-                data = self.stream.read(self.chunk, exception_on_overflow = False)
+                data = self.stream.read(self.chunk, exception_on_overflow=False)
                 self.frames += data
 
                 audio_array = Client.bytes_to_float_array(data)
@@ -527,7 +480,7 @@ class Client:
 
     def write_srt_file(self, output_path="output.srt"):
         self.transcript.append(self.last_segment)
-        create_srt_file(self.transcript, output_path)
+        utils.create_srt_file(self.transcript, output_path)
 
 
 class TranscriptionClient:
@@ -553,14 +506,8 @@ class TranscriptionClient:
         transcription_client()
         ```
     """
-    def __init__(self,
+    def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True):
-        host,
+        self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
-        port,
-        lang=None,
-        translate=False,
-        model="small",
-    ):
-        self.client = Client(host, port, lang, translate, model)
 
     def __call__(self, audio=None, hls_url=None):
         """
@@ -584,7 +531,7 @@ class TranscriptionClient:
         if hls_url is not None:
             self.client.process_hls_stream(hls_url)
         elif audio is not None:
-            resampled_file = resample(audio)
+            resampled_file = utils.resample(audio)
             self.client.play_file(resampled_file)
         else:
             self.client.record()

whisper_live/tensorrt_utils.py

@@ -214,7 +214,7 @@ def store_transcripts(filename: Pathlike, texts: Iterable[Tuple[str, str,
             print(f"{cut_id}:\thyp={hyp}", file=f)
 
 
-def write_error_stats(
+def write_error_stats(                                              # noqa: C901
     f: TextIO,
     test_set_name: str,
     results: List[Tuple[str, str]],

whisper_live/transcriber.py

@@ -400,7 +400,7 @@ class WhisperModel:
 
         return segments, info
 
-    def generate_segments(
+    def generate_segments(                                                  # noqa: C901
         self,
         features: np.ndarray,
         tokenizer: Tokenizer,
@@ -425,7 +425,7 @@ class WhisperModel:
         all_segments = []
         while seek < content_frames:
             time_offset = seek * self.feature_extractor.time_per_frame
-            segment = features[:, seek : seek + self.feature_extractor.nb_max_frames]
+            segment = features[:, seek:seek + self.feature_extractor.nb_max_frames]
             segment_size = min(
                 self.feature_extractor.nb_max_frames, content_frames - seek
             )
@@ -749,7 +749,7 @@ class WhisperModel:
 
         if previous_tokens:
             prompt.append(tokenizer.sot_prev)
-            prompt.extend(previous_tokens[-(self.max_length // 2 - 1) :])
+            prompt.extend(previous_tokens[-(self.max_length // 2 - 1):])
 
         prompt.extend(tokenizer.sot_sequence)
 
@@ -766,7 +766,7 @@ class WhisperModel:
 
         return prompt
 
-    def add_word_timestamps(
+    def add_word_timestamps(                                                    # noqa: C901
         self,
         segments: List[dict],
         tokenizer: Tokenizer,

whisper_live/transcriber_tensorrt.py

@@ -1,17 +1,14 @@
-import argparse
 import json
 import re
-import time
 from collections import OrderedDict
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional, TextIO, Tuple, Union
+from typing import Union
 
 import torch
 import numpy as np
+import torch.nn.functional as F
 from whisper.tokenizer import get_tokenizer
-from whisper_live.tensorrt_utils import (mel_filters, store_transcripts,
+from whisper_live.tensorrt_utils import (mel_filters, load_audio_wav_format, pad_or_trim, load_audio)
-                           write_error_stats, load_audio_wav_format,
-                           pad_or_trim, load_audio)
 
 import tensorrt_llm
 import tensorrt_llm.logger as logger
@@ -38,8 +35,6 @@ class WhisperEncoding:
         with open(config_path, 'r') as f:
             config = json.load(f)
 
-        use_gpt_attention_plugin = config['plugin_config'][
-            'gpt_attention_plugin']
         dtype = config['builder_config']['precision']
         n_mels = config['builder_config']['n_mels']
         num_languages = config['builder_config']['num_languages']
@@ -176,16 +171,8 @@ class WhisperDecoding:
 
 class WhisperTRTLLM(object):
 
-    def __init__(
+    def __init__(self, engine_dir, assets_dir=None, device=None, is_multilingual=False,
-        self,
+                 language="en", task="transcribe"):
-        engine_dir,
-        debug_mode=False,
-        assets_dir=None,
-        device=None,
-        is_multilingual=False,
-        language="en",
-        task="transcribe"
-        ):
         world_size = 1
         runtime_rank = tensorrt_llm.mpi_rank()
         runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank)
@@ -212,7 +199,7 @@ class WhisperTRTLLM(object):
         self,
         audio: Union[str, np.ndarray, torch.Tensor],
         padding: int = 0,
-        return_duration = True
+        return_duration=True
     ):
         """
         Compute the log-Mel spectrogram of
@@ -242,8 +229,7 @@ class WhisperTRTLLM(object):
                     audio, _ = load_audio_wav_format(audio)
                 else:
                     audio = load_audio(audio)
-            assert isinstance(audio,
+            assert isinstance(audio, np.ndarray), f"Unsupported audio type: {type(audio)}"
-                            np.ndarray), f"Unsupported audio type: {type(audio)}"
             duration = audio.shape[-1] / SAMPLE_RATE
             audio = pad_or_trim(audio, N_SAMPLES)
             audio = audio.astype(np.float32)
@@ -254,14 +240,9 @@ class WhisperTRTLLM(object):
         if padding > 0:
             audio = F.pad(audio, (0, padding))
         window = torch.hann_window(N_FFT).to(audio.device)
-        stft = torch.stft(audio,
+        stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
-                        N_FFT,
-                        HOP_LENGTH,
-                        window=window,
-                        return_complex=True)
         magnitudes = stft[..., :-1].abs()**2
 
-        
         mel_spec = self.filters @ magnitudes
 
         log_spec = torch.clamp(mel_spec, min=1e-10).log10()
@@ -272,7 +253,6 @@ class WhisperTRTLLM(object):
         else:
             return log_spec
 
-
     def process_batch(
             self,
             mel,

whisper_live/utils.py

@@ -0,0 +1,71 @@
+import os
+import textwrap
+import scipy
+import ffmpeg
+import numpy as np
+
+
+def clear_screen():
+    """Clears the console screen."""
+    os.system("cls" if os.name == "nt" else "clear")
+
+
+def print_transcript(text):
+    """Prints formatted transcript text."""
+    wrapper = textwrap.TextWrapper(width=60)
+    for line in wrapper.wrap(text="".join(text)):
+        print(line)
+
+
+def format_time(s):
+    """Convert seconds (float) to SRT time format."""
+    hours = int(s // 3600)
+    minutes = int((s % 3600) // 60)
+    seconds = int(s % 60)
+    milliseconds = int((s - int(s)) * 1000)
+    return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
+
+
+def create_srt_file(segments, output_file):
+    with open(output_file, 'w', encoding='utf-8') as srt_file:
+        segment_number = 1
+        for segment in segments:
+            start_time = format_time(float(segment['start']))
+            end_time = format_time(float(segment['end']))
+            text = segment['text']
+
+            srt_file.write(f"{segment_number}\n")
+            srt_file.write(f"{start_time} --> {end_time}\n")
+            srt_file.write(f"{text}\n\n")
+
+            segment_number += 1
+
+
+def resample(file: str, sr: int = 16000):
+    """
+    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
+    Open an audio file and read as mono waveform, resampling as necessary,
+    save the resampled audio
+
+    Args:
+        file (str): The audio file to open
+        sr (int): The sample rate to resample the audio if necessary
+
+    Returns:
+        resampled_file (str): The resampled audio file
+    """
+    try:
+        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+        out, _ = (
+            ffmpeg.input(file, threads=0)
+            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
+            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+        )
+    except ffmpeg.Error as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+    np_buffer = np.frombuffer(out, dtype=np.int16)
+
+    resampled_file = f"{file.split('.')[0]}_resampled.wav"
+    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
+    return resampled_file

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]
 
@@ -38,7 +34,7 @@ class VoiceActivityDetection():
 
         if sr != 16000 and (sr % 16000 == 0):
             step = sr // 16000
-            x = x[:,::step]
+            x = x[:, ::step]
             sr = 16000
 
         if sr not in self.sample_rates:
@@ -110,9 +106,37 @@ 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:
                 print("Failed to download the model using wget.")
         return model_filename
+
+
+class VoiceActivityDetector:
+    def __init__(self, threshold=0.5, frame_rate=16000):
+        """
+        Initializes the VoiceActivityDetector with a voice activity detection model and a threshold.
+
+        Args:
+            threshold (float, optional): The probability threshold for detecting voice activity. Defaults to 0.5.
+        """
+        self.model = VoiceActivityDetection()
+        self.threshold = threshold
+        self.frame_rate = frame_rate
+
+    def __call__(self, audio_frame):
+        """
+        Determines if the given audio frame contains speech by comparing the detected speech probability against
+        the threshold.
+
+        Args:
+            audio_frame (np.ndarray): The audio frame to be analyzed for voice activity. It is expected to be a
+                                      NumPy array of audio samples.
+
+        Returns:
+            bool: True if the speech probability exceeds the threshold, indicating the presence of voice activity;
+                  False otherwise.
+        """
+        speech_prob = self.model(torch.from_numpy(audio_frame), self.frame_rate).item()
+        return speech_prob > self.threshold