Compare commits
72 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bc070d6688 | |||
| 8e7e329a39 | |||
| 380f07394b | |||
| 30f78a2cc6 | |||
| 01c6bc1ecd | |||
| bdaed45820 | |||
| 4870e9fb9e | |||
| ccb183b4d8 | |||
| fac62aaccc | |||
| aade67736a | |||
| abfe830eee | |||
| cb392cbb93 | |||
| 42733da59a | |||
| 26c517021f | |||
| cf721e8b53 | |||
| 5985ec82b6 | |||
| 2f1c934ea2 | |||
| b220ccb330 | |||
| 5e3906fc7b | |||
| a8b9275013 | |||
| 815441e8bb | |||
| 5b9bc2bc0e | |||
| ee132517fa | |||
| 761bb61e87 | |||
| 14077315ae | |||
| ab17c4dbc6 | |||
| 5e2421118d | |||
| d1de2ec3ce | |||
| 22a37e7843 | |||
| e4579ef291 | |||
| f73a146eb9 | |||
| cfba5b3e54 | |||
| 1ac7a278bb | |||
| 3a96f60006 | |||
| 3c09289dea | |||
| e1a42c22d2 | |||
| 3d043dc906 | |||
| 399e9e7efe | |||
| 9d2ea75247 | |||
| 225a98be0c | |||
| 8f373c3537 | |||
| 61d07edabb | |||
| 03e30e1fed | |||
| c0a947a8f6 | |||
| 819ab35b28 | |||
| 615c9c7aed | |||
| a9683319e0 | |||
| 0a2d92c5b8 | |||
| dccfce2a3c | |||
| f78fc473c5 | |||
| 0dfbdb2477 | |||
| e171b9c460 | |||
| 66b5dc7c15 | |||
| 0f1d36fc06 | |||
| d24c53198c | |||
| fe1640695c | |||
| 8d77f0fa5a | |||
| 2e37216282 | |||
| 7c7a446478 | |||
| 37d7f2ed66 | |||
| 754f22dfae | |||
| ebd2dc9568 | |||
| 9b2e17ec4d | |||
| 5b32dc4130 | |||
| c0f37c77e9 | |||
| 3b15dc76b4 | |||
| 4d477e35e7 | |||
| a17f4041de | |||
| 8a06ba802b | |||
| 02d4566289 | |||
| acd4902bec | |||
| a495a49b06 |
@@ -77,7 +77,7 @@ jobs:
|
|||||||
build-and-push-docker-cpu:
|
build-and-push-docker-cpu:
|
||||||
needs: [run-tests, check-code-format]
|
needs: [run-tests, check-code-format]
|
||||||
runs-on: ubuntu-22.04
|
runs-on: ubuntu-22.04
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v2
|
- uses: actions/checkout@v2
|
||||||
|
|
||||||
@@ -99,11 +99,40 @@ jobs:
|
|||||||
push: true
|
push: true
|
||||||
tags: ghcr.io/collabora/whisperlive-cpu:latest
|
tags: ghcr.io/collabora/whisperlive-cpu:latest
|
||||||
|
|
||||||
|
build-and-push-docker-tensorrt:
|
||||||
|
needs: [run-tests, check-code-format]
|
||||||
|
timeout-minutes: 20
|
||||||
|
runs-on: ubuntu-22.04
|
||||||
|
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || 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.tensorrt
|
||||||
|
push: true
|
||||||
|
tags: ghcr.io/collabora/whisperlive-tensorrt:latest
|
||||||
|
|
||||||
build-and-push-docker-gpu:
|
build-and-push-docker-gpu:
|
||||||
needs: [run-tests, check-code-format, build-and-push-docker-cpu]
|
needs: [run-tests, check-code-format, build-and-push-docker-cpu]
|
||||||
timeout-minutes: 20
|
timeout-minutes: 20
|
||||||
runs-on: ubuntu-22.04
|
runs-on: ubuntu-22.04
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v2
|
- uses: actions/checkout@v2
|
||||||
|
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ python3 run_server.py --port 9090 \
|
|||||||
|
|
||||||
# running with custom model
|
# running with custom model
|
||||||
python3 run_server.py --port 9090 \
|
python3 run_server.py --port 9090 \
|
||||||
--backend faster_whisper
|
--backend faster_whisper \
|
||||||
-fw "/path/to/custom/faster/whisper/model"
|
-fw "/path/to/custom/faster/whisper/model"
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -53,10 +53,30 @@ python3 run_server.py -p 9090 \
|
|||||||
-trt /home/TensorRT-LLM/examples/whisper/whisper_small \
|
-trt /home/TensorRT-LLM/examples/whisper/whisper_small \
|
||||||
-m
|
-m
|
||||||
```
|
```
|
||||||
|
#### Controlling OpenMP Threads
|
||||||
|
To control the number of threads used by OpenMP, you can set the `OMP_NUM_THREADS` environment variable. This is useful for managing CPU resources and ensuring consistent performance. If not specified, `OMP_NUM_THREADS` is set to `1` by default. You can change this by using the `--omp_num_threads` argument:
|
||||||
|
```bash
|
||||||
|
python3 run_server.py --port 9090 \
|
||||||
|
--backend faster_whisper \
|
||||||
|
--omp_num_threads 4
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Single model mode
|
||||||
|
By default, when running the server without specifying a model, the server will instantiate a new whisper model for every client connection. This has the advantage, that the server can use different model sizes, based on the client's requested model size. On the other hand, it also means you have to wait for the model to be loaded upon client connection and you will have increased (V)RAM usage.
|
||||||
|
|
||||||
|
When serving a custom TensorRT model using the `-trt` or a custom faster_whisper model using the `-fw` option, the server will instead only instantiate the custom model once and then reuse it for all client connections.
|
||||||
|
|
||||||
|
If you don't want this, set `--no_single_model`.
|
||||||
|
|
||||||
|
|
||||||
### Running the Client
|
### Running the Client
|
||||||
- Initializing the client:
|
- Initializing the client with below parameters:
|
||||||
|
- `lang`: Language of the input audio, applicable only if using a multilingual model.
|
||||||
|
- `translate`: If set to `True` then translate from any language to `en`.
|
||||||
|
- `model`: Whisper model size.
|
||||||
|
- `use_vad`: Whether to use `Voice Activity Detection` on the server.
|
||||||
|
- `save_output_recording`: Set to True to save the microphone input as a `.wav` file during live transcription. This option is helpful for recording sessions for later playback or analysis. Defaults to `False`.
|
||||||
|
- `output_recording_filename`: Specifies the `.wav` file path where the microphone input will be saved if `save_output_recording` is set to `True`.
|
||||||
```python
|
```python
|
||||||
from whisper_live.client import TranscriptionClient
|
from whisper_live.client import TranscriptionClient
|
||||||
client = TranscriptionClient(
|
client = TranscriptionClient(
|
||||||
@@ -66,11 +86,13 @@ client = TranscriptionClient(
|
|||||||
translate=False,
|
translate=False,
|
||||||
model="small",
|
model="small",
|
||||||
use_vad=False,
|
use_vad=False,
|
||||||
|
save_output_recording=True, # Only used for microphone input, False by Default
|
||||||
|
output_recording_filename="./output_recording.wav" # Only used for microphone input
|
||||||
)
|
)
|
||||||
```
|
```
|
||||||
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.
|
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:
|
- Transcribe an audio file:
|
||||||
```python
|
```python
|
||||||
client("tests/jfk.wav")
|
client("tests/jfk.wav")
|
||||||
```
|
```
|
||||||
@@ -80,9 +102,14 @@ client("tests/jfk.wav")
|
|||||||
client()
|
client()
|
||||||
```
|
```
|
||||||
|
|
||||||
|
- To transcribe from a RTSP stream:
|
||||||
|
```python
|
||||||
|
client(rtsp_url="rtsp://admin:admin@192.168.0.1/rtsp")
|
||||||
|
```
|
||||||
|
|
||||||
- To transcribe from a HLS stream:
|
- To transcribe from a HLS stream:
|
||||||
```python
|
```python
|
||||||
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")
|
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")
|
||||||
```
|
```
|
||||||
|
|
||||||
## Browser Extensions
|
## Browser Extensions
|
||||||
@@ -96,7 +123,18 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
|
|||||||
docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu: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.
|
- TensorRT.
|
||||||
|
```bash
|
||||||
|
docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
|
||||||
|
|
||||||
|
# Build tiny.en engine
|
||||||
|
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
|
||||||
|
|
||||||
|
# Run server with tiny.en
|
||||||
|
python3 run_server.py --port 9090 \
|
||||||
|
--backend tensorrt \
|
||||||
|
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
|
||||||
|
```
|
||||||
|
|
||||||
- CPU
|
- CPU
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
+9
-34
@@ -1,6 +1,6 @@
|
|||||||
# Whisper-TensorRT
|
# WhisperLive-TensorRT
|
||||||
We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
|
We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
|
||||||
**Note**: We use [our fork to setup TensorRT](https://github.com/makaveli10/TensorRT-LLM)
|
**Note**: We use `tensorrt_llm==0.9.0`
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
- Install [docker](https://docs.docker.com/engine/install/)
|
- Install [docker](https://docs.docker.com/engine/install/)
|
||||||
@@ -12,56 +12,31 @@ git clone https://github.com/collabora/WhisperLive.git
|
|||||||
cd WhisperLive
|
cd WhisperLive
|
||||||
```
|
```
|
||||||
|
|
||||||
- Pull the TensorRT-LLM docker image which we prebuilt for WhisperLive TensorRT backend.
|
- Run WhisperLive TensorRT in docker
|
||||||
```bash
|
```bash
|
||||||
docker pull ghcr.io/collabora/whisperbot-base:latest
|
docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
|
||||||
```
|
```
|
||||||
|
|
||||||
- Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
|
|
||||||
```bash
|
|
||||||
docker run -it --gpus all --shm-size=8g \
|
|
||||||
--ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
|
|
||||||
-p 9090:9090 -v /path/to/WhisperLive:/home/WhisperLive \
|
|
||||||
ghcr.io/collabora/whisperbot-base:latest
|
|
||||||
```
|
|
||||||
|
|
||||||
- Make sure to test the installation.
|
|
||||||
```bash
|
|
||||||
# export ENV=${ENV:-/etc/shinit_v2}
|
|
||||||
# source $ENV
|
|
||||||
python -c "import torch; import tensorrt; import tensorrt_llm"
|
|
||||||
```
|
|
||||||
**NOTE**: Uncomment and update library paths if imports fail.
|
|
||||||
|
|
||||||
## Whisper TensorRT Engine
|
## Whisper TensorRT Engine
|
||||||
- We build `small.en` and `small` multilingual TensorRT engine. The script logs the path of the directory with Whisper TensorRT engine. We need the model_path to run the server.
|
- We build `small.en` and `small` multilingual TensorRT engine as examples below. The script logs the path of the directory with Whisper TensorRT engine. We need that model_path to run the server.
|
||||||
```bash
|
```bash
|
||||||
# convert small.en
|
# convert small.en
|
||||||
bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small.en
|
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
|
||||||
|
|
||||||
# convert small multilingual model
|
# convert small multilingual model
|
||||||
bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small
|
bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
|
||||||
```
|
```
|
||||||
|
|
||||||
## Run WhisperLive Server with TensorRT Backend
|
## Run WhisperLive Server with TensorRT Backend
|
||||||
```bash
|
```bash
|
||||||
cd /home/WhisperLive
|
|
||||||
|
|
||||||
# Install requirements
|
|
||||||
apt update && bash scripts/setup.sh
|
|
||||||
pip install -r requirements/server.txt
|
|
||||||
|
|
||||||
# Required to create mel spectogram
|
|
||||||
wget --directory-prefix=assets assets/mel_filters.npz https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
|
|
||||||
|
|
||||||
# Run English only model
|
# Run English only model
|
||||||
python3 run_server.py --port 9090 \
|
python3 run_server.py --port 9090 \
|
||||||
--backend tensorrt \
|
--backend tensorrt \
|
||||||
--trt_model_path "path/to/whisper_trt/from/build/step"
|
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
|
||||||
|
|
||||||
# Run Multilingual model
|
# Run Multilingual model
|
||||||
python3 run_server.py --port 9090 \
|
python3 run_server.py --port 9090 \
|
||||||
--backend tensorrt \
|
--backend tensorrt \
|
||||||
--trt_model_path "path/to/whisper_trt/from/build/step" \
|
--trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
|
||||||
--trt_multilingual
|
--trt_multilingual
|
||||||
```
|
```
|
||||||
|
|||||||
+13
-12
@@ -1,22 +1,23 @@
|
|||||||
FROM python:3.8-slim-buster
|
FROM python:3.10-bookworm
|
||||||
|
|
||||||
ARG DEBIAN_FRONTEND=noninteractive
|
ARG DEBIAN_FRONTEND=noninteractive
|
||||||
|
|
||||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
# install lib required for pyaudio
|
||||||
curl \
|
RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||||
ca-certificates \
|
|
||||||
sudo \
|
|
||||||
git \
|
|
||||||
bzip2 \
|
|
||||||
libx11-6 \
|
|
||||||
&& apt-get clean \
|
|
||||||
&& rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
|
# update pip to support for whl.metadata -> less downloading
|
||||||
|
RUN pip install --no-cache-dir -U "pip>=24"
|
||||||
|
|
||||||
|
# create a working directory
|
||||||
|
RUN mkdir /app
|
||||||
WORKDIR /app
|
WORKDIR /app
|
||||||
|
|
||||||
COPY scripts/setup.sh requirements/server.txt /app/
|
# install pytorch, but without the nvidia-libs that are only necessary for gpu
|
||||||
|
RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
|
||||||
|
|
||||||
RUN apt update && bash setup.sh && pip install -r server.txt
|
# install the requirements for running the whisper-live server
|
||||||
|
COPY requirements/server.txt /app/
|
||||||
|
RUN pip install --no-cache-dir -r server.txt && rm server.txt
|
||||||
|
|
||||||
COPY whisper_live /app/whisper_live
|
COPY whisper_live /app/whisper_live
|
||||||
COPY run_server.py /app
|
COPY run_server.py /app
|
||||||
|
|||||||
+14
-21
@@ -1,33 +1,26 @@
|
|||||||
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
|
FROM python:3.10-bookworm
|
||||||
|
|
||||||
ARG DEBIAN_FRONTEND=noninteractive
|
ARG DEBIAN_FRONTEND=noninteractive
|
||||||
|
|
||||||
# Remove any third-party apt sources to avoid issues with expiring keys.
|
# install lib required for pyaudio
|
||||||
RUN rm -f /etc/apt/sources.list.d/*.list
|
RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
# Install some basic utilities.
|
# update pip to support for whl.metadata -> less downloading
|
||||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
RUN pip install --no-cache-dir -U "pip>=24"
|
||||||
curl \
|
|
||||||
ca-certificates \
|
|
||||||
sudo \
|
|
||||||
git \
|
|
||||||
bzip2 \
|
|
||||||
libx11-6 \
|
|
||||||
python3-dev \
|
|
||||||
python3-pip \
|
|
||||||
&& python3 -m pip install --upgrade pip \
|
|
||||||
&& rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
# Create a working directory.
|
# create a working directory
|
||||||
RUN mkdir /app
|
RUN mkdir /app
|
||||||
WORKDIR /app
|
WORKDIR /app
|
||||||
|
|
||||||
COPY scripts/setup.sh requirements/server.txt /app
|
# install the requirements for running the whisper-live server
|
||||||
|
COPY requirements/server.txt /app/
|
||||||
|
RUN pip install --no-cache-dir -r server.txt && rm server.txt
|
||||||
|
|
||||||
RUN apt update && bash setup.sh && rm setup.sh
|
# make the paths of the nvidia libs installed as wheels visible. equivalent to:
|
||||||
RUN pip install -r server.txt && rm server.txt
|
# export LD_LIBRARY_PATH=`python3 -c 'import os; import nvidia.cublas.lib; import nvidia.cudnn.lib; print(os.path.dirname(nvidia.cublas.lib.__file__) + ":" + os.path.dirname(nvidia.cudnn.lib.__file__))'`
|
||||||
|
ENV LD_LIBRARY_PATH="/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib:/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib"
|
||||||
|
|
||||||
COPY whisper_live /app/whisper_live
|
COPY whisper_live /app/whisper_live
|
||||||
|
|
||||||
COPY run_server.py /app
|
COPY run_server.py /app
|
||||||
|
|
||||||
CMD ["python3", "run_server.py"]
|
CMD ["python", "run_server.py"]
|
||||||
|
|||||||
@@ -0,0 +1,28 @@
|
|||||||
|
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
|
||||||
|
|
||||||
|
ARG DEBIAN_FRONTEND=noninteractive
|
||||||
|
|
||||||
|
RUN apt-get update && apt-get install -y \
|
||||||
|
python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
|
||||||
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
|
||||||
|
mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
|
||||||
|
rm -rf TensorRT-LLM
|
||||||
|
|
||||||
|
COPY assets/ ./assets
|
||||||
|
RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
|
||||||
|
|
||||||
|
COPY scripts/setup.sh ./
|
||||||
|
RUN apt update && bash setup.sh && rm setup.sh
|
||||||
|
|
||||||
|
COPY requirements/server.txt .
|
||||||
|
RUN pip install --no-cache-dir -r server.txt && rm server.txt
|
||||||
|
|
||||||
|
COPY whisper_live ./whisper_live
|
||||||
|
COPY scripts/build_whisper_tensorrt.sh .
|
||||||
|
COPY run_server.py .
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
faster-whisper==0.10.0
|
faster-whisper==1.0.1
|
||||||
torch
|
torch==2.3.0
|
||||||
websockets
|
websockets
|
||||||
onnxruntime==1.16.0
|
onnxruntime==1.16.0
|
||||||
numba
|
numba
|
||||||
@@ -9,4 +9,5 @@ soundfile
|
|||||||
ffmpeg-python
|
ffmpeg-python
|
||||||
scipy
|
scipy
|
||||||
jiwer
|
jiwer
|
||||||
evaluate
|
evaluate
|
||||||
|
numpy<2
|
||||||
+14
-2
@@ -1,5 +1,5 @@
|
|||||||
import argparse
|
import argparse
|
||||||
from whisper_live.server import TranscriptionServer
|
import os
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
@@ -21,12 +21,23 @@ if __name__ == "__main__":
|
|||||||
parser.add_argument('--trt_multilingual', '-m',
|
parser.add_argument('--trt_multilingual', '-m',
|
||||||
action="store_true",
|
action="store_true",
|
||||||
help='Boolean only for TensorRT model. True if multilingual.')
|
help='Boolean only for TensorRT model. True if multilingual.')
|
||||||
|
parser.add_argument('--omp_num_threads', '-omp',
|
||||||
|
type=int,
|
||||||
|
default=1,
|
||||||
|
help="Number of threads to use for OpenMP")
|
||||||
|
parser.add_argument('--no_single_model', '-nsm',
|
||||||
|
action='store_true',
|
||||||
|
help='Set this if every connection should instantiate its own model. Only relevant for custom model, passed using -trt or -fw.')
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.backend == "tensorrt":
|
if args.backend == "tensorrt":
|
||||||
if args.trt_model_path is None:
|
if args.trt_model_path is None:
|
||||||
raise ValueError("Please Provide a valid tensorrt model path")
|
raise ValueError("Please Provide a valid tensorrt model path")
|
||||||
|
|
||||||
|
if "OMP_NUM_THREADS" not in os.environ:
|
||||||
|
os.environ["OMP_NUM_THREADS"] = str(args.omp_num_threads)
|
||||||
|
|
||||||
|
from whisper_live.server import TranscriptionServer
|
||||||
server = TranscriptionServer()
|
server = TranscriptionServer()
|
||||||
server.run(
|
server.run(
|
||||||
"0.0.0.0",
|
"0.0.0.0",
|
||||||
@@ -34,5 +45,6 @@ if __name__ == "__main__":
|
|||||||
backend=args.backend,
|
backend=args.backend,
|
||||||
faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
|
faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
|
||||||
whisper_tensorrt_path=args.trt_model_path,
|
whisper_tensorrt_path=args.trt_model_path,
|
||||||
trt_multilingual=args.trt_multilingual
|
trt_multilingual=args.trt_multilingual,
|
||||||
|
single_model=not args.no_single_model,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ download_and_build_model() {
|
|||||||
local output_dir="whisper_${model_name//./_}"
|
local output_dir="whisper_${model_name//./_}"
|
||||||
echo "$output_dir"
|
echo "$output_dir"
|
||||||
echo "Running build script for $model_name with output directory $output_dir"
|
echo "Running build script for $model_name with output directory $output_dir"
|
||||||
python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --model_name "$model_name"
|
python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --enable_context_fmha --model_name "$model_name"
|
||||||
echo "Whisper $model_name TensorRT engine built."
|
echo "Whisper $model_name TensorRT engine built."
|
||||||
echo "========================================="
|
echo "========================================="
|
||||||
echo "Model is located at: $(pwd)/$output_dir"
|
echo "Model is located at: $(pwd)/$output_dir"
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ setup(
|
|||||||
),
|
),
|
||||||
install_requires=[
|
install_requires=[
|
||||||
"PyAudio",
|
"PyAudio",
|
||||||
"faster-whisper==0.10.0",
|
"faster-whisper==1.0.1",
|
||||||
"torch",
|
"torch",
|
||||||
"torchaudio",
|
"torchaudio",
|
||||||
"websockets",
|
"websockets",
|
||||||
|
|||||||
+50
-5
@@ -2,10 +2,12 @@ import json
|
|||||||
import os
|
import os
|
||||||
import scipy
|
import scipy
|
||||||
import websocket
|
import websocket
|
||||||
|
import copy
|
||||||
import unittest
|
import unittest
|
||||||
from unittest.mock import patch, MagicMock
|
from unittest.mock import patch, MagicMock
|
||||||
from whisper_live.client import TranscriptionClient
|
from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
|
||||||
from whisper_live.utils import resample
|
from whisper_live.utils import resample
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
class BaseTestCase(unittest.TestCase):
|
class BaseTestCase(unittest.TestCase):
|
||||||
@@ -24,6 +26,7 @@ class BaseTestCase(unittest.TestCase):
|
|||||||
|
|
||||||
self.mock_pyaudio = mock_pyaudio
|
self.mock_pyaudio = mock_pyaudio
|
||||||
self.mock_websocket = mock_websocket
|
self.mock_websocket = mock_websocket
|
||||||
|
self.mock_audio_packet = b'\x00\x01\x02\x03'
|
||||||
|
|
||||||
def tearDown(self):
|
def tearDown(self):
|
||||||
self.client.close_websocket()
|
self.client.close_websocket()
|
||||||
@@ -31,7 +34,6 @@ class BaseTestCase(unittest.TestCase):
|
|||||||
self.mock_websocket.stop()
|
self.mock_websocket.stop()
|
||||||
del self.client
|
del self.client
|
||||||
|
|
||||||
|
|
||||||
class TestClientWebSocketCommunication(BaseTestCase):
|
class TestClientWebSocketCommunication(BaseTestCase):
|
||||||
def test_websocket_communication(self):
|
def test_websocket_communication(self):
|
||||||
expected_url = 'ws://localhost:9090'
|
expected_url = 'ws://localhost:9090'
|
||||||
@@ -106,6 +108,49 @@ class TestAudioResampling(unittest.TestCase):
|
|||||||
|
|
||||||
class TestSendingAudioPacket(BaseTestCase):
|
class TestSendingAudioPacket(BaseTestCase):
|
||||||
def test_send_packet(self):
|
def test_send_packet(self):
|
||||||
mock_audio_packet = b'\x00\x01\x02\x03'
|
self.client.send_packet_to_server(self.mock_audio_packet)
|
||||||
self.client.send_packet_to_server(mock_audio_packet)
|
self.client.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
|
||||||
self.client.client_socket.send.assert_called_with(mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
|
|
||||||
|
class TestTee(BaseTestCase):
|
||||||
|
@patch('whisper_live.client.websocket.WebSocketApp')
|
||||||
|
@patch('whisper_live.client.pyaudio.PyAudio')
|
||||||
|
def setUp(self, mock_audio, mock_websocket):
|
||||||
|
super().setUp()
|
||||||
|
self.client2 = Client(host='localhost', port=9090, lang="es", translate=False, srt_file_path="transcript.srt")
|
||||||
|
self.client3 = Client(host='localhost', port=9090, lang="es", translate=True, srt_file_path="translation.srt")
|
||||||
|
# need a separate mock for each websocket
|
||||||
|
self.client3.client_socket = copy.deepcopy(self.client3.client_socket)
|
||||||
|
self.tee = TranscriptionTeeClient([self.client2, self.client3])
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
self.tee.close_all_clients()
|
||||||
|
del self.tee
|
||||||
|
super().tearDown()
|
||||||
|
|
||||||
|
def test_invalid_constructor(self):
|
||||||
|
with self.assertRaises(Exception) as context:
|
||||||
|
TranscriptionTeeClient([])
|
||||||
|
|
||||||
|
def test_multicast_unconditional(self):
|
||||||
|
self.tee.multicast_packet(self.mock_audio_packet, True)
|
||||||
|
for client in self.tee.clients:
|
||||||
|
client.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
|
||||||
|
|
||||||
|
def test_multicast_conditional(self):
|
||||||
|
self.client2.recording = False
|
||||||
|
self.client3.recording = True
|
||||||
|
self.tee.multicast_packet(self.mock_audio_packet, False)
|
||||||
|
self.client2.client_socket.send.assert_not_called()
|
||||||
|
self.client3.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
|
||||||
|
|
||||||
|
def test_close_all(self):
|
||||||
|
self.tee.close_all_clients()
|
||||||
|
for client in self.tee.clients:
|
||||||
|
client.client_socket.close.assert_called()
|
||||||
|
|
||||||
|
def test_write_all_srt(self):
|
||||||
|
for client in self.tee.clients:
|
||||||
|
client.server_backend = "faster_whisper"
|
||||||
|
self.tee.write_all_clients_srt()
|
||||||
|
self.assertTrue(Path("transcript.srt").is_file())
|
||||||
|
self.assertTrue(Path("translation.srt").is_file())
|
||||||
|
|||||||
+25
-12
@@ -9,7 +9,7 @@ import evaluate
|
|||||||
|
|
||||||
from websockets.exceptions import ConnectionClosed
|
from websockets.exceptions import ConnectionClosed
|
||||||
from whisper_live.server import TranscriptionServer
|
from whisper_live.server import TranscriptionServer
|
||||||
from whisper_live.client import TranscriptionClient
|
from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
|
||||||
from whisper.normalizers import EnglishTextNormalizer
|
from whisper.normalizers import EnglishTextNormalizer
|
||||||
|
|
||||||
|
|
||||||
@@ -69,6 +69,10 @@ class TestServerConnection(unittest.TestCase):
|
|||||||
class TestServerInferenceAccuracy(unittest.TestCase):
|
class TestServerInferenceAccuracy(unittest.TestCase):
|
||||||
@classmethod
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
|
cls.mock_pyaudio_patch = mock.patch('pyaudio.PyAudio')
|
||||||
|
cls.mock_pyaudio = cls.mock_pyaudio_patch.start()
|
||||||
|
cls.mock_pyaudio.return_value.open.return_value = mock.MagicMock()
|
||||||
|
|
||||||
cls.server_process = subprocess.Popen(["python", "run_server.py"])
|
cls.server_process = subprocess.Popen(["python", "run_server.py"])
|
||||||
time.sleep(2)
|
time.sleep(2)
|
||||||
|
|
||||||
@@ -77,21 +81,13 @@ class TestServerInferenceAccuracy(unittest.TestCase):
|
|||||||
cls.server_process.terminate()
|
cls.server_process.terminate()
|
||||||
cls.server_process.wait()
|
cls.server_process.wait()
|
||||||
|
|
||||||
@mock.patch('pyaudio.PyAudio')
|
def setUp(self):
|
||||||
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.metric = evaluate.load("wer")
|
||||||
self.normalizer = EnglishTextNormalizer()
|
self.normalizer = EnglishTextNormalizer()
|
||||||
self.client = TranscriptionClient(
|
|
||||||
"localhost", "9090", model="base.en", lang="en",
|
|
||||||
)
|
|
||||||
|
|
||||||
def test_inference(self):
|
def check_prediction(self, srt_path):
|
||||||
gt = "And so my fellow Americans, ask not, what your country can do for you. Ask what you can do for your country!"
|
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(srt_path, "r") as f:
|
||||||
with open("output.srt", "r") as f:
|
|
||||||
lines = f.readlines()
|
lines = f.readlines()
|
||||||
prediction = " ".join([line.strip() for line in lines[2::4]])
|
prediction = " ".join([line.strip() for line in lines[2::4]])
|
||||||
prediction_normalized = self.normalizer(prediction)
|
prediction_normalized = self.normalizer(prediction)
|
||||||
@@ -104,6 +100,23 @@ class TestServerInferenceAccuracy(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
self.assertLess(wer, 0.05)
|
self.assertLess(wer, 0.05)
|
||||||
|
|
||||||
|
def test_inference(self):
|
||||||
|
client = TranscriptionClient(
|
||||||
|
"localhost", "9090", model="base.en", lang="en",
|
||||||
|
)
|
||||||
|
client("assets/jfk.flac")
|
||||||
|
self.check_prediction("output.srt")
|
||||||
|
|
||||||
|
def test_simultaneous_inference(self):
|
||||||
|
client1 = Client(
|
||||||
|
"localhost", "9090", model="base.en", lang="en", srt_file_path="transcript1.srt")
|
||||||
|
client2 = Client(
|
||||||
|
"localhost", "9090", model="base.en", lang="en", srt_file_path="transcript2.srt")
|
||||||
|
tee = TranscriptionTeeClient([client1, client2])
|
||||||
|
tee("assets/jfk.flac")
|
||||||
|
self.check_prediction("transcript1.srt")
|
||||||
|
self.check_prediction("transcript2.srt")
|
||||||
|
|
||||||
|
|
||||||
class TestExceptionHandling(unittest.TestCase):
|
class TestExceptionHandling(unittest.TestCase):
|
||||||
def setUp(self):
|
def setUp(self):
|
||||||
|
|||||||
@@ -1 +1 @@
|
|||||||
__version__ = "0.3.0"
|
__version__ = "0.5.1"
|
||||||
|
|||||||
+393
-235
@@ -1,6 +1,8 @@
|
|||||||
import os
|
import os
|
||||||
|
import shutil
|
||||||
import wave
|
import wave
|
||||||
|
|
||||||
|
import logging
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pyaudio
|
import pyaudio
|
||||||
import threading
|
import threading
|
||||||
@@ -14,7 +16,7 @@ import whisper_live.utils as utils
|
|||||||
|
|
||||||
class Client:
|
class Client:
|
||||||
"""
|
"""
|
||||||
Handles audio recording, streaming, and communication with a server using WebSocket.
|
Handles communication with a server using WebSocket.
|
||||||
"""
|
"""
|
||||||
INSTANCES = {}
|
INSTANCES = {}
|
||||||
END_OF_AUDIO = "END_OF_AUDIO"
|
END_OF_AUDIO = "END_OF_AUDIO"
|
||||||
@@ -27,7 +29,8 @@ class Client:
|
|||||||
translate=False,
|
translate=False,
|
||||||
model="small",
|
model="small",
|
||||||
srt_file_path="output.srt",
|
srt_file_path="output.srt",
|
||||||
use_vad=True
|
use_vad=True,
|
||||||
|
log_transcription=True
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initializes a Client instance for audio recording and streaming to a server.
|
Initializes a Client instance for audio recording and streaming to a server.
|
||||||
@@ -42,37 +45,25 @@ class Client:
|
|||||||
lang (str, optional): The selected language for transcription. Default is None.
|
lang (str, optional): The selected language for transcription. Default is None.
|
||||||
translate (bool, optional): Specifies if the task is translation. Default is False.
|
translate (bool, optional): Specifies if the task is translation. Default is False.
|
||||||
"""
|
"""
|
||||||
self.chunk = 4096
|
|
||||||
self.format = pyaudio.paInt16
|
|
||||||
self.channels = 1
|
|
||||||
self.rate = 16000
|
|
||||||
self.record_seconds = 60000
|
|
||||||
self.recording = False
|
self.recording = False
|
||||||
self.task = "transcribe"
|
self.task = "transcribe"
|
||||||
self.uid = str(uuid.uuid4())
|
self.uid = str(uuid.uuid4())
|
||||||
self.waiting = False
|
self.waiting = False
|
||||||
self.last_response_recieved = None
|
self.last_response_received = None
|
||||||
self.disconnect_if_no_response_for = 15
|
self.disconnect_if_no_response_for = 15
|
||||||
self.language = lang
|
self.language = lang
|
||||||
self.model = model
|
self.model = model
|
||||||
self.server_error = False
|
self.server_error = False
|
||||||
self.srt_file_path = srt_file_path
|
self.srt_file_path = srt_file_path
|
||||||
self.use_vad = use_vad
|
self.use_vad = use_vad
|
||||||
self.last_recieved_segment = None
|
self.last_segment = None
|
||||||
|
self.last_received_segment = None
|
||||||
|
self.log_transcription = log_transcription
|
||||||
|
|
||||||
if translate:
|
if translate:
|
||||||
self.task = "translate"
|
self.task = "translate"
|
||||||
|
|
||||||
self.timestamp_offset = 0.0
|
|
||||||
self.audio_bytes = None
|
self.audio_bytes = None
|
||||||
self.p = pyaudio.PyAudio()
|
|
||||||
self.stream = self.p.open(
|
|
||||||
format=self.format,
|
|
||||||
channels=self.channels,
|
|
||||||
rate=self.rate,
|
|
||||||
input=True,
|
|
||||||
frames_per_buffer=self.chunk,
|
|
||||||
)
|
|
||||||
|
|
||||||
if host is not None and port is not None:
|
if host is not None and port is not None:
|
||||||
socket_url = f"ws://{host}:{port}"
|
socket_url = f"ws://{host}:{port}"
|
||||||
@@ -96,7 +87,6 @@ class Client:
|
|||||||
self.ws_thread.setDaemon(True)
|
self.ws_thread.setDaemon(True)
|
||||||
self.ws_thread.start()
|
self.ws_thread.start()
|
||||||
|
|
||||||
self.frames = b""
|
|
||||||
self.transcript = []
|
self.transcript = []
|
||||||
print("[INFO]: * recording")
|
print("[INFO]: * recording")
|
||||||
|
|
||||||
@@ -124,15 +114,16 @@ class Client:
|
|||||||
(not self.transcript or
|
(not self.transcript or
|
||||||
float(seg['start']) >= float(self.transcript[-1]['end']))):
|
float(seg['start']) >= float(self.transcript[-1]['end']))):
|
||||||
self.transcript.append(seg)
|
self.transcript.append(seg)
|
||||||
# update last received segment and last valild responsne time
|
# update last received segment and last valid response time
|
||||||
if self.last_recieved_segment is None or self.last_recieved_segment != segments[-1]["text"]:
|
if self.last_received_segment is None or self.last_received_segment != segments[-1]["text"]:
|
||||||
self.last_response_recieved = time.time()
|
self.last_response_received = time.time()
|
||||||
self.last_recieved_segment = segments[-1]["text"]
|
self.last_received_segment = segments[-1]["text"]
|
||||||
|
|
||||||
# Truncate to last 3 entries for brevity.
|
if self.log_transcription:
|
||||||
text = text[-3:]
|
# Truncate to last 3 entries for brevity.
|
||||||
utils.clear_screen()
|
text = text[-3:]
|
||||||
utils.print_transcript(text)
|
utils.clear_screen()
|
||||||
|
utils.print_transcript(text)
|
||||||
|
|
||||||
def on_message(self, ws, message):
|
def on_message(self, ws, message):
|
||||||
"""
|
"""
|
||||||
@@ -162,7 +153,7 @@ class Client:
|
|||||||
self.recording = False
|
self.recording = False
|
||||||
|
|
||||||
if "message" in message.keys() and message["message"] == "SERVER_READY":
|
if "message" in message.keys() and message["message"] == "SERVER_READY":
|
||||||
self.last_response_recieved = time.time()
|
self.last_response_received = time.time()
|
||||||
self.recording = True
|
self.recording = True
|
||||||
self.server_backend = message["backend"]
|
self.server_backend = message["backend"]
|
||||||
print(f"[INFO]: Server Running with backend {self.server_backend}")
|
print(f"[INFO]: Server Running with backend {self.server_backend}")
|
||||||
@@ -187,7 +178,6 @@ class Client:
|
|||||||
def on_close(self, ws, close_status_code, close_msg):
|
def on_close(self, ws, close_status_code, close_msg):
|
||||||
print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
|
print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
|
||||||
self.recording = False
|
self.recording = False
|
||||||
self.server_error = False
|
|
||||||
self.waiting = False
|
self.waiting = False
|
||||||
|
|
||||||
def on_open(self, ws):
|
def on_open(self, ws):
|
||||||
@@ -214,23 +204,6 @@ class Client:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
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
|
|
||||||
have values between -1 and 1.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
audio_bytes (bytes): Audio data in bytes.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
|
|
||||||
"""
|
|
||||||
raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
|
|
||||||
return raw_data.astype(np.float32) / 32768.0
|
|
||||||
|
|
||||||
def send_packet_to_server(self, message):
|
def send_packet_to_server(self, message):
|
||||||
"""
|
"""
|
||||||
Send an audio packet to the server using WebSocket.
|
Send an audio packet to the server using WebSocket.
|
||||||
@@ -244,62 +217,6 @@ class Client:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(e)
|
print(e)
|
||||||
|
|
||||||
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
|
|
||||||
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.
|
|
||||||
|
|
||||||
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(
|
|
||||||
format=self.p.get_format_from_width(wavfile.getsampwidth()),
|
|
||||||
channels=wavfile.getnchannels(),
|
|
||||||
rate=wavfile.getframerate(),
|
|
||||||
input=True,
|
|
||||||
output=True,
|
|
||||||
frames_per_buffer=self.chunk,
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
while self.recording:
|
|
||||||
data = wavfile.readframes(self.chunk)
|
|
||||||
if data == b"":
|
|
||||||
break
|
|
||||||
|
|
||||||
audio_array = self.bytes_to_float_array(data)
|
|
||||||
self.send_packet_to_server(audio_array.tobytes())
|
|
||||||
self.stream.write(data)
|
|
||||||
|
|
||||||
wavfile.close()
|
|
||||||
|
|
||||||
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()
|
|
||||||
self.close_websocket()
|
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
|
||||||
wavfile.close()
|
|
||||||
self.stream.stop_stream()
|
|
||||||
self.stream.close()
|
|
||||||
self.p.terminate()
|
|
||||||
self.close_websocket()
|
|
||||||
if self.server_backend == "faster_whisper":
|
|
||||||
self.write_srt_file(self.srt_file_path)
|
|
||||||
print("[INFO]: Keyboard interrupt.")
|
|
||||||
|
|
||||||
def close_websocket(self):
|
def close_websocket(self):
|
||||||
"""
|
"""
|
||||||
Close the WebSocket connection and join the WebSocket thread.
|
Close the WebSocket connection and join the WebSocket thread.
|
||||||
@@ -327,6 +244,330 @@ class Client:
|
|||||||
"""
|
"""
|
||||||
return self.client_socket
|
return self.client_socket
|
||||||
|
|
||||||
|
def write_srt_file(self, output_path="output.srt"):
|
||||||
|
"""
|
||||||
|
Writes out the transcript in .srt format.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
message (output_path, optional): The path to the target file. Default is "output.srt".
|
||||||
|
|
||||||
|
"""
|
||||||
|
if self.server_backend == "faster_whisper":
|
||||||
|
if (self.last_segment):
|
||||||
|
self.transcript.append(self.last_segment)
|
||||||
|
utils.create_srt_file(self.transcript, output_path)
|
||||||
|
|
||||||
|
def wait_before_disconnect(self):
|
||||||
|
"""Waits a bit before disconnecting in order to process pending responses."""
|
||||||
|
assert self.last_response_received
|
||||||
|
while time.time() - self.last_response_received < self.disconnect_if_no_response_for:
|
||||||
|
continue
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionTeeClient:
|
||||||
|
"""
|
||||||
|
Client for handling audio recording, streaming, and transcription tasks via one or more
|
||||||
|
WebSocket connections.
|
||||||
|
|
||||||
|
Acts as a high-level client for audio transcription tasks using a WebSocket connection. It can be used
|
||||||
|
to send audio data for transcription to one or more servers, and receive transcribed text segments.
|
||||||
|
Args:
|
||||||
|
clients (list): one or more previously initialized Client instances
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
clients (list): the underlying Client instances responsible for handling WebSocket connections.
|
||||||
|
"""
|
||||||
|
def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav"):
|
||||||
|
self.clients = clients
|
||||||
|
if not self.clients:
|
||||||
|
raise Exception("At least one client is required.")
|
||||||
|
self.chunk = 4096
|
||||||
|
self.format = pyaudio.paInt16
|
||||||
|
self.channels = 1
|
||||||
|
self.rate = 16000
|
||||||
|
self.record_seconds = 60000
|
||||||
|
self.save_output_recording = save_output_recording
|
||||||
|
self.output_recording_filename = output_recording_filename
|
||||||
|
self.frames = b""
|
||||||
|
self.p = pyaudio.PyAudio()
|
||||||
|
try:
|
||||||
|
self.stream = self.p.open(
|
||||||
|
format=self.format,
|
||||||
|
channels=self.channels,
|
||||||
|
rate=self.rate,
|
||||||
|
input=True,
|
||||||
|
frames_per_buffer=self.chunk,
|
||||||
|
)
|
||||||
|
except OSError as error:
|
||||||
|
print(f"[WARN]: Unable to access microphone. {error}")
|
||||||
|
self.stream = None
|
||||||
|
|
||||||
|
def __call__(self, audio=None, rtsp_url=None, hls_url=None, save_file=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
|
||||||
|
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.
|
||||||
|
|
||||||
|
"""
|
||||||
|
assert sum(
|
||||||
|
source is not None for source in [audio, rtsp_url, hls_url]
|
||||||
|
) <= 1, 'You must provide only one selected source'
|
||||||
|
|
||||||
|
print("[INFO]: Waiting for server ready ...")
|
||||||
|
for client in self.clients:
|
||||||
|
while not client.recording:
|
||||||
|
if client.waiting or client.server_error:
|
||||||
|
self.close_all_clients()
|
||||||
|
return
|
||||||
|
|
||||||
|
print("[INFO]: Server Ready!")
|
||||||
|
if hls_url is not None:
|
||||||
|
self.process_hls_stream(hls_url, save_file)
|
||||||
|
elif audio is not None:
|
||||||
|
resampled_file = utils.resample(audio)
|
||||||
|
self.play_file(resampled_file)
|
||||||
|
elif rtsp_url is not None:
|
||||||
|
self.process_rtsp_stream(rtsp_url)
|
||||||
|
else:
|
||||||
|
self.record()
|
||||||
|
|
||||||
|
def close_all_clients(self):
|
||||||
|
"""Closes all client websockets."""
|
||||||
|
for client in self.clients:
|
||||||
|
client.close_websocket()
|
||||||
|
|
||||||
|
def write_all_clients_srt(self):
|
||||||
|
"""Writes out .srt files for all clients."""
|
||||||
|
for client in self.clients:
|
||||||
|
client.write_srt_file(client.srt_file_path)
|
||||||
|
|
||||||
|
def multicast_packet(self, packet, unconditional=False):
|
||||||
|
"""
|
||||||
|
Sends an identical packet via all clients.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
packet (bytes): The audio data packet in bytes to be sent.
|
||||||
|
unconditional (bool, optional): If true, send regardless of whether clients are recording. Default is False.
|
||||||
|
"""
|
||||||
|
for client in self.clients:
|
||||||
|
if (unconditional or client.recording):
|
||||||
|
client.send_packet_to_server(packet)
|
||||||
|
|
||||||
|
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
|
||||||
|
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.
|
||||||
|
|
||||||
|
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(
|
||||||
|
format=self.p.get_format_from_width(wavfile.getsampwidth()),
|
||||||
|
channels=wavfile.getnchannels(),
|
||||||
|
rate=wavfile.getframerate(),
|
||||||
|
input=True,
|
||||||
|
output=True,
|
||||||
|
frames_per_buffer=self.chunk,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
while any(client.recording for client in self.clients):
|
||||||
|
data = wavfile.readframes(self.chunk)
|
||||||
|
if data == b"":
|
||||||
|
break
|
||||||
|
|
||||||
|
audio_array = self.bytes_to_float_array(data)
|
||||||
|
self.multicast_packet(audio_array.tobytes())
|
||||||
|
self.stream.write(data)
|
||||||
|
|
||||||
|
wavfile.close()
|
||||||
|
|
||||||
|
for client in self.clients:
|
||||||
|
client.wait_before_disconnect()
|
||||||
|
self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
|
||||||
|
self.write_all_clients_srt()
|
||||||
|
self.stream.close()
|
||||||
|
self.close_all_clients()
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
wavfile.close()
|
||||||
|
self.stream.stop_stream()
|
||||||
|
self.stream.close()
|
||||||
|
self.p.terminate()
|
||||||
|
self.close_all_clients()
|
||||||
|
self.write_all_clients_srt()
|
||||||
|
print("[INFO]: Keyboard interrupt.")
|
||||||
|
|
||||||
|
def process_rtsp_stream(self, rtsp_url):
|
||||||
|
"""
|
||||||
|
Connect to an RTSP source, process the audio stream, and send it for trascription.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
rtsp_url (str): The URL of the RTSP stream source.
|
||||||
|
"""
|
||||||
|
process = self.get_rtsp_ffmpeg_process(rtsp_url)
|
||||||
|
self.handle_ffmpeg_process(process, stream_type='RTSP')
|
||||||
|
|
||||||
|
def process_hls_stream(self, hls_url, save_file):
|
||||||
|
"""
|
||||||
|
Connect to an HLS source, process the audio stream, and send it for transcription.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
hls_url (str): The URL of the HLS stream source.
|
||||||
|
save_file (str, optional): Local path to save the network stream.
|
||||||
|
"""
|
||||||
|
process = self.get_hls_ffmpeg_process(hls_url, save_file)
|
||||||
|
self.handle_ffmpeg_process(process, stream_type='HLS')
|
||||||
|
|
||||||
|
def handle_ffmpeg_process(self, process, stream_type):
|
||||||
|
print(f"[INFO]: Connecting to {stream_type} stream...")
|
||||||
|
stderr_thread = threading.Thread(target=self.consume_stderr, args=(process,))
|
||||||
|
stderr_thread.start()
|
||||||
|
try:
|
||||||
|
# Process the stream
|
||||||
|
while True:
|
||||||
|
in_bytes = process.stdout.read(self.chunk * 2) # 2 bytes per sample
|
||||||
|
if not in_bytes:
|
||||||
|
break
|
||||||
|
audio_array = self.bytes_to_float_array(in_bytes)
|
||||||
|
self.multicast_packet(audio_array.tobytes())
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
|
||||||
|
finally:
|
||||||
|
self.close_all_clients()
|
||||||
|
self.write_all_clients_srt()
|
||||||
|
if process:
|
||||||
|
process.kill()
|
||||||
|
|
||||||
|
print(f"[INFO]: {stream_type} stream processing finished.")
|
||||||
|
|
||||||
|
def get_rtsp_ffmpeg_process(self, rtsp_url):
|
||||||
|
return (
|
||||||
|
ffmpeg
|
||||||
|
.input(rtsp_url, threads=0)
|
||||||
|
.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
|
||||||
|
.run_async(pipe_stdout=True, pipe_stderr=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_hls_ffmpeg_process(self, hls_url, save_file):
|
||||||
|
if save_file is None:
|
||||||
|
process = (
|
||||||
|
ffmpeg
|
||||||
|
.input(hls_url, threads=0)
|
||||||
|
.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
|
||||||
|
.run_async(pipe_stdout=True, pipe_stderr=True)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
input = ffmpeg.input(hls_url, threads=0)
|
||||||
|
output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
|
||||||
|
output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
|
||||||
|
process = (
|
||||||
|
ffmpeg.merge_outputs(output_file, output_std)
|
||||||
|
.run_async(pipe_stdout=True, pipe_stderr=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
return process
|
||||||
|
|
||||||
|
def consume_stderr(self, process):
|
||||||
|
"""
|
||||||
|
Consume and log the stderr output of a process in a separate thread.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
process (subprocess.Popen): The process whose stderr output will be logged.
|
||||||
|
"""
|
||||||
|
for line in iter(process.stderr.readline, b""):
|
||||||
|
logging.debug(f'[STDERR]: {line.decode()}')
|
||||||
|
|
||||||
|
def save_chunk(self, n_audio_file):
|
||||||
|
"""
|
||||||
|
Saves the current audio frames to a WAV file in a separate thread.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
n_audio_file (int): The index of the audio file which determines the filename.
|
||||||
|
This helps in maintaining the order and uniqueness of each chunk.
|
||||||
|
"""
|
||||||
|
t = threading.Thread(
|
||||||
|
target=self.write_audio_frames_to_file,
|
||||||
|
args=(self.frames[:], f"chunks/{n_audio_file}.wav",),
|
||||||
|
)
|
||||||
|
t.start()
|
||||||
|
|
||||||
|
def finalize_recording(self, n_audio_file):
|
||||||
|
"""
|
||||||
|
Finalizes the recording process by saving any remaining audio frames,
|
||||||
|
closing the audio stream, and terminating the process.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
n_audio_file (int): The file index to be used if there are remaining audio frames to be saved.
|
||||||
|
This index is incremented before use if the last chunk is saved.
|
||||||
|
"""
|
||||||
|
if self.save_output_recording and len(self.frames):
|
||||||
|
self.write_audio_frames_to_file(
|
||||||
|
self.frames[:], f"chunks/{n_audio_file}.wav"
|
||||||
|
)
|
||||||
|
n_audio_file += 1
|
||||||
|
self.stream.stop_stream()
|
||||||
|
self.stream.close()
|
||||||
|
self.p.terminate()
|
||||||
|
self.close_all_clients()
|
||||||
|
if self.save_output_recording:
|
||||||
|
self.write_output_recording(n_audio_file)
|
||||||
|
self.write_all_clients_srt()
|
||||||
|
|
||||||
|
def record(self):
|
||||||
|
"""
|
||||||
|
Record audio data from the input stream and save it to a WAV file.
|
||||||
|
|
||||||
|
Continuously records audio data from the input stream, sends it to the server via a WebSocket
|
||||||
|
connection, and simultaneously saves it to multiple WAV files in chunks. It stops recording when
|
||||||
|
the `RECORD_SECONDS` duration is reached or when the `RECORDING` flag is set to `False`.
|
||||||
|
|
||||||
|
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 method combines all the saved audio chunks into the specified `out_file`.
|
||||||
|
"""
|
||||||
|
n_audio_file = 0
|
||||||
|
if self.save_output_recording:
|
||||||
|
if os.path.exists("chunks"):
|
||||||
|
shutil.rmtree("chunks")
|
||||||
|
os.makedirs("chunks")
|
||||||
|
try:
|
||||||
|
for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
|
||||||
|
if not any(client.recording for client in self.clients):
|
||||||
|
break
|
||||||
|
data = self.stream.read(self.chunk, exception_on_overflow=False)
|
||||||
|
self.frames += data
|
||||||
|
|
||||||
|
audio_array = self.bytes_to_float_array(data)
|
||||||
|
|
||||||
|
self.multicast_packet(audio_array.tobytes())
|
||||||
|
|
||||||
|
# save frames if more than a minute
|
||||||
|
if len(self.frames) > 60 * self.rate:
|
||||||
|
if self.save_output_recording:
|
||||||
|
self.save_chunk(n_audio_file)
|
||||||
|
n_audio_file += 1
|
||||||
|
self.frames = b""
|
||||||
|
self.write_all_clients_srt()
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
self.finalize_recording(n_audio_file)
|
||||||
|
|
||||||
def write_audio_frames_to_file(self, frames, file_name):
|
def write_audio_frames_to_file(self, frames, file_name):
|
||||||
"""
|
"""
|
||||||
Write audio frames to a WAV file.
|
Write audio frames to a WAV file.
|
||||||
@@ -346,104 +587,7 @@ class Client:
|
|||||||
wavfile.setframerate(self.rate)
|
wavfile.setframerate(self.rate)
|
||||||
wavfile.writeframes(frames)
|
wavfile.writeframes(frames)
|
||||||
|
|
||||||
def process_hls_stream(self, hls_url):
|
def write_output_recording(self, n_audio_file):
|
||||||
"""
|
|
||||||
Connect to an HLS source, process the audio stream, and send it for transcription.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
hls_url (str): The URL of the HLS stream source.
|
|
||||||
"""
|
|
||||||
print("[INFO]: Connecting to HLS stream...")
|
|
||||||
process = None # Initialize process to None
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Connecting to the HLS stream using ffmpeg-python
|
|
||||||
process = (
|
|
||||||
ffmpeg
|
|
||||||
.input(hls_url, threads=0)
|
|
||||||
.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
|
|
||||||
.run_async(pipe_stdout=True, pipe_stderr=True)
|
|
||||||
)
|
|
||||||
|
|
||||||
# Process the stream
|
|
||||||
while True:
|
|
||||||
in_bytes = process.stdout.read(self.chunk * 2) # 2 bytes per sample
|
|
||||||
if not in_bytes:
|
|
||||||
break
|
|
||||||
audio_array = self.bytes_to_float_array(in_bytes)
|
|
||||||
self.send_packet_to_server(audio_array.tobytes())
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"[ERROR]: Failed to connect to HLS stream: {e}")
|
|
||||||
finally:
|
|
||||||
if process:
|
|
||||||
process.kill()
|
|
||||||
|
|
||||||
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.
|
|
||||||
|
|
||||||
Continuously records audio data from the input stream, sends it to the server via a WebSocket
|
|
||||||
connection, and simultaneously saves it to multiple WAV files in chunks. It stops recording when
|
|
||||||
the `RECORD_SECONDS` duration is reached or when the `RECORDING` flag is set to `False`.
|
|
||||||
|
|
||||||
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 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".
|
|
||||||
|
|
||||||
"""
|
|
||||||
n_audio_file = 0
|
|
||||||
if not os.path.exists("chunks"):
|
|
||||||
os.makedirs("chunks", exist_ok=True)
|
|
||||||
try:
|
|
||||||
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)
|
|
||||||
self.frames += data
|
|
||||||
|
|
||||||
audio_array = Client.bytes_to_float_array(data)
|
|
||||||
|
|
||||||
self.send_packet_to_server(audio_array.tobytes())
|
|
||||||
|
|
||||||
# save frames if more than a minute
|
|
||||||
if len(self.frames) > 60 * self.rate:
|
|
||||||
t = threading.Thread(
|
|
||||||
target=self.write_audio_frames_to_file,
|
|
||||||
args=(
|
|
||||||
self.frames[:],
|
|
||||||
f"chunks/{n_audio_file}.wav",
|
|
||||||
),
|
|
||||||
)
|
|
||||||
t.start()
|
|
||||||
n_audio_file += 1
|
|
||||||
self.frames = b""
|
|
||||||
if self.server_backend == "faster_whisper":
|
|
||||||
self.write_srt_file(self.srt_file_path)
|
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
|
||||||
if len(self.frames):
|
|
||||||
self.write_audio_frames_to_file(
|
|
||||||
self.frames[:], f"chunks/{n_audio_file}.wav"
|
|
||||||
)
|
|
||||||
n_audio_file += 1
|
|
||||||
self.stream.stop_stream()
|
|
||||||
self.stream.close()
|
|
||||||
self.p.terminate()
|
|
||||||
self.close_websocket()
|
|
||||||
|
|
||||||
self.write_output_recording(n_audio_file, out_file)
|
|
||||||
if self.server_backend == "faster_whisper":
|
|
||||||
self.write_srt_file(self.srt_file_path)
|
|
||||||
|
|
||||||
def write_output_recording(self, n_audio_file, out_file):
|
|
||||||
"""
|
"""
|
||||||
Combine and save recorded audio chunks into a single WAV file.
|
Combine and save recorded audio chunks into a single WAV file.
|
||||||
|
|
||||||
@@ -462,7 +606,7 @@ class Client:
|
|||||||
for i in range(n_audio_file)
|
for i in range(n_audio_file)
|
||||||
if os.path.exists(f"chunks/{i}.wav")
|
if os.path.exists(f"chunks/{i}.wav")
|
||||||
]
|
]
|
||||||
with wave.open(out_file, "wb") as wavfile:
|
with wave.open(self.output_recording_filename, "wb") as wavfile:
|
||||||
wavfile: wave.Wave_write
|
wavfile: wave.Wave_write
|
||||||
wavfile.setnchannels(self.channels)
|
wavfile.setnchannels(self.channels)
|
||||||
wavfile.setsampwidth(2)
|
wavfile.setsampwidth(2)
|
||||||
@@ -477,15 +621,31 @@ class Client:
|
|||||||
# remove this file
|
# remove this file
|
||||||
os.remove(in_file)
|
os.remove(in_file)
|
||||||
wavfile.close()
|
wavfile.close()
|
||||||
|
# clean up temporary directory to store chunks
|
||||||
|
if os.path.exists("chunks"):
|
||||||
|
shutil.rmtree("chunks")
|
||||||
|
|
||||||
def write_srt_file(self, output_path="output.srt"):
|
@staticmethod
|
||||||
self.transcript.append(self.last_segment)
|
def bytes_to_float_array(audio_bytes):
|
||||||
utils.create_srt_file(self.transcript, output_path)
|
"""
|
||||||
|
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
|
||||||
|
have values between -1 and 1.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
audio_bytes (bytes): Audio data in bytes.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
|
||||||
|
"""
|
||||||
|
raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
|
||||||
|
return raw_data.astype(np.float32) / 32768.0
|
||||||
|
|
||||||
|
|
||||||
class TranscriptionClient:
|
class TranscriptionClient(TranscriptionTeeClient):
|
||||||
"""
|
"""
|
||||||
Client for handling audio transcription tasks via a WebSocket connection.
|
Client for handling audio transcription tasks via a single WebSocket connection.
|
||||||
|
|
||||||
Acts as a high-level client for audio transcription tasks using a WebSocket connection. It can be used
|
Acts as a high-level client for audio transcription tasks using a WebSocket connection. It can be used
|
||||||
to send audio data for transcription to a server and receive transcribed text segments.
|
to send audio data for transcription to a server and receive transcribed text segments.
|
||||||
@@ -495,6 +655,9 @@ class TranscriptionClient:
|
|||||||
port (int): The port number to connect to on the server.
|
port (int): The port number to connect to on the server.
|
||||||
lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
|
lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
|
||||||
translate (bool, optional): Indicates whether translation tasks are required (default is False).
|
translate (bool, optional): Indicates whether translation tasks are required (default is False).
|
||||||
|
save_output_recording (bool, optional): Indicates whether to save recording from microphone.
|
||||||
|
output_recording_filename (str, optional): File to save the output recording.
|
||||||
|
output_transcription_path (str, optional): File to save the output transcription.
|
||||||
|
|
||||||
Attributes:
|
Attributes:
|
||||||
client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
|
client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
|
||||||
@@ -506,32 +669,27 @@ class TranscriptionClient:
|
|||||||
transcription_client()
|
transcription_client()
|
||||||
```
|
```
|
||||||
"""
|
"""
|
||||||
def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True):
|
def __init__(
|
||||||
self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
|
self,
|
||||||
|
host,
|
||||||
def __call__(self, audio=None, hls_url=None):
|
port,
|
||||||
"""
|
lang=None,
|
||||||
Start the transcription process.
|
translate=False,
|
||||||
|
model="small",
|
||||||
Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
|
use_vad=True,
|
||||||
to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it
|
save_output_recording=False,
|
||||||
will be played and streamed to the server; otherwise, it will perform live recording.
|
output_recording_filename="./output_recording.wav",
|
||||||
|
output_transcription_path="./output.srt",
|
||||||
Args:
|
log_transcription=True,
|
||||||
audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
|
):
|
||||||
|
self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad, log_transcription=log_transcription)
|
||||||
"""
|
if save_output_recording and not output_recording_filename.endswith(".wav"):
|
||||||
print("[INFO]: Waiting for server ready ...")
|
raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
|
||||||
while not self.client.recording:
|
if not output_transcription_path.endswith(".srt"):
|
||||||
if self.client.waiting or self.client.server_error:
|
raise ValueError(f"Please provide a valid `output_transcription_path`: {output_transcription_path}. The file extension should be `.srt`.")
|
||||||
self.client.close_websocket()
|
TranscriptionTeeClient.__init__(
|
||||||
return
|
self,
|
||||||
|
[self.client],
|
||||||
print("[INFO]: Server Ready!")
|
save_output_recording=save_output_recording,
|
||||||
if hls_url is not None:
|
output_recording_filename=output_recording_filename
|
||||||
self.client.process_hls_stream(hls_url)
|
)
|
||||||
elif audio is not None:
|
|
||||||
resampled_file = utils.resample(audio)
|
|
||||||
self.client.play_file(resampled_file)
|
|
||||||
else:
|
|
||||||
self.client.record()
|
|
||||||
|
|||||||
+135
-33
@@ -4,6 +4,9 @@ import threading
|
|||||||
import json
|
import json
|
||||||
import functools
|
import functools
|
||||||
import logging
|
import logging
|
||||||
|
from enum import Enum
|
||||||
|
from typing import List, Optional
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from websockets.sync.server import serve
|
from websockets.sync.server import serve
|
||||||
@@ -121,6 +124,25 @@ class ClientManager:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
class BackendType(Enum):
|
||||||
|
FASTER_WHISPER = "faster_whisper"
|
||||||
|
TENSORRT = "tensorrt"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def valid_types() -> List[str]:
|
||||||
|
return [backend_type.value for backend_type in BackendType]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def is_valid(backend: str) -> bool:
|
||||||
|
return backend in BackendType.valid_types()
|
||||||
|
|
||||||
|
def is_faster_whisper(self) -> bool:
|
||||||
|
return self == BackendType.FASTER_WHISPER
|
||||||
|
|
||||||
|
def is_tensorrt(self) -> bool:
|
||||||
|
return self == BackendType.TENSORRT
|
||||||
|
|
||||||
|
|
||||||
class TranscriptionServer:
|
class TranscriptionServer:
|
||||||
RATE = 16000
|
RATE = 16000
|
||||||
|
|
||||||
@@ -128,12 +150,15 @@ class TranscriptionServer:
|
|||||||
self.client_manager = ClientManager()
|
self.client_manager = ClientManager()
|
||||||
self.no_voice_activity_chunks = 0
|
self.no_voice_activity_chunks = 0
|
||||||
self.use_vad = True
|
self.use_vad = True
|
||||||
|
self.single_model = False
|
||||||
|
|
||||||
def initialize_client(
|
def initialize_client(
|
||||||
self, websocket, options, faster_whisper_custom_model_path,
|
self, websocket, options, faster_whisper_custom_model_path,
|
||||||
whisper_tensorrt_path, trt_multilingual
|
whisper_tensorrt_path, trt_multilingual
|
||||||
):
|
):
|
||||||
if self.backend == "tensorrt":
|
client: Optional[ServeClientBase] = None
|
||||||
|
|
||||||
|
if self.backend.is_tensorrt():
|
||||||
try:
|
try:
|
||||||
client = ServeClientTensorRT(
|
client = ServeClientTensorRT(
|
||||||
websocket,
|
websocket,
|
||||||
@@ -141,7 +166,8 @@ class TranscriptionServer:
|
|||||||
language=options["language"],
|
language=options["language"],
|
||||||
task=options["task"],
|
task=options["task"],
|
||||||
client_uid=options["uid"],
|
client_uid=options["uid"],
|
||||||
model=whisper_tensorrt_path
|
model=whisper_tensorrt_path,
|
||||||
|
single_model=self.single_model,
|
||||||
)
|
)
|
||||||
logging.info("Running TensorRT backend.")
|
logging.info("Running TensorRT backend.")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -153,9 +179,9 @@ class TranscriptionServer:
|
|||||||
"message": "TensorRT-LLM not supported on Server yet. "
|
"message": "TensorRT-LLM not supported on Server yet. "
|
||||||
"Reverting to available backend: 'faster_whisper'"
|
"Reverting to available backend: 'faster_whisper'"
|
||||||
}))
|
}))
|
||||||
self.backend = "faster_whisper"
|
self.backend = BackendType.FASTER_WHISPER
|
||||||
|
|
||||||
if self.backend == "faster_whisper":
|
if self.backend.is_faster_whisper():
|
||||||
if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
|
if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
|
||||||
logging.info(f"Using custom model {faster_whisper_custom_model_path}")
|
logging.info(f"Using custom model {faster_whisper_custom_model_path}")
|
||||||
options["model"] = faster_whisper_custom_model_path
|
options["model"] = faster_whisper_custom_model_path
|
||||||
@@ -168,9 +194,13 @@ class TranscriptionServer:
|
|||||||
initial_prompt=options.get("initial_prompt"),
|
initial_prompt=options.get("initial_prompt"),
|
||||||
vad_parameters=options.get("vad_parameters"),
|
vad_parameters=options.get("vad_parameters"),
|
||||||
use_vad=self.use_vad,
|
use_vad=self.use_vad,
|
||||||
|
single_model=self.single_model,
|
||||||
)
|
)
|
||||||
logging.info("Running faster_whisper backend.")
|
logging.info("Running faster_whisper backend.")
|
||||||
|
|
||||||
|
if client is None:
|
||||||
|
raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
|
||||||
|
|
||||||
self.client_manager.add_client(websocket, client)
|
self.client_manager.add_client(websocket, client)
|
||||||
|
|
||||||
def get_audio_from_websocket(self, websocket):
|
def get_audio_from_websocket(self, websocket):
|
||||||
@@ -199,7 +229,7 @@ class TranscriptionServer:
|
|||||||
websocket.close()
|
websocket.close()
|
||||||
return False # Indicates that the connection should not continue
|
return False # Indicates that the connection should not continue
|
||||||
|
|
||||||
if self.backend == "tensorrt":
|
if self.backend.is_tensorrt():
|
||||||
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
|
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
|
||||||
self.initialize_client(websocket, options, faster_whisper_custom_model_path,
|
self.initialize_client(websocket, options, faster_whisper_custom_model_path,
|
||||||
whisper_tensorrt_path, trt_multilingual)
|
whisper_tensorrt_path, trt_multilingual)
|
||||||
@@ -218,11 +248,11 @@ class TranscriptionServer:
|
|||||||
frame_np = self.get_audio_from_websocket(websocket)
|
frame_np = self.get_audio_from_websocket(websocket)
|
||||||
client = self.client_manager.get_client(websocket)
|
client = self.client_manager.get_client(websocket)
|
||||||
if frame_np is False:
|
if frame_np is False:
|
||||||
if self.backend == "tensorrt":
|
if self.backend.is_tensorrt():
|
||||||
client.set_eos(True)
|
client.set_eos(True)
|
||||||
return False
|
return False
|
||||||
|
|
||||||
if self.backend == "tensorrt":
|
if self.backend.is_tensorrt():
|
||||||
voice_active = self.voice_activity(websocket, frame_np)
|
voice_active = self.voice_activity(websocket, frame_np)
|
||||||
if voice_active:
|
if voice_active:
|
||||||
self.no_voice_activity_chunks = 0
|
self.no_voice_activity_chunks = 0
|
||||||
@@ -235,7 +265,7 @@ class TranscriptionServer:
|
|||||||
|
|
||||||
def recv_audio(self,
|
def recv_audio(self,
|
||||||
websocket,
|
websocket,
|
||||||
backend="faster_whisper",
|
backend: BackendType = BackendType.FASTER_WHISPER,
|
||||||
faster_whisper_custom_model_path=None,
|
faster_whisper_custom_model_path=None,
|
||||||
whisper_tensorrt_path=None,
|
whisper_tensorrt_path=None,
|
||||||
trt_multilingual=False):
|
trt_multilingual=False):
|
||||||
@@ -288,7 +318,8 @@ class TranscriptionServer:
|
|||||||
backend="tensorrt",
|
backend="tensorrt",
|
||||||
faster_whisper_custom_model_path=None,
|
faster_whisper_custom_model_path=None,
|
||||||
whisper_tensorrt_path=None,
|
whisper_tensorrt_path=None,
|
||||||
trt_multilingual=False):
|
trt_multilingual=False,
|
||||||
|
single_model=False):
|
||||||
"""
|
"""
|
||||||
Run the transcription server.
|
Run the transcription server.
|
||||||
|
|
||||||
@@ -296,10 +327,23 @@ class TranscriptionServer:
|
|||||||
host (str): The host address to bind the server.
|
host (str): The host address to bind the server.
|
||||||
port (int): The port number to bind the server.
|
port (int): The port number to bind the server.
|
||||||
"""
|
"""
|
||||||
|
if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
|
||||||
|
raise ValueError(f"Custom faster_whisper model '{faster_whisper_custom_model_path}' is not a valid path.")
|
||||||
|
if whisper_tensorrt_path is not None and not os.path.exists(whisper_tensorrt_path):
|
||||||
|
raise ValueError(f"TensorRT model '{whisper_tensorrt_path}' is not a valid path.")
|
||||||
|
if single_model:
|
||||||
|
if faster_whisper_custom_model_path or whisper_tensorrt_path:
|
||||||
|
logging.info("Custom model option was provided. Switching to single model mode.")
|
||||||
|
self.single_model = True
|
||||||
|
# TODO: load model initially
|
||||||
|
else:
|
||||||
|
logging.info("Single model mode currently only works with custom models.")
|
||||||
|
if not BackendType.is_valid(backend):
|
||||||
|
raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
|
||||||
with serve(
|
with serve(
|
||||||
functools.partial(
|
functools.partial(
|
||||||
self.recv_audio,
|
self.recv_audio,
|
||||||
backend=backend,
|
backend=BackendType(backend),
|
||||||
faster_whisper_custom_model_path=faster_whisper_custom_model_path,
|
faster_whisper_custom_model_path=faster_whisper_custom_model_path,
|
||||||
whisper_tensorrt_path=whisper_tensorrt_path,
|
whisper_tensorrt_path=whisper_tensorrt_path,
|
||||||
trt_multilingual=trt_multilingual
|
trt_multilingual=trt_multilingual
|
||||||
@@ -532,7 +576,11 @@ class ServeClientBase(object):
|
|||||||
|
|
||||||
|
|
||||||
class ServeClientTensorRT(ServeClientBase):
|
class ServeClientTensorRT(ServeClientBase):
|
||||||
def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None):
|
|
||||||
|
SINGLE_MODEL = None
|
||||||
|
SINGLE_MODEL_LOCK = threading.Lock()
|
||||||
|
|
||||||
|
def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None, single_model=False):
|
||||||
"""
|
"""
|
||||||
Initialize a ServeClient instance.
|
Initialize a ServeClient instance.
|
||||||
The Whisper model is initialized based on the client's language and device availability.
|
The Whisper model is initialized based on the client's language and device availability.
|
||||||
@@ -546,21 +594,22 @@ class ServeClientTensorRT(ServeClientBase):
|
|||||||
multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
|
multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
|
||||||
language (str, optional): The language for transcription. Defaults to None.
|
language (str, optional): The language for transcription. Defaults to None.
|
||||||
client_uid (str, optional): A unique identifier for the client. Defaults to None.
|
client_uid (str, optional): A unique identifier for the client. Defaults to None.
|
||||||
|
single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
super().__init__(client_uid, websocket)
|
super().__init__(client_uid, websocket)
|
||||||
self.language = language if multilingual else "en"
|
self.language = language if multilingual else "en"
|
||||||
self.task = task
|
self.task = task
|
||||||
self.eos = False
|
self.eos = False
|
||||||
self.transcriber = WhisperTRTLLM(
|
|
||||||
model,
|
if single_model:
|
||||||
assets_dir="assets",
|
if ServeClientTensorRT.SINGLE_MODEL is None:
|
||||||
device="cuda",
|
self.create_model(model, multilingual)
|
||||||
is_multilingual=multilingual,
|
ServeClientTensorRT.SINGLE_MODEL = self.transcriber
|
||||||
language=self.language,
|
else:
|
||||||
task=self.task
|
self.transcriber = ServeClientTensorRT.SINGLE_MODEL
|
||||||
)
|
else:
|
||||||
self.warmup()
|
self.create_model(model, multilingual)
|
||||||
|
|
||||||
# threading
|
# threading
|
||||||
self.trans_thread = threading.Thread(target=self.speech_to_text)
|
self.trans_thread = threading.Thread(target=self.speech_to_text)
|
||||||
@@ -572,6 +621,21 @@ class ServeClientTensorRT(ServeClientBase):
|
|||||||
"backend": "tensorrt"
|
"backend": "tensorrt"
|
||||||
}))
|
}))
|
||||||
|
|
||||||
|
def create_model(self, model, multilingual, warmup=True):
|
||||||
|
"""
|
||||||
|
Instantiates a new model, sets it as the transcriber and does warmup if desired.
|
||||||
|
"""
|
||||||
|
self.transcriber = WhisperTRTLLM(
|
||||||
|
model,
|
||||||
|
assets_dir="assets",
|
||||||
|
device="cuda",
|
||||||
|
is_multilingual=multilingual,
|
||||||
|
language=self.language,
|
||||||
|
task=self.task
|
||||||
|
)
|
||||||
|
if warmup:
|
||||||
|
self.warmup()
|
||||||
|
|
||||||
def warmup(self, warmup_steps=10):
|
def warmup(self, warmup_steps=10):
|
||||||
"""
|
"""
|
||||||
Warmup TensorRT since first few inferences are slow.
|
Warmup TensorRT since first few inferences are slow.
|
||||||
@@ -616,12 +680,16 @@ class ServeClientTensorRT(ServeClientBase):
|
|||||||
Args:
|
Args:
|
||||||
input_bytes (np.array): The audio chunk to transcribe.
|
input_bytes (np.array): The audio chunk to transcribe.
|
||||||
"""
|
"""
|
||||||
|
if ServeClientTensorRT.SINGLE_MODEL:
|
||||||
|
ServeClientTensorRT.SINGLE_MODEL_LOCK.acquire()
|
||||||
logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
|
logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
|
||||||
mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
|
mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
|
||||||
last_segment = self.transcriber.transcribe(
|
last_segment = self.transcriber.transcribe(
|
||||||
mel,
|
mel,
|
||||||
text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>"
|
text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>"
|
||||||
)
|
)
|
||||||
|
if ServeClientTensorRT.SINGLE_MODEL:
|
||||||
|
ServeClientTensorRT.SINGLE_MODEL_LOCK.release()
|
||||||
if last_segment:
|
if last_segment:
|
||||||
self.handle_transcription_output(last_segment, duration)
|
self.handle_transcription_output(last_segment, duration)
|
||||||
|
|
||||||
@@ -681,8 +749,12 @@ class ServeClientTensorRT(ServeClientBase):
|
|||||||
|
|
||||||
|
|
||||||
class ServeClientFasterWhisper(ServeClientBase):
|
class ServeClientFasterWhisper(ServeClientBase):
|
||||||
|
|
||||||
|
SINGLE_MODEL = None
|
||||||
|
SINGLE_MODEL_LOCK = threading.Lock()
|
||||||
|
|
||||||
def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
|
def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
|
||||||
initial_prompt=None, vad_parameters=None, use_vad=True):
|
initial_prompt=None, vad_parameters=None, use_vad=True, single_model=False):
|
||||||
"""
|
"""
|
||||||
Initialize a ServeClient instance.
|
Initialize a ServeClient instance.
|
||||||
The Whisper model is initialized based on the client's language and device availability.
|
The Whisper model is initialized based on the client's language and device availability.
|
||||||
@@ -697,6 +769,7 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
client_uid (str, optional): A unique identifier for the client. Defaults to None.
|
client_uid (str, optional): A unique identifier for the client. Defaults to None.
|
||||||
model (str, optional): The whisper model size. Defaults to 'small.en'
|
model (str, optional): The whisper model size. Defaults to 'small.en'
|
||||||
initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
|
initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
|
||||||
|
single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
|
||||||
"""
|
"""
|
||||||
super().__init__(client_uid, websocket)
|
super().__init__(client_uid, websocket)
|
||||||
self.model_sizes = [
|
self.model_sizes = [
|
||||||
@@ -714,16 +787,25 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
self.no_speech_thresh = 0.45
|
self.no_speech_thresh = 0.45
|
||||||
|
|
||||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||||
|
if device == "cuda":
|
||||||
|
major, _ = torch.cuda.get_device_capability(device)
|
||||||
|
self.compute_type = "float16" if major >= 7 else "float32"
|
||||||
|
else:
|
||||||
|
self.compute_type = "int8"
|
||||||
|
|
||||||
if self.model_size_or_path is None:
|
if self.model_size_or_path is None:
|
||||||
return
|
return
|
||||||
|
logging.info(f"Using Device={device} with precision {self.compute_type}")
|
||||||
|
|
||||||
|
if single_model:
|
||||||
|
if ServeClientFasterWhisper.SINGLE_MODEL is None:
|
||||||
|
self.create_model(device)
|
||||||
|
ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
|
||||||
|
else:
|
||||||
|
self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
|
||||||
|
else:
|
||||||
|
self.create_model(device)
|
||||||
|
|
||||||
self.transcriber = WhisperModel(
|
|
||||||
self.model_size_or_path,
|
|
||||||
device=device,
|
|
||||||
compute_type="int8" if device == "cpu" else "float16",
|
|
||||||
local_files_only=False,
|
|
||||||
)
|
|
||||||
self.use_vad = use_vad
|
self.use_vad = use_vad
|
||||||
|
|
||||||
# threading
|
# threading
|
||||||
@@ -739,6 +821,17 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def create_model(self, device):
|
||||||
|
"""
|
||||||
|
Instantiates a new model, sets it as the transcriber.
|
||||||
|
"""
|
||||||
|
self.transcriber = WhisperModel(
|
||||||
|
self.model_size_or_path,
|
||||||
|
device=device,
|
||||||
|
compute_type=self.compute_type,
|
||||||
|
local_files_only=False,
|
||||||
|
)
|
||||||
|
|
||||||
def check_valid_model(self, model_size):
|
def check_valid_model(self, model_size):
|
||||||
"""
|
"""
|
||||||
Check if it's a valid whisper model size.
|
Check if it's a valid whisper model size.
|
||||||
@@ -794,6 +887,8 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
depends on the implementation of the `transcriber.transcribe` method but typically
|
depends on the implementation of the `transcriber.transcribe` method but typically
|
||||||
includes the transcribed text.
|
includes the transcribed text.
|
||||||
"""
|
"""
|
||||||
|
if ServeClientFasterWhisper.SINGLE_MODEL:
|
||||||
|
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.acquire()
|
||||||
result, info = self.transcriber.transcribe(
|
result, info = self.transcriber.transcribe(
|
||||||
input_sample,
|
input_sample,
|
||||||
initial_prompt=self.initial_prompt,
|
initial_prompt=self.initial_prompt,
|
||||||
@@ -801,6 +896,8 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
task=self.task,
|
task=self.task,
|
||||||
vad_filter=self.use_vad,
|
vad_filter=self.use_vad,
|
||||||
vad_parameters=self.vad_parameters if self.use_vad else None)
|
vad_parameters=self.vad_parameters if self.use_vad else None)
|
||||||
|
if ServeClientFasterWhisper.SINGLE_MODEL:
|
||||||
|
ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
|
||||||
|
|
||||||
if self.language is None and info is not None:
|
if self.language is None and info is not None:
|
||||||
self.set_language(info)
|
self.set_language(info)
|
||||||
@@ -882,6 +979,7 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
|
|
||||||
input_bytes, duration = self.get_audio_chunk_for_processing()
|
input_bytes, duration = self.get_audio_chunk_for_processing()
|
||||||
if duration < 1.0:
|
if duration < 1.0:
|
||||||
|
time.sleep(0.1) # wait for audio chunks to arrive
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
input_sample = input_bytes.copy()
|
input_sample = input_bytes.copy()
|
||||||
@@ -940,6 +1038,8 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
"""
|
"""
|
||||||
offset = None
|
offset = None
|
||||||
self.current_out = ''
|
self.current_out = ''
|
||||||
|
last_segment = None
|
||||||
|
|
||||||
# process complete segments
|
# process complete segments
|
||||||
if len(segments) > 1:
|
if len(segments) > 1:
|
||||||
for i, s in enumerate(segments[:-1]):
|
for i, s in enumerate(segments[:-1]):
|
||||||
@@ -955,12 +1055,14 @@ class ServeClientFasterWhisper(ServeClientBase):
|
|||||||
self.transcript.append(self.format_segment(start, end, text_))
|
self.transcript.append(self.format_segment(start, end, text_))
|
||||||
offset = min(duration, s.end)
|
offset = min(duration, s.end)
|
||||||
|
|
||||||
self.current_out += segments[-1].text
|
# only process the segments if it satisfies the no_speech_thresh
|
||||||
last_segment = self.format_segment(
|
if segments[-1].no_speech_prob <= self.no_speech_thresh:
|
||||||
self.timestamp_offset + segments[-1].start,
|
self.current_out += segments[-1].text
|
||||||
self.timestamp_offset + min(duration, segments[-1].end),
|
last_segment = self.format_segment(
|
||||||
self.current_out
|
self.timestamp_offset + segments[-1].start,
|
||||||
)
|
self.timestamp_offset + min(duration, segments[-1].end),
|
||||||
|
self.current_out
|
||||||
|
)
|
||||||
|
|
||||||
# if same incomplete segment is seen multiple times then update the offset
|
# if same incomplete segment is seen multiple times then update the offset
|
||||||
# and append the segment to the list
|
# and append the segment to the list
|
||||||
|
|||||||
+194
-29
@@ -1,22 +1,22 @@
|
|||||||
# original https://github.com/guillaumekln/faster-whisper/blob/master/faster_whisper/transcribe.py
|
# original https://github.com/guillaumekln/faster-whisper/blob/master/faster_whisper/transcribe.py
|
||||||
|
|
||||||
import itertools
|
import itertools
|
||||||
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import zlib
|
import zlib
|
||||||
import json
|
|
||||||
from inspect import signature
|
|
||||||
|
|
||||||
|
from inspect import signature
|
||||||
from typing import BinaryIO, Iterable, List, NamedTuple, Optional, Tuple, Union
|
from typing import BinaryIO, Iterable, List, NamedTuple, Optional, Tuple, Union
|
||||||
|
|
||||||
import ctranslate2
|
import ctranslate2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import tokenizers
|
import tokenizers
|
||||||
|
|
||||||
from faster_whisper.audio import decode_audio
|
from faster_whisper.audio import decode_audio, pad_or_trim
|
||||||
from faster_whisper.feature_extractor import FeatureExtractor
|
from faster_whisper.feature_extractor import FeatureExtractor
|
||||||
from faster_whisper.tokenizer import _LANGUAGE_CODES, Tokenizer
|
from faster_whisper.tokenizer import _LANGUAGE_CODES, Tokenizer
|
||||||
from faster_whisper.utils import download_model, format_timestamp, get_logger
|
from faster_whisper.utils import download_model, format_timestamp, get_end, get_logger
|
||||||
from faster_whisper.vad import (
|
from faster_whisper.vad import (
|
||||||
SpeechTimestampsMap,
|
SpeechTimestampsMap,
|
||||||
VadOptions,
|
VadOptions,
|
||||||
@@ -68,6 +68,9 @@ class TranscriptionOptions(NamedTuple):
|
|||||||
word_timestamps: bool
|
word_timestamps: bool
|
||||||
prepend_punctuations: str
|
prepend_punctuations: str
|
||||||
append_punctuations: str
|
append_punctuations: str
|
||||||
|
max_new_tokens: Optional[int]
|
||||||
|
clip_timestamps: Union[str, List[float]]
|
||||||
|
hallucination_silence_threshold: Optional[float]
|
||||||
|
|
||||||
|
|
||||||
class TranscriptionInfo(NamedTuple):
|
class TranscriptionInfo(NamedTuple):
|
||||||
@@ -96,8 +99,8 @@ class WhisperModel:
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
model_size_or_path: Size of the model to use (tiny, tiny.en, base, base.en,
|
model_size_or_path: Size of the model to use (tiny, tiny.en, base, base.en,
|
||||||
small, small.en, medium, medium.en, large-v1, large-v2, large-v3, or large), a path to a converted
|
small, small.en, medium, medium.en, large-v1, large-v2, large-v3, or large), a path to a
|
||||||
model directory, or a CTranslate2-converted Whisper model ID from the Hugging Face Hub.
|
converted model directory, or a CTranslate2-converted Whisper model ID from the HF Hub.
|
||||||
When a size or a model ID is configured, the converted model is downloaded
|
When a size or a model ID is configured, the converted model is downloaded
|
||||||
from the Hugging Face Hub.
|
from the Hugging Face Hub.
|
||||||
device: Device to use for computation ("cpu", "cuda", "auto").
|
device: Device to use for computation ("cpu", "cuda", "auto").
|
||||||
@@ -215,6 +218,10 @@ class WhisperModel:
|
|||||||
append_punctuations: str = "\"'.。,,!!??::”)]}、",
|
append_punctuations: str = "\"'.。,,!!??::”)]}、",
|
||||||
vad_filter: bool = False,
|
vad_filter: bool = False,
|
||||||
vad_parameters: Optional[Union[dict, VadOptions]] = None,
|
vad_parameters: Optional[Union[dict, VadOptions]] = None,
|
||||||
|
max_new_tokens: Optional[int] = None,
|
||||||
|
chunk_length: Optional[int] = None,
|
||||||
|
clip_timestamps: Union[str, List[float]] = "0",
|
||||||
|
hallucination_silence_threshold: Optional[float] = None,
|
||||||
) -> Tuple[Iterable[Segment], TranscriptionInfo]:
|
) -> Tuple[Iterable[Segment], TranscriptionInfo]:
|
||||||
"""Transcribes an input file.
|
"""Transcribes an input file.
|
||||||
|
|
||||||
@@ -266,6 +273,16 @@ class WhisperModel:
|
|||||||
https://github.com/snakers4/silero-vad.
|
https://github.com/snakers4/silero-vad.
|
||||||
vad_parameters: Dictionary of Silero VAD parameters or VadOptions class (see available
|
vad_parameters: Dictionary of Silero VAD parameters or VadOptions class (see available
|
||||||
parameters and default values in the class `VadOptions`).
|
parameters and default values in the class `VadOptions`).
|
||||||
|
max_new_tokens: Maximum number of new tokens to generate per-chunk. If not set,
|
||||||
|
the maximum will be set by the default max_length.
|
||||||
|
chunk_length: The length of audio segments. If it is not None, it will overwrite the
|
||||||
|
default chunk_length of the FeatureExtractor.
|
||||||
|
clip_timestamps: Union[str, List[float]]
|
||||||
|
Comma-separated list start,end,start,end,... timestamps (in seconds) of clips to
|
||||||
|
process. The last end timestamp defaults to the end of the file.
|
||||||
|
hallucination_silence_threshold: Optional[float]
|
||||||
|
When word_timestamps is True, skip silent periods longer than this threshold
|
||||||
|
(in seconds) when a possible hallucination is detected
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
A tuple with:
|
A tuple with:
|
||||||
@@ -318,7 +335,7 @@ class WhisperModel:
|
|||||||
if audio.shape[0] == 0:
|
if audio.shape[0] == 0:
|
||||||
return None, None
|
return None, None
|
||||||
|
|
||||||
features = self.feature_extractor(audio)
|
features = self.feature_extractor(audio, chunk_length=chunk_length)
|
||||||
|
|
||||||
encoder_output = None
|
encoder_output = None
|
||||||
all_language_probs = None
|
all_language_probs = None
|
||||||
@@ -384,6 +401,9 @@ class WhisperModel:
|
|||||||
word_timestamps=word_timestamps,
|
word_timestamps=word_timestamps,
|
||||||
prepend_punctuations=prepend_punctuations,
|
prepend_punctuations=prepend_punctuations,
|
||||||
append_punctuations=append_punctuations,
|
append_punctuations=append_punctuations,
|
||||||
|
max_new_tokens=max_new_tokens,
|
||||||
|
clip_timestamps=clip_timestamps,
|
||||||
|
hallucination_silence_threshold=hallucination_silence_threshold,
|
||||||
)
|
)
|
||||||
|
|
||||||
segments = self.generate_segments(features, tokenizer, options, encoder_output)
|
segments = self.generate_segments(features, tokenizer, options, encoder_output)
|
||||||
@@ -403,7 +423,7 @@ class WhisperModel:
|
|||||||
|
|
||||||
return segments, info
|
return segments, info
|
||||||
|
|
||||||
def generate_segments( # noqa: C901
|
def generate_segments(
|
||||||
self,
|
self,
|
||||||
features: np.ndarray,
|
features: np.ndarray,
|
||||||
tokenizer: Tokenizer,
|
tokenizer: Tokenizer,
|
||||||
@@ -411,8 +431,33 @@ class WhisperModel:
|
|||||||
encoder_output: Optional[ctranslate2.StorageView] = None,
|
encoder_output: Optional[ctranslate2.StorageView] = None,
|
||||||
) -> Iterable[Segment]:
|
) -> Iterable[Segment]:
|
||||||
content_frames = features.shape[-1] - self.feature_extractor.nb_max_frames
|
content_frames = features.shape[-1] - self.feature_extractor.nb_max_frames
|
||||||
|
content_duration = float(content_frames * self.feature_extractor.time_per_frame)
|
||||||
|
|
||||||
|
if isinstance(options.clip_timestamps, str):
|
||||||
|
TranscriptionOptions.clip_timestamps = [
|
||||||
|
float(ts)
|
||||||
|
for ts in (
|
||||||
|
options.clip_timestamps.split(",")
|
||||||
|
if options.clip_timestamps
|
||||||
|
else []
|
||||||
|
)
|
||||||
|
]
|
||||||
|
seek_points: List[int] = [
|
||||||
|
round(ts * self.frames_per_second) for ts in options.clip_timestamps
|
||||||
|
]
|
||||||
|
if len(seek_points) == 0:
|
||||||
|
seek_points.append(0)
|
||||||
|
if len(seek_points) % 2 == 1:
|
||||||
|
seek_points.append(content_frames)
|
||||||
|
seek_clips: List[Tuple[int, int]] = list(
|
||||||
|
zip(seek_points[::2], seek_points[1::2])
|
||||||
|
)
|
||||||
|
|
||||||
|
punctuation = "\"'“¿([{-\"'.。,,!!??::”)]}、"
|
||||||
|
|
||||||
idx = 0
|
idx = 0
|
||||||
seek = 0
|
clip_idx = 0
|
||||||
|
seek = seek_clips[clip_idx][0]
|
||||||
all_tokens = []
|
all_tokens = []
|
||||||
prompt_reset_since = 0
|
prompt_reset_since = 0
|
||||||
|
|
||||||
@@ -426,13 +471,34 @@ class WhisperModel:
|
|||||||
|
|
||||||
last_speech_timestamp = 0.0
|
last_speech_timestamp = 0.0
|
||||||
all_segments = []
|
all_segments = []
|
||||||
while seek < content_frames:
|
# NOTE: This loop is obscurely flattened to make the diff readable.
|
||||||
|
# A later commit should turn this into a simpler nested loop.
|
||||||
|
# for seek_clip_start, seek_clip_end in seek_clips:
|
||||||
|
# while seek < seek_clip_end
|
||||||
|
while clip_idx < len(seek_clips):
|
||||||
|
seek_clip_start, seek_clip_end = seek_clips[clip_idx]
|
||||||
|
if seek_clip_end > content_frames:
|
||||||
|
seek_clip_end = content_frames
|
||||||
|
if seek < seek_clip_start:
|
||||||
|
seek = seek_clip_start
|
||||||
|
if seek >= seek_clip_end:
|
||||||
|
clip_idx += 1
|
||||||
|
if clip_idx < len(seek_clips):
|
||||||
|
seek = seek_clips[clip_idx][0]
|
||||||
|
continue
|
||||||
time_offset = seek * self.feature_extractor.time_per_frame
|
time_offset = seek * self.feature_extractor.time_per_frame
|
||||||
segment = features[:, seek:seek + self.feature_extractor.nb_max_frames]
|
window_end_time = float(
|
||||||
segment_size = min(
|
(seek + self.feature_extractor.nb_max_frames)
|
||||||
self.feature_extractor.nb_max_frames, content_frames - seek
|
* self.feature_extractor.time_per_frame
|
||||||
)
|
)
|
||||||
|
segment_size = min(
|
||||||
|
self.feature_extractor.nb_max_frames,
|
||||||
|
content_frames - seek,
|
||||||
|
seek_clip_end - seek,
|
||||||
|
)
|
||||||
|
segment = features[:, seek:seek + segment_size]
|
||||||
segment_duration = segment_size * self.feature_extractor.time_per_frame
|
segment_duration = segment_size * self.feature_extractor.time_per_frame
|
||||||
|
segment = pad_or_trim(segment, self.feature_extractor.nb_max_frames)
|
||||||
|
|
||||||
if self.logger.isEnabledFor(logging.DEBUG):
|
if self.logger.isEnabledFor(logging.DEBUG):
|
||||||
self.logger.debug(
|
self.logger.debug(
|
||||||
@@ -484,10 +550,33 @@ class WhisperModel:
|
|||||||
previous_seek = seek
|
previous_seek = seek
|
||||||
current_segments = []
|
current_segments = []
|
||||||
|
|
||||||
|
# anomalous words are very long/short/improbable
|
||||||
|
def word_anomaly_score(word: dict) -> float:
|
||||||
|
probability = word.get("probability", 0.0)
|
||||||
|
duration = word["end"] - word["start"]
|
||||||
|
score = 0.0
|
||||||
|
if probability < 0.15:
|
||||||
|
score += 1.0
|
||||||
|
if duration < 0.133:
|
||||||
|
score += (0.133 - duration) * 15
|
||||||
|
if duration > 2.0:
|
||||||
|
score += duration - 2.0
|
||||||
|
return score
|
||||||
|
|
||||||
|
def is_segment_anomaly(segment: Optional[dict]) -> bool:
|
||||||
|
if segment is None or not segment["words"]:
|
||||||
|
return False
|
||||||
|
words = [w for w in segment["words"] if w["word"] not in punctuation]
|
||||||
|
words = words[:8]
|
||||||
|
score = sum(word_anomaly_score(w) for w in words)
|
||||||
|
return score >= 3 or score + 0.01 >= len(words)
|
||||||
|
|
||||||
|
def next_words_segment(segments: List[dict]) -> Optional[dict]:
|
||||||
|
return next((s for s in segments if s["words"]), None)
|
||||||
|
|
||||||
single_timestamp_ending = (
|
single_timestamp_ending = (
|
||||||
len(tokens) >= 2
|
len(tokens) >= 2
|
||||||
and tokens[-2] < tokenizer.timestamp_begin
|
and tokens[-2] < tokenizer.timestamp_begin <= tokens[-1]
|
||||||
and tokens[-1] >= tokenizer.timestamp_begin
|
|
||||||
)
|
)
|
||||||
|
|
||||||
consecutive_timestamps = [
|
consecutive_timestamps = [
|
||||||
@@ -570,18 +659,62 @@ class WhisperModel:
|
|||||||
last_speech_timestamp=last_speech_timestamp,
|
last_speech_timestamp=last_speech_timestamp,
|
||||||
)
|
)
|
||||||
|
|
||||||
word_end_timestamps = [
|
if not single_timestamp_ending:
|
||||||
w["end"] for s in current_segments for w in s["words"]
|
last_word_end = get_end(current_segments)
|
||||||
]
|
if last_word_end is not None and last_word_end > time_offset:
|
||||||
if len(word_end_timestamps) > 0:
|
seek = round(last_word_end * self.frames_per_second)
|
||||||
last_speech_timestamp = word_end_timestamps[-1]
|
|
||||||
if not single_timestamp_ending and len(word_end_timestamps) > 0:
|
|
||||||
seek_shift = round(
|
|
||||||
(word_end_timestamps[-1] - time_offset) * self.frames_per_second
|
|
||||||
)
|
|
||||||
|
|
||||||
if seek_shift > 0:
|
# skip silence before possible hallucinations
|
||||||
seek = previous_seek + seek_shift
|
if options.hallucination_silence_threshold is not None:
|
||||||
|
threshold = options.hallucination_silence_threshold
|
||||||
|
|
||||||
|
# if first segment might be a hallucination, skip leading silence
|
||||||
|
first_segment = next_words_segment(current_segments)
|
||||||
|
if first_segment is not None and is_segment_anomaly(first_segment):
|
||||||
|
gap = first_segment["start"] - time_offset
|
||||||
|
if gap > threshold:
|
||||||
|
seek = previous_seek + round(gap * self.frames_per_second)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# skip silence before any possible hallucination that is surrounded
|
||||||
|
# by silence or more hallucinations
|
||||||
|
hal_last_end = last_speech_timestamp
|
||||||
|
for si in range(len(current_segments)):
|
||||||
|
segment = current_segments[si]
|
||||||
|
if not segment["words"]:
|
||||||
|
continue
|
||||||
|
if is_segment_anomaly(segment):
|
||||||
|
next_segment = next_words_segment(
|
||||||
|
current_segments[si + 1:]
|
||||||
|
)
|
||||||
|
if next_segment is not None:
|
||||||
|
hal_next_start = next_segment["words"][0]["start"]
|
||||||
|
else:
|
||||||
|
hal_next_start = time_offset + segment_duration
|
||||||
|
silence_before = (
|
||||||
|
segment["start"] - hal_last_end > threshold
|
||||||
|
or segment["start"] < threshold
|
||||||
|
or segment["start"] - time_offset < 2.0
|
||||||
|
)
|
||||||
|
silence_after = (
|
||||||
|
hal_next_start - segment["end"] > threshold
|
||||||
|
or is_segment_anomaly(next_segment)
|
||||||
|
or window_end_time - segment["end"] < 2.0
|
||||||
|
)
|
||||||
|
if silence_before and silence_after:
|
||||||
|
seek = round(
|
||||||
|
max(time_offset + 1, segment["start"])
|
||||||
|
* self.frames_per_second
|
||||||
|
)
|
||||||
|
if content_duration - segment["end"] < threshold:
|
||||||
|
seek = content_frames
|
||||||
|
current_segments[si:] = []
|
||||||
|
break
|
||||||
|
hal_last_end = segment["end"]
|
||||||
|
|
||||||
|
last_word_end = get_end(current_segments)
|
||||||
|
if last_word_end is not None:
|
||||||
|
last_speech_timestamp = last_word_end
|
||||||
|
|
||||||
for segment in current_segments:
|
for segment in current_segments:
|
||||||
tokens = segment["tokens"]
|
tokens = segment["tokens"]
|
||||||
@@ -608,7 +741,7 @@ class WhisperModel:
|
|||||||
[Word(**word) for word in segment["words"]]
|
[Word(**word) for word in segment["words"]]
|
||||||
if options.word_timestamps
|
if options.word_timestamps
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
))
|
))
|
||||||
|
|
||||||
if (
|
if (
|
||||||
@@ -649,6 +782,21 @@ class WhisperModel:
|
|||||||
max_initial_timestamp_index = int(
|
max_initial_timestamp_index = int(
|
||||||
round(options.max_initial_timestamp / self.time_precision)
|
round(options.max_initial_timestamp / self.time_precision)
|
||||||
)
|
)
|
||||||
|
if options.max_new_tokens is not None:
|
||||||
|
max_length = len(prompt) + options.max_new_tokens
|
||||||
|
else:
|
||||||
|
max_length = self.max_length
|
||||||
|
|
||||||
|
if max_length > self.max_length:
|
||||||
|
raise ValueError(
|
||||||
|
f"The length of the prompt is {len(prompt)}, and the `max_new_tokens` "
|
||||||
|
f"{max_length - len(prompt)}. Thus, the combined length of the prompt "
|
||||||
|
f"and `max_new_tokens` is: {max_length}. This exceeds the "
|
||||||
|
f"`max_length` of the Whisper model: {self.max_length}. "
|
||||||
|
"You should either reduce the length of your prompt, or "
|
||||||
|
"reduce the value of `max_new_tokens`, "
|
||||||
|
f"so that their combined length is less that {self.max_length}."
|
||||||
|
)
|
||||||
|
|
||||||
for temperature in options.temperatures:
|
for temperature in options.temperatures:
|
||||||
if temperature > 0:
|
if temperature > 0:
|
||||||
@@ -670,7 +818,7 @@ class WhisperModel:
|
|||||||
length_penalty=options.length_penalty,
|
length_penalty=options.length_penalty,
|
||||||
repetition_penalty=options.repetition_penalty,
|
repetition_penalty=options.repetition_penalty,
|
||||||
no_repeat_ngram_size=options.no_repeat_ngram_size,
|
no_repeat_ngram_size=options.no_repeat_ngram_size,
|
||||||
max_length=self.max_length,
|
max_length=max_length,
|
||||||
return_scores=True,
|
return_scores=True,
|
||||||
return_no_speech_prob=True,
|
return_no_speech_prob=True,
|
||||||
suppress_blank=options.suppress_blank,
|
suppress_blank=options.suppress_blank,
|
||||||
@@ -728,6 +876,8 @@ class WhisperModel:
|
|||||||
if (
|
if (
|
||||||
options.no_speech_threshold is not None
|
options.no_speech_threshold is not None
|
||||||
and result.no_speech_prob > options.no_speech_threshold
|
and result.no_speech_prob > options.no_speech_threshold
|
||||||
|
and options.log_prob_threshold is not None
|
||||||
|
and avg_logprob < options.log_prob_threshold
|
||||||
):
|
):
|
||||||
needs_fallback = False # silence
|
needs_fallback = False # silence
|
||||||
|
|
||||||
@@ -738,6 +888,13 @@ class WhisperModel:
|
|||||||
decode_result = max(
|
decode_result = max(
|
||||||
below_cr_threshold_results or all_results, key=lambda x: x[1]
|
below_cr_threshold_results or all_results, key=lambda x: x[1]
|
||||||
)
|
)
|
||||||
|
# to pass final temperature for prompt_reset_on_temperature
|
||||||
|
decode_result = (
|
||||||
|
decode_result[0],
|
||||||
|
decode_result[1],
|
||||||
|
temperature,
|
||||||
|
decode_result[3],
|
||||||
|
)
|
||||||
|
|
||||||
return decode_result
|
return decode_result
|
||||||
|
|
||||||
@@ -769,7 +926,7 @@ class WhisperModel:
|
|||||||
|
|
||||||
return prompt
|
return prompt
|
||||||
|
|
||||||
def add_word_timestamps( # noqa: C901
|
def add_word_timestamps(
|
||||||
self,
|
self,
|
||||||
segments: List[dict],
|
segments: List[dict],
|
||||||
tokenizer: Tokenizer,
|
tokenizer: Tokenizer,
|
||||||
@@ -794,6 +951,7 @@ class WhisperModel:
|
|||||||
word_durations = np.array([word["end"] - word["start"] for word in alignment])
|
word_durations = np.array([word["end"] - word["start"] for word in alignment])
|
||||||
word_durations = word_durations[word_durations.nonzero()]
|
word_durations = word_durations[word_durations.nonzero()]
|
||||||
median_duration = np.median(word_durations) if len(word_durations) > 0 else 0.0
|
median_duration = np.median(word_durations) if len(word_durations) > 0 else 0.0
|
||||||
|
median_duration = min(0.7, float(median_duration))
|
||||||
max_duration = median_duration * 2
|
max_duration = median_duration * 2
|
||||||
|
|
||||||
# hack: truncate long words at sentence boundaries.
|
# hack: truncate long words at sentence boundaries.
|
||||||
@@ -915,6 +1073,13 @@ class WhisperModel:
|
|||||||
words, word_tokens = tokenizer.split_to_word_tokens(
|
words, word_tokens = tokenizer.split_to_word_tokens(
|
||||||
text_tokens + [tokenizer.eot]
|
text_tokens + [tokenizer.eot]
|
||||||
)
|
)
|
||||||
|
if len(word_tokens) <= 1:
|
||||||
|
# return on eot only
|
||||||
|
# >>> np.pad([], (1, 0))
|
||||||
|
# array([0.])
|
||||||
|
# This results in crashes when we lookup jump_times with float, like
|
||||||
|
# IndexError: arrays used as indices must be of integer (or boolean) type
|
||||||
|
return []
|
||||||
word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens[:-1]]), (1, 0))
|
word_boundaries = np.pad(np.cumsum([len(t) for t in word_tokens[:-1]]), (1, 0))
|
||||||
if len(word_boundaries) <= 1:
|
if len(word_boundaries) <= 1:
|
||||||
return []
|
return []
|
||||||
|
|||||||
@@ -35,6 +35,8 @@ class WhisperEncoding:
|
|||||||
with open(config_path, 'r') as f:
|
with open(config_path, 'r') as f:
|
||||||
config = json.load(f)
|
config = json.load(f)
|
||||||
|
|
||||||
|
use_gpt_attention_plugin = config['plugin_config'][
|
||||||
|
'gpt_attention_plugin']
|
||||||
dtype = config['builder_config']['precision']
|
dtype = config['builder_config']['precision']
|
||||||
n_mels = config['builder_config']['n_mels']
|
n_mels = config['builder_config']['n_mels']
|
||||||
num_languages = config['builder_config']['num_languages']
|
num_languages = config['builder_config']['num_languages']
|
||||||
@@ -51,12 +53,21 @@ class WhisperEncoding:
|
|||||||
return session
|
return session
|
||||||
|
|
||||||
def get_audio_features(self, mel):
|
def get_audio_features(self, mel):
|
||||||
inputs = OrderedDict()
|
|
||||||
output_list = []
|
|
||||||
|
|
||||||
inputs.update({'x': mel})
|
input_lengths = torch.tensor(
|
||||||
output_list.append(
|
[mel.shape[2] // 2 for _ in range(mel.shape[0])],
|
||||||
TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape))
|
dtype=torch.int32,
|
||||||
|
device=mel.device)
|
||||||
|
|
||||||
|
inputs = OrderedDict()
|
||||||
|
inputs['x'] = mel
|
||||||
|
inputs['input_lengths'] = input_lengths
|
||||||
|
|
||||||
|
output_list = [
|
||||||
|
TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
|
||||||
|
TensorInfo('input_lengths', str_dtype_to_trt('int32'),
|
||||||
|
input_lengths.shape)
|
||||||
|
]
|
||||||
|
|
||||||
output_info = (self.session).infer_shapes(output_list)
|
output_info = (self.session).infer_shapes(output_list)
|
||||||
|
|
||||||
@@ -101,6 +112,8 @@ class WhisperDecoding:
|
|||||||
decoder_engine_buffer = f.read()
|
decoder_engine_buffer = f.read()
|
||||||
|
|
||||||
decoder_model_config = ModelConfig(
|
decoder_model_config = ModelConfig(
|
||||||
|
max_batch_size=self.decoder_config['max_batch_size'],
|
||||||
|
max_beam_width=self.decoder_config['max_beam_width'],
|
||||||
num_heads=self.decoder_config['num_heads'],
|
num_heads=self.decoder_config['num_heads'],
|
||||||
num_kv_heads=self.decoder_config['num_heads'],
|
num_kv_heads=self.decoder_config['num_heads'],
|
||||||
hidden_size=self.decoder_config['hidden_size'],
|
hidden_size=self.decoder_config['hidden_size'],
|
||||||
@@ -141,6 +154,10 @@ class WhisperDecoding:
|
|||||||
device='cuda')
|
device='cuda')
|
||||||
decoder_max_input_length = torch.max(decoder_input_lengths).item()
|
decoder_max_input_length = torch.max(decoder_input_lengths).item()
|
||||||
|
|
||||||
|
cross_attention_mask = torch.ones(
|
||||||
|
[encoder_outputs.shape[0], 1,
|
||||||
|
encoder_outputs.shape[1]]).int().cuda()
|
||||||
|
|
||||||
# generation config
|
# generation config
|
||||||
sampling_config = SamplingConfig(end_id=eot_id,
|
sampling_config = SamplingConfig(end_id=eot_id,
|
||||||
pad_id=eot_id,
|
pad_id=eot_id,
|
||||||
@@ -161,6 +178,7 @@ class WhisperDecoding:
|
|||||||
sampling_config,
|
sampling_config,
|
||||||
encoder_output=encoder_outputs,
|
encoder_output=encoder_outputs,
|
||||||
encoder_input_lengths=encoder_input_lengths,
|
encoder_input_lengths=encoder_input_lengths,
|
||||||
|
cross_attention_mask=cross_attention_mask,
|
||||||
)
|
)
|
||||||
torch.cuda.synchronize()
|
torch.cuda.synchronize()
|
||||||
|
|
||||||
|
|||||||
+1
-1
@@ -94,7 +94,7 @@ class VoiceActivityDetection():
|
|||||||
return stacked.cpu()
|
return stacked.cpu()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def download(model_url="https://github.com/snakers4/silero-vad/raw/master/files/silero_vad.onnx"):
|
def download(model_url="https://github.com/snakers4/silero-vad/raw/v4.0/files/silero_vad.onnx"):
|
||||||
target_dir = os.path.expanduser("~/.cache/whisper-live/")
|
target_dir = os.path.expanduser("~/.cache/whisper-live/")
|
||||||
|
|
||||||
# Ensure the target directory exists
|
# Ensure the target directory exists
|
||||||
|
|||||||
Reference in New Issue
Block a user