bump version v0.4.1

Fix client close

.github/workflows/ci.yml

@@ -77,7 +77,7 @@ jobs:
   build-and-push-docker-cpu:
     needs: [run-tests, check-code-format]
     runs-on: ubuntu-22.04
-    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
     steps:
       - uses: actions/checkout@v2
       
@@ -103,7 +103,7 @@ jobs:
     needs: [run-tests, check-code-format, build-and-push-docker-cpu]
     timeout-minutes: 20
     runs-on: ubuntu-22.04
-    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
     steps:
       - uses: actions/checkout@v2
 
@@ -157,6 +157,7 @@ jobs:
         run: |
           pip install -r requirements/server.txt
           pip install -r requirements/client.txt
+          pip install wheel
 
       - name: Build package
         run: python setup.py sdist bdist_wheel

docker/Dockerfile.gpu

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
+FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
 ARG DEBIAN_FRONTEND=noninteractive
 
 # Remove any third-party apt sources to avoid issues with expiring keys.
@@ -30,4 +30,4 @@ COPY whisper_live /app/whisper_live
 
 COPY run_server.py /app
 
-CMD ["python", "run_server.py"]
+CMD ["python3", "run_server.py"]

requirements/server.txt

@@ -1,4 +1,4 @@
-faster-whisper==0.10.0
+faster-whisper==1.0.1
 torch
 websockets
 onnxruntime==1.16.0

setup.py

@@ -43,7 +43,7 @@ setup(
     ),
     install_requires=[
         "PyAudio",
-        "faster-whisper==0.10.0",
+        "faster-whisper==1.0.1",
         "torch",
         "torchaudio",
         "websockets",

tests/test_client.py

@@ -2,10 +2,12 @@ import json
 import os
 import scipy
 import websocket
+import copy
 import unittest
 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 pathlib import Path
 
 
 class BaseTestCase(unittest.TestCase):
@@ -24,6 +26,7 @@ class BaseTestCase(unittest.TestCase):
 
         self.mock_pyaudio = mock_pyaudio
         self.mock_websocket = mock_websocket
+        self.mock_audio_packet = b'\x00\x01\x02\x03'
 
     def tearDown(self):
         self.client.close_websocket()
@@ -31,7 +34,6 @@ class BaseTestCase(unittest.TestCase):
         self.mock_websocket.stop()
         del self.client
 
-
 class TestClientWebSocketCommunication(BaseTestCase):
     def test_websocket_communication(self):
         expected_url = 'ws://localhost:9090'
@@ -106,6 +108,49 @@ class TestAudioResampling(unittest.TestCase):
 
 class TestSendingAudioPacket(BaseTestCase):
     def test_send_packet(self):
-        mock_audio_packet = b'\x00\x01\x02\x03'
-        self.client.send_packet_to_server(mock_audio_packet)
-        self.client.client_socket.send.assert_called_with(mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
+        self.client.send_packet_to_server(self.mock_audio_packet)
+        self.client.client_socket.send.assert_called_with(self.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())

tests/test_server.py

@@ -9,7 +9,7 @@ import evaluate
 
 from websockets.exceptions import ConnectionClosed
 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
 
 
@@ -69,6 +69,10 @@ class TestServerConnection(unittest.TestCase):
 class TestServerInferenceAccuracy(unittest.TestCase):
     @classmethod
     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"])
         time.sleep(2)
 
@@ -77,21 +81,13 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         cls.server_process.terminate()
         cls.server_process.wait()
 
-    @mock.patch('pyaudio.PyAudio')
-    def setUp(self, mock_pyaudio):
-        self.mock_pyaudio = mock_pyaudio.return_value
-        self.mock_stream = mock.MagicMock()
-        self.mock_pyaudio.open.return_value = self.mock_stream
+    def setUp(self):
         self.metric = evaluate.load("wer")
         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!"
-        self.client("assets/jfk.flac")
-        with open("output.srt", "r") as f:
+        with open(srt_path, "r") as f:
             lines = f.readlines()
             prediction = " ".join([line.strip() for line in lines[2::4]])
         prediction_normalized = self.normalizer(prediction)
@@ -104,6 +100,23 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         )
         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):
     def setUp(self):

whisper_live/__version__.py

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

whisper_live/client.py

@@ -14,7 +14,7 @@ import whisper_live.utils as utils
 
 class Client:
     """
-    Handles audio recording, streaming, and communication with a server using WebSocket.
+    Handles communication with a server using WebSocket.
     """
     INSTANCES = {}
     END_OF_AUDIO = "END_OF_AUDIO"
@@ -42,37 +42,25 @@ class Client:
             lang (str, optional): The selected language for transcription. Default is None.
             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.task = "transcribe"
         self.uid = str(uuid.uuid4())
         self.waiting = False
-        self.last_response_recieved = None
+        self.last_response_received = None
         self.disconnect_if_no_response_for = 15
         self.language = lang
         self.model = model
         self.server_error = False
         self.srt_file_path = srt_file_path
         self.use_vad = use_vad
-        self.last_recieved_segment = None
+        self.last_segment = None
+        self.last_received_segment = None
 
         if translate:
             self.task = "translate"
 
         self.timestamp_offset = 0.0
         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:
             socket_url = f"ws://{host}:{port}"
@@ -96,7 +84,6 @@ class Client:
         self.ws_thread.setDaemon(True)
         self.ws_thread.start()
 
-        self.frames = b""
         self.transcript = []
         print("[INFO]: * recording")
 
@@ -124,10 +111,10 @@ class Client:
                       (not self.transcript or
                         float(seg['start']) >= float(self.transcript[-1]['end']))):
                     self.transcript.append(seg)
-        # update last received segment and last valild responsne time
-        if self.last_recieved_segment is None or self.last_recieved_segment != segments[-1]["text"]:
-            self.last_response_recieved = time.time()
-            self.last_recieved_segment = segments[-1]["text"]
+        # update last received segment and last valid response time
+        if self.last_received_segment is None or self.last_received_segment != segments[-1]["text"]:
+            self.last_response_received = time.time()
+            self.last_received_segment = segments[-1]["text"]
 
         # Truncate to last 3 entries for brevity.
         text = text[-3:]
@@ -162,7 +149,7 @@ class Client:
             self.recording = False
 
         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.server_backend = message["backend"]
             print(f"[INFO]: Server Running with backend {self.server_backend}")
@@ -187,7 +174,6 @@ class Client:
     def on_close(self, ws, close_status_code, close_msg):
         print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
         self.recording = False
-        self.server_error = False
         self.waiting = False
 
     def on_open(self, ws):
@@ -214,23 +200,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):
         """
         Send an audio packet to the server using WebSocket.
@@ -244,62 +213,6 @@ class Client:
         except Exception as 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):
         """
         Close the WebSocket connection and join the WebSocket thread.
@@ -327,24 +240,163 @@ class Client:
         """
         return self.client_socket
 
-    def write_audio_frames_to_file(self, frames, file_name):
+    def write_srt_file(self, output_path="output.srt"):
         """
-        Write audio frames to a WAV file.
-
-        The WAV file is created or overwritten with the specified name. The audio frames should be
-        in the correct format and match the specified channel, sample width, and sample rate.
+        Writes out the transcript in .srt format.
 
         Args:
-            frames (bytes): The audio frames to be written to the file.
-            file_name (str): The name of the WAV file to which the frames will be written.
+            message (output_path, optional): The path to the target file.  Default is "output.srt".
 
         """
-        with wave.open(file_name, "wb") as wavfile:
-            wavfile: wave.Wave_write
-            wavfile.setnchannels(self.channels)
-            wavfile.setsampwidth(2)
-            wavfile.setframerate(self.rate)
-            wavfile.writeframes(frames)
+        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):
+        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.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, hls_url=None):
+        """
+        Start the transcription process.
+
+        Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
+        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it
+        will be played and streamed to the server; otherwise, it will perform live recording.
+
+        Args:
+            audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
+
+        """
+        print("[INFO]: Waiting for server ready ...")
+        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)
+        elif audio is not None:
+            resampled_file = utils.resample(audio)
+            self.play_file(resampled_file)
+        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_hls_stream(self, hls_url):
         """
@@ -371,7 +423,7 @@ class Client:
                 if not in_bytes:
                     break
                 audio_array = self.bytes_to_float_array(in_bytes)
-                self.send_packet_to_server(audio_array.tobytes())
+                self.multicast_packet(audio_array.tobytes())
 
         except Exception as e:
             print(f"[ERROR]: Failed to connect to HLS stream: {e}")
@@ -404,14 +456,14 @@ class Client:
             os.makedirs("chunks", exist_ok=True)
         try:
             for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
-                if not self.recording:
+                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 = Client.bytes_to_float_array(data)
+                audio_array = self.bytes_to_float_array(data)
 
-                self.send_packet_to_server(audio_array.tobytes())
+                self.multicast_packet(audio_array.tobytes())
 
                 # save frames if more than a minute
                 if len(self.frames) > 60 * self.rate:
@@ -425,8 +477,7 @@ class Client:
                     t.start()
                     n_audio_file += 1
                     self.frames = b""
-            if self.server_backend == "faster_whisper":
-                self.write_srt_file(self.srt_file_path)
+            self.write_all_clients_srt()
 
         except KeyboardInterrupt:
             if len(self.frames):
@@ -437,11 +488,29 @@ class Client:
             self.stream.stop_stream()
             self.stream.close()
             self.p.terminate()
-            self.close_websocket()
+            self.close_all_clients()
 
             self.write_output_recording(n_audio_file, out_file)
-            if self.server_backend == "faster_whisper":
-                self.write_srt_file(self.srt_file_path)
+            self.write_all_clients_srt()
+
+    def write_audio_frames_to_file(self, frames, file_name):
+        """
+        Write audio frames to a WAV file.
+
+        The WAV file is created or overwritten with the specified name. The audio frames should be
+        in the correct format and match the specified channel, sample width, and sample rate.
+
+        Args:
+            frames (bytes): The audio frames to be written to the file.
+            file_name (str): The name of the WAV file to which the frames will be written.
+
+        """
+        with wave.open(file_name, "wb") as wavfile:
+            wavfile: wave.Wave_write
+            wavfile.setnchannels(self.channels)
+            wavfile.setsampwidth(2)
+            wavfile.setframerate(self.rate)
+            wavfile.writeframes(frames)
 
     def write_output_recording(self, n_audio_file, out_file):
         """
@@ -478,14 +547,26 @@ class Client:
                 os.remove(in_file)
         wavfile.close()
 
-    def write_srt_file(self, output_path="output.srt"):
-        self.transcript.append(self.last_segment)
-        utils.create_srt_file(self.transcript, output_path)
+    @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.
 
-class TranscriptionClient:
+        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(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
     to send audio data for transcription to a server and receive transcribed text segments.
@@ -508,30 +589,4 @@ class TranscriptionClient:
     """
     def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True):
         self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
-
-    def __call__(self, audio=None, hls_url=None):
-        """
-        Start the transcription process.
-
-        Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
-        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it
-        will be played and streamed to the server; otherwise, it will perform live recording.
-
-        Args:
-            audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
-
-        """
-        print("[INFO]: Waiting for server ready ...")
-        while not self.client.recording:
-            if self.client.waiting or self.client.server_error:
-                self.client.close_websocket()
-                return
-
-        print("[INFO]: Server Ready!")
-        if hls_url is not None:
-            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()
+        TranscriptionTeeClient.__init__(self, [self.client])

whisper_live/server.py

@@ -408,6 +408,11 @@ class ServeClientBase(object):
         if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
             self.frames_offset += 30.0
             self.frames_np = self.frames_np[int(30*self.RATE):]
+            # check timestamp offset(should be >= self.frame_offset)
+            # this basically means that there is no speech as timestamp offset hasnt updated
+            # and is less than frame_offset
+            if self.timestamp_offset < self.frames_offset:
+                self.timestamp_offset = self.frames_offset
         if self.frames_np is None:
             self.frames_np = frame_np.copy()
         else:
@@ -575,7 +580,7 @@ class ServeClientTensorRT(ServeClientBase):
             warmup_steps (int): Number of steps to warm up the model for.
         """
         logging.info("[INFO:] Warming up TensorRT engine..")
-        mel, _ = self.transcriber.log_mel_spectrogram("tests/jfk.flac")
+        mel, _ = self.transcriber.log_mel_spectrogram("assets/jfk.flac")
         for i in range(warmup_steps):
             self.transcriber.transcribe(mel)
 
@@ -613,7 +618,10 @@ class ServeClientTensorRT(ServeClientBase):
         """
         logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
         mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
-        last_segment = self.transcriber.transcribe(mel)
+        last_segment = self.transcriber.transcribe(
+            mel,
+            text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>"
+        )
         if last_segment:
             self.handle_transcription_output(last_segment, duration)
 
@@ -793,7 +801,8 @@ class ServeClientFasterWhisper(ServeClientBase):
             task=self.task,
             vad_filter=self.use_vad,
             vad_parameters=self.vad_parameters if self.use_vad else None)
-        if self.language is None:
+
+        if self.language is None and info is not None:
             self.set_language(info)
         return result
 
@@ -878,7 +887,9 @@ class ServeClientFasterWhisper(ServeClientBase):
                 input_sample = input_bytes.copy()
                 result = self.transcribe_audio(input_sample)
 
-                if self.language is None:
+                if result is None or self.language is None:
+                    self.timestamp_offset += duration
+                    time.sleep(0.25)    # wait for voice activity, result is None when no voice activity
                     continue
                 self.handle_transcription_output(result, duration)
 
@@ -929,7 +940,6 @@ class ServeClientFasterWhisper(ServeClientBase):
         """
         offset = None
         self.current_out = ''
-        last_segment = None
         # process complete segments
         if len(segments) > 1:
             for i, s in enumerate(segments[:-1]):

whisper_live/transcriber.py

@@ -1,22 +1,22 @@
 # original https://github.com/guillaumekln/faster-whisper/blob/master/faster_whisper/transcribe.py
 
 import itertools
+import json
 import logging
 import os
 import zlib
-import json
-from inspect import signature
 
+from inspect import signature
 from typing import BinaryIO, Iterable, List, NamedTuple, Optional, Tuple, Union
 
 import ctranslate2
 import numpy as np
 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.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 (
     SpeechTimestampsMap,
     VadOptions,
@@ -68,6 +68,9 @@ class TranscriptionOptions(NamedTuple):
     word_timestamps: bool
     prepend_punctuations: str
     append_punctuations: str
+    max_new_tokens: Optional[int]
+    clip_timestamps: Union[str, List[float]]
+    hallucination_silence_threshold: Optional[float]
 
 
 class TranscriptionInfo(NamedTuple):
@@ -96,8 +99,8 @@ class WhisperModel:
 
         Args:
           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
-            model directory, or a CTranslate2-converted Whisper model ID from the Hugging Face Hub.
+            small, small.en, medium, medium.en, large-v1, large-v2, large-v3, or large), a path to a
+            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
             from the Hugging Face Hub.
           device: Device to use for computation ("cpu", "cuda", "auto").
@@ -180,7 +183,7 @@ class WhisperModel:
 
         return config
 
-    def transcribe(
+    def transcribe(                                                         # noqa: C901
         self,
         audio: Union[str, BinaryIO, np.ndarray],
         language: Optional[str] = None,
@@ -215,6 +218,10 @@ class WhisperModel:
         append_punctuations: str = "\"'.。,，!！?？:：”)]}、",
         vad_filter: bool = False,
         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]:
         """Transcribes an input file.
 
@@ -266,6 +273,16 @@ class WhisperModel:
             https://github.com/snakers4/silero-vad.
           vad_parameters: Dictionary of Silero VAD parameters or VadOptions class (see available
             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:
           A tuple with:
@@ -315,7 +332,10 @@ class WhisperModel:
         else:
             speech_chunks = None
 
-        features = self.feature_extractor(audio)
+        if audio.shape[0] == 0:
+            return None, None
+
+        features = self.feature_extractor(audio, chunk_length=chunk_length)
 
         encoder_output = None
         all_language_probs = None
@@ -381,6 +401,9 @@ class WhisperModel:
             word_timestamps=word_timestamps,
             prepend_punctuations=prepend_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)
@@ -400,7 +423,7 @@ class WhisperModel:
 
         return segments, info
 
-    def generate_segments(                                                  # noqa: C901
+    def generate_segments(
         self,
         features: np.ndarray,
         tokenizer: Tokenizer,
@@ -408,8 +431,33 @@ class WhisperModel:
         encoder_output: Optional[ctranslate2.StorageView] = None,
     ) -> Iterable[Segment]:
         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
-        seek = 0
+        clip_idx = 0
+        seek = seek_clips[clip_idx][0]
         all_tokens = []
         prompt_reset_since = 0
 
@@ -423,13 +471,34 @@ class WhisperModel:
 
         last_speech_timestamp = 0.0
         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
-            segment = features[:, seek:seek + self.feature_extractor.nb_max_frames]
-            segment_size = min(
-                self.feature_extractor.nb_max_frames, content_frames - seek
+            window_end_time = float(
+                (seek + self.feature_extractor.nb_max_frames)
+                * 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 = pad_or_trim(segment, self.feature_extractor.nb_max_frames)
 
             if self.logger.isEnabledFor(logging.DEBUG):
                 self.logger.debug(
@@ -481,10 +550,33 @@ class WhisperModel:
             previous_seek = seek
             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 = (
                 len(tokens) >= 2
-                and tokens[-2] < tokenizer.timestamp_begin
-                and tokens[-1] >= tokenizer.timestamp_begin
+                and tokens[-2] < tokenizer.timestamp_begin <= tokens[-1]
             )
 
             consecutive_timestamps = [
@@ -567,18 +659,62 @@ class WhisperModel:
                     last_speech_timestamp=last_speech_timestamp,
                 )
 
-                word_end_timestamps = [
-                    w["end"] for s in current_segments for w in s["words"]
-                ]
-                if len(word_end_timestamps) > 0:
-                    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 not single_timestamp_ending:
+                    last_word_end = get_end(current_segments)
+                    if last_word_end is not None and last_word_end > time_offset:
+                        seek = round(last_word_end * self.frames_per_second)
 
-                    if seek_shift > 0:
-                        seek = previous_seek + seek_shift
+                # skip silence before possible hallucinations
+                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:
                 tokens = segment["tokens"]
@@ -605,7 +741,7 @@ class WhisperModel:
                         [Word(**word) for word in segment["words"]]
                         if options.word_timestamps
                         else None
-                        ),
+                    ),
                 ))
 
             if (
@@ -646,6 +782,21 @@ class WhisperModel:
         max_initial_timestamp_index = int(
             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:
             if temperature > 0:
@@ -667,7 +818,7 @@ class WhisperModel:
                 length_penalty=options.length_penalty,
                 repetition_penalty=options.repetition_penalty,
                 no_repeat_ngram_size=options.no_repeat_ngram_size,
-                max_length=self.max_length,
+                max_length=max_length,
                 return_scores=True,
                 return_no_speech_prob=True,
                 suppress_blank=options.suppress_blank,
@@ -725,6 +876,8 @@ class WhisperModel:
             if (
                 options.no_speech_threshold is not None
                 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
 
@@ -735,6 +888,13 @@ class WhisperModel:
             decode_result = max(
                 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
 
@@ -749,7 +909,7 @@ class WhisperModel:
 
         if previous_tokens:
             prompt.append(tokenizer.sot_prev)
-            prompt.extend(previous_tokens[-(self.max_length // 2 - 1):])
+            prompt.extend(previous_tokens[-(self.max_length // 2 - 1) :])
 
         prompt.extend(tokenizer.sot_sequence)
 
@@ -791,6 +951,7 @@ class WhisperModel:
         word_durations = np.array([word["end"] - word["start"] for word in alignment])
         word_durations = word_durations[word_durations.nonzero()]
         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
 
         # hack: truncate long words at sentence boundaries.
@@ -912,6 +1073,13 @@ class WhisperModel:
         words, word_tokens = tokenizer.split_to_word_tokens(
             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))
         if len(word_boundaries) <= 1:
             return []