bump version 0.3.0

fix: cpu usage issue.

.github/workflows/ci.yml

@@ -1,4 +1,4 @@
-name: CI
+name: Test & Build CI/CD
 
 on:
   push:
@@ -7,46 +7,163 @@ on:
     tags:
       - v*
   pull_request:
-    branches:
-      - main
+    branches: [ main ]
+    types: [opened, synchronize, reopened]
 
 jobs:
-  build-and-push-package:
-    runs-on: ubuntu-latest
+  run-tests:
+    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
-      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: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
 
-    - name: Install Client Requirements
-      run: pip install -r requirements/client.txt
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install flake8
 
-    - name: Install Server Requirements
-      run: pip install -r requirements/server.txt
+      - 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
 
-    - 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 }}
+  build-and-push-docker-cpu:
+    needs: [run-tests, check-code-format]
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+    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' && github.ref == 'refs/heads/main'
+    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: 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
+          pip install wheel
+
+      - 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 }}

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 } }, () => {
@@ -118,7 +126,8 @@ document.addEventListener("DOMContentLoaded", function () {
   function toggleCaptureButtons(isCapturing) {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
-    useServerCheckbox.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
-A nearly-live implementation of OpenAI's Whisper.
+# WhisperLive
 
-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
-- Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server)
-
-### 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.
+## 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 [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) for setup instructions.
 
 ## 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 whisper-live:latest
+  docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu: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 whisper-live:latest
+docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu: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 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 \
+ && 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 requirements/ /app
+COPY scripts/setup.sh requirements/server.txt /app/
 
-RUN bash setup.sh
-RUN pip install -r server.txt
+RUN apt update && bash setup.sh && 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,47 +1,33 @@
-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 \
- && 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
+    python3-dev \
+    python3-pip \
+    && python3 -m pip install --upgrade pip \
+    && rm -rf /var/lib/apt/lists/*
 
 # Create a working directory.
 RUN mkdir /app
 WORKDIR /app
 
-COPY scripts/setup.sh /app
-COPY requirements/ /app
+COPY scripts/setup.sh requirements/server.txt /app
 
-RUN apt update --fix-missing
-RUN bash setup.sh
-RUN pip install -r server.txt
+RUN apt update && bash setup.sh && rm setup.sh
+RUN pip install -r server.txt && rm server.txt
 
 COPY whisper_live /app/whisper_live
 
 COPY run_server.py /app
 
-CMD ["python", "run_server.py"]
+CMD ["python3", "run_server.py"]

requirements/server.txt

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

run_server.py

@@ -4,15 +4,15 @@ from whisper_live.server import TranscriptionServer
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--port', '-p',
-                        type=int, 
+                        type=int,
                         default=9090,
                         help="Websocket port to run the server on.")
     parser.add_argument('--backend', '-b',
-                        type=str, 
-                        default='faster_whisper', 
+                        type=str,
+                        default='faster_whisper',
                         help='Backends from ["tensorrt", "faster_whisper"]')
     parser.add_argument('--faster_whisper_custom_model_path', '-fw',
-                        type=str, default=None, 
+                        type=str, default=None,
                         help="Custom Faster Whisper Model")
     parser.add_argument('--trt_model_path', '-trt',
                         type=str,
@@ -30,7 +30,7 @@ if __name__ == "__main__":
     server = TranscriptionServer()
     server.run(
         "0.0.0.0",
-        port=args.port, 
+        port=args.port,
         backend=args.backend,
         faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
         whisper_tensorrt_path=args.trt_model_path,

setup.py

@@ -10,36 +10,38 @@ HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
 
 # This call to setup() does all the work
-setup(name="whisper-live",
-      version=__version__,
-      description="A nearly-live implementation of OpenAI's Whisper.",
-      long_description=README,
-      long_description_content_type="text/markdown",
-      include_package_data=True,
-      url="https://github.com/collabora/WhisperLive",
-      author="Collabora Ltd",
-      author_email="vineet.suryan@collabora.com",
-      license="MIT",
-      classifiers=[
-          "Development Status :: 4 - Beta",
-          "Intended Audience :: Developers",
-          "Intended Audience :: Science/Research",
-          "License :: OSI Approved :: MIT License",
-          "Programming Language :: Python :: 3",
-          "Programming Language :: Python :: 3 :: Only",
-          "Programming Language :: Python :: 3.8",
-          "Programming Language :: Python :: 3.9",
-          "Topic :: Scientific/Engineering :: Artificial Intelligence",
-      ],
-      packages=find_packages(
-        exclude=("examples",
-                 "Audio-Transcription-Chrome",
-                 "Audio-Transcription-Firefox",
-                 "requirements",
-                 "whisper-finetuning"
-                )
-      ),
-      install_requires=[
+setup(
+    name="whisper-live",
+    version=__version__,
+    description="A nearly-live implementation of OpenAI's Whisper.",
+    long_description=README,
+    long_description_content_type="text/markdown",
+    include_package_data=True,
+    url="https://github.com/collabora/WhisperLive",
+    author="Collabora Ltd",
+    author_email="vineet.suryan@collabora.com",
+    license="MIT",
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3 :: Only",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+    packages=find_packages(
+        exclude=(
+            "examples",
+            "Audio-Transcription-Chrome",
+            "Audio-Transcription-Firefox",
+            "requirements",
+            "whisper-finetuning"
+        )
+    ),
+    install_requires=[
         "PyAudio",
         "faster-whisper==0.10.0",
         "torch",
@@ -53,6 +55,6 @@ setup(name="whisper-live",
         "openai-whisper",
         "kaldialign",
         "soundfile",
-      ],
-      python_requires=">=3.8"
+    ],
+    python_requires=">=3.8"
 )

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.3.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
-
-
-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
+import ffmpeg
+import whisper_live.utils as utils
 
 
 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,10 +100,44 @@ 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.
-        
+
         It updates various attributes of the client based on the received message, including
         recording status, language detection, and server messages. If a disconnect message
         is received, it sets the recording status to False.
@@ -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.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
+            self.handle_status_messages(message)
             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,49 +176,24 @@ class Client:
             )
             return
 
-        if "segments" not in message.keys():
-            return
-
-        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)
+        if "segments" in message.keys():
+            self.process_segments(message["segments"])
 
     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):
         """
         Callback function called when the WebSocket connection is successfully opened.
-        
+
         Sends an initial configuration message to the server, including client UID,
         language selection, and task type.
 
@@ -257,6 +209,7 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
+                    "use_vad": self.use_vad
                 }
             )
         )
@@ -265,8 +218,8 @@ class Client:
     def bytes_to_float_array(audio_bytes):
         """
         Convert audio data from bytes to a NumPy float array.
-        
-        It assumes that the audio data is in 16-bit PCM format. The audio data is normalized to 
+
+        It assumes that the audio data is in 16-bit PCM format. The audio data is normalized to
         have values between -1 and 1.
 
         Args:
@@ -294,10 +247,10 @@ class Client:
     def play_file(self, filename):
         """
         Play an audio file and send it to the server for processing.
-        
+
         Reads an audio file, plays it through the audio output, and simultaneously sends
-        the audio data to the server for processing. It uses PyAudio to create an audio 
-        stream for playback. The audio data is read from the file in chunks, converted to 
+        the audio data to the server for processing. It uses PyAudio to create an audio
+        stream for playback. The audio data is read from the file in chunks, converted to
         floating-point format, and sent to the server using WebSocket communication.
         This method is typically used when you want to process pre-recorded audio and send it
         to the server in real-time.
@@ -305,7 +258,7 @@ class Client:
         Args:
             filename (str): The path to the audio file to be played and sent to the server.
         """
-        
+
         # read audio and create pyaudio stream
         with wave.open(filename, "rb") as wavfile:
             self.stream = self.p.open(
@@ -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()
@@ -351,7 +304,7 @@ class Client:
         """
         Close the WebSocket connection and join the WebSocket thread.
 
-        First attempts to close the WebSocket connection using `self.client_socket.close()`. After 
+        First attempts to close the WebSocket connection using `self.client_socket.close()`. After
         closing the connection, it joins the WebSocket thread to ensure proper termination.
 
         """
@@ -378,7 +331,7 @@ class Client:
         """
         Write audio frames to a WAV file.
 
-        The WAV file is created or overwritten with the specified name. The audio frames should be 
+        The WAV file is created or overwritten with the specified name. The audio frames should be
         in the correct format and match the specified channel, sample width, and sample rate.
 
         Args:
@@ -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.
@@ -439,11 +391,12 @@ class Client:
 
         Audio data is saved in chunks to the "chunks" directory. Each chunk is saved as a separate WAV file.
         The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`.
-        The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording, 
+        The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording,
         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)
@@ -493,8 +446,8 @@ class Client:
     def write_output_recording(self, n_audio_file, out_file):
         """
         Combine and save recorded audio chunks into a single WAV file.
-        
-        The individual audio chunk files are expected to be located in the "chunks" directory. Reads each chunk 
+
+        The individual audio chunk files are expected to be located in the "chunks" directory. Reads each chunk
         file, appends its audio data to the final recording, and then deletes the chunk file. After combining
         and saving, the final recording is stored in the specified `out_file`.
 
@@ -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,26 +506,20 @@ class TranscriptionClient:
         transcription_client()
         ```
     """
-    def __init__(self,
-        host,
-        port,
-        lang=None,
-        translate=False,
-        model="small",
-    ):
-        self.client = Client(host, port, lang, translate, model)
+    def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True):
+        self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
 
     def __call__(self, audio=None, hls_url=None):
         """
         Start the transcription process.
 
         Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
-        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it 
+        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it
         will be played and streamed to the server; otherwise, it will perform live recording.
 
         Args:
             audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
-                   
+
         """
         print("[INFO]: Waiting for server ready ...")
         while not self.client.recording:
@@ -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()
+            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]],
@@ -362,4 +362,4 @@ def write_error_stats(
         hyp_count = corr + hyp_sub + ins
 
         print(f"{word}   {corr} {tot_errs} {ref_count} {hyp_count}", file=f)
-    return float(tot_err_rate)
+    return float(tot_err_rate)

whisper_live/transcriber.py

@@ -180,7 +180,7 @@ class WhisperModel:
 
         return config
 
-    def transcribe(
+    def transcribe(                                                         # noqa: C901
         self,
         audio: Union[str, BinaryIO, np.ndarray],
         language: Optional[str] = None,
@@ -315,6 +315,9 @@ class WhisperModel:
         else:
             speech_chunks = None
 
+        if audio.shape[0] == 0:
+            return None, None
+
         features = self.feature_extractor(audio)
 
         encoder_output = None
@@ -400,7 +403,7 @@ class WhisperModel:
 
         return segments, info
 
-    def generate_segments(
+    def generate_segments(                                                  # noqa: C901
         self,
         features: np.ndarray,
         tokenizer: Tokenizer,
@@ -425,7 +428,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 +752,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 +769,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,
-                           write_error_stats, load_audio_wav_format,
-                           pad_or_trim, load_audio)
+from whisper_live.tensorrt_utils import (mel_filters, 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__(
-        self,
-        engine_dir,
-        debug_mode=False,
-        assets_dir=None,
-        device=None,
-        is_multilingual=False,
-        language="en",
-        task="transcribe"
-        ):
+    def __init__(self, engine_dir, 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,
-                            np.ndarray), f"Unsupported audio type: {type(audio)}"
+            assert isinstance(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,
-                        N_FFT,
-                        HOP_LENGTH,
-                        window=window,
-                        return_complex=True)
+        stft = torch.stft(audio, 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,
@@ -296,7 +276,7 @@ class WhisperTRTLLM(object):
             text = self.tokenizer.decode(output_ids[i][0]).strip()
             texts.append(text)
         return texts
-    
+
     def transcribe(
             self,
             mel,
@@ -336,5 +316,5 @@ def decode_wav_file(
     prediction = re.sub(r'<\|.*?\|>', '', prediction)
     if normalizer:
         prediction = normalizer(prediction)
-    
+
     return prediction.strip()

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
+        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