add eos to faster_whisper server

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>

README.md

@@ -36,7 +36,7 @@ python3 run_server.py --port 9090 \
   
 # running with custom model
 python3 run_server.py --port 9090 \
-                      --backend faster_whisper
+                      --backend faster_whisper \
                       -fw "/path/to/custom/faster/whisper/model"
 ```
 

requirements/server.txt

@@ -9,4 +9,5 @@ soundfile
 ffmpeg-python
 scipy
 jiwer
-evaluate
+evaluate
+numpy<2

whisper_live/server.py

@@ -4,6 +4,9 @@ import threading
 import json
 import functools
 import logging
+from enum import Enum
+from typing import List, Optional
+
 import torch
 import numpy as np
 from websockets.sync.server import serve
@@ -121,6 +124,25 @@ class ClientManager:
         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:
     RATE = 16000
 
@@ -134,7 +156,9 @@ class TranscriptionServer:
         self, websocket, options, faster_whisper_custom_model_path,
         whisper_tensorrt_path, trt_multilingual
     ):
-        if self.backend == "tensorrt":
+        client: Optional[ServeClientBase] = None
+
+        if self.backend.is_tensorrt():
             try:
                 client = ServeClientTensorRT(
                     websocket,
@@ -155,9 +179,9 @@ class TranscriptionServer:
                     "message": "TensorRT-LLM not supported on Server yet. "
                                "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):
                 logging.info(f"Using custom model {faster_whisper_custom_model_path}")
                 options["model"] = faster_whisper_custom_model_path
@@ -174,6 +198,9 @@ class TranscriptionServer:
             )
             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)
 
     def get_audio_from_websocket(self, websocket):
@@ -202,8 +229,7 @@ class TranscriptionServer:
                 websocket.close()
                 return False  # Indicates that the connection should not continue
 
-            if self.backend == "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,
                                    whisper_tensorrt_path, trt_multilingual)
             return True
@@ -221,24 +247,22 @@ class TranscriptionServer:
         frame_np = self.get_audio_from_websocket(websocket)
         client = self.client_manager.get_client(websocket)
         if frame_np is False:
-            if self.backend == "tensorrt":
+            client.set_eos(True)
-                client.set_eos(True)
             return False
 
-        if self.backend == "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
+            client.set_eos(False)
-                client.set_eos(False)
+        if self.use_vad and not voice_active:
-            if self.use_vad and not voice_active:
+            return True
-                return True
 
         client.add_frames(frame_np)
         return True
 
     def recv_audio(self,
                    websocket,
-                   backend="faster_whisper",
+                   backend: BackendType = BackendType.FASTER_WHISPER,
                    faster_whisper_custom_model_path=None,
                    whisper_tensorrt_path=None,
                    trt_multilingual=False):
@@ -304,17 +328,14 @@ class TranscriptionServer:
             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:
+        self.single_model = single_model
-                logging.info("Custom model option was provided. Switching to single model mode.")
+        if not BackendType.is_valid(backend):
-                self.single_model = True
+            raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
-                # TODO: load model initially
-            else:
-                logging.info("Single model mode currently only works with custom models.")
         with serve(
             functools.partial(
                 self.recv_audio,
-                backend=backend,
+                backend=BackendType(backend),
                 faster_whisper_custom_model_path=faster_whisper_custom_model_path,
                 whisper_tensorrt_path=whisper_tensorrt_path,
                 trt_multilingual=trt_multilingual
@@ -387,6 +408,7 @@ class ServeClientBase(object):
         self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
         self.transcript = []
         self.send_last_n_segments = 10
+        self.eos = False
 
         # text formatting
         self.pick_previous_segments = 2
@@ -394,6 +416,18 @@ class ServeClientBase(object):
         # threading
         self.lock = threading.Lock()
 
+    def set_eos(self, eos):
+        """
+        Sets the End of Speech (EOS) flag.
+
+        Args:
+            eos (bool): The value to set for the EOS flag.
+        """
+        self.lock.acquire()
+        self.eos = eos
+        self.lock.release()
+
+
     def speech_to_text(self):
         raise NotImplementedError
 
@@ -514,7 +548,8 @@ class ServeClientBase(object):
             self.websocket.send(
                 json.dumps({
                     "uid": self.client_uid,
-                    "segments": segments,
+                    "text": segments,
+                    "eos": self.eos
                 })
             )
         except Exception as e:
@@ -619,17 +654,6 @@ class ServeClientTensorRT(ServeClientBase):
         for i in range(warmup_steps):
             self.transcriber.transcribe(mel)
 
-    def set_eos(self, eos):
-        """
-        Sets the End of Speech (EOS) flag.
-
-        Args:
-            eos (bool): The value to set for the EOS flag.
-        """
-        self.lock.acquire()
-        self.eos = eos
-        self.lock.release()
-
     def handle_transcription_output(self, last_segment, duration):
         """
         Handle the transcription output, updating the transcript and sending data to the client.
@@ -755,7 +779,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         self.task = task
         self.initial_prompt = initial_prompt
         self.vad_parameters = vad_parameters or {"threshold": 0.5}
-        self.no_speech_thresh = 0.45
+        self.no_speech_thresh = 0.35
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -767,6 +791,7 @@ class ServeClientFasterWhisper(ServeClientBase):
                 self.create_model(device)
                 ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
             else:
+                print("Re-using already initialized model.")
                 self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
         else:
             self.create_model(device)
@@ -859,8 +884,9 @@ class ServeClientFasterWhisper(ServeClientBase):
             initial_prompt=self.initial_prompt,
             language=self.language,
             task=self.task,
-            vad_filter=self.use_vad,
+            vad_filter=False,
-            vad_parameters=self.vad_parameters if self.use_vad else None)
+            vad_parameters=self.vad_parameters if self.use_vad else None,
+            beam_size=5)
         if ServeClientFasterWhisper.SINGLE_MODEL:
             ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
 
@@ -903,17 +929,16 @@ class ServeClientFasterWhisper(ServeClientBase):
             result (str): The result from whisper inference i.e. the list of segments.
             duration (float): Duration of the transcribed audio chunk.
         """
-        segments = []
         if len(result):
-            self.t_start = None
             last_segment = self.update_segments(result, duration)
-            segments = self.prepare_segments(last_segment)
-        else:
-            # show previous output if there is pause i.e. no output from whisper
-            segments = self.get_previous_output()
 
-        if len(segments):
+        if len(self.text):
-            self.send_transcription_to_client(segments)
+            if self.eos and last_segment is None:
+                self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
+                self.set_eos(False)
+                self.text = []
+            elif not self.eos:
+                self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
 
     def speech_to_text(self):
         """
@@ -943,7 +968,12 @@ class ServeClientFasterWhisper(ServeClientBase):
             self.clip_audio_if_no_valid_segment()
 
             input_bytes, duration = self.get_audio_chunk_for_processing()
-            if duration < 1.0:
+            if duration < 0.6:
+                if len(self.text) and self.eos:
+                    self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
+                    self.set_eos(False)
+                    self.text = []
+                time.sleep(0.1)
                 continue
             try:
                 input_sample = input_bytes.copy()
@@ -951,7 +981,7 @@ class ServeClientFasterWhisper(ServeClientBase):
 
                 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
+                    time.sleep(0.1)    # wait for voice activity, result is None when no voice activity
                     continue
                 self.handle_transcription_output(result, duration)
 
@@ -1000,13 +1030,13 @@ class ServeClientFasterWhisper(ServeClientBase):
             dict or None: The last processed segment with its start time, end time, and transcribed text.
                      Returns None if there are no valid segments to process.
         """
+        last_segment = None
         offset = None
         self.current_out = ''
         # process complete segments
         if len(segments) > 1:
             for i, s in enumerate(segments[:-1]):
                 text_ = s.text
-                self.text.append(text_)
                 start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
 
                 if start >= end:
@@ -1014,15 +1044,17 @@ class ServeClientFasterWhisper(ServeClientBase):
                 if s.no_speech_prob > self.no_speech_thresh:
                     continue
 
+                self.text.append(text_)
                 self.transcript.append(self.format_segment(start, end, text_))
                 offset = min(duration, s.end)
 
-        self.current_out += segments[-1].text
+        if segments[-1].no_speech_prob <= self.no_speech_thresh:
-        last_segment = self.format_segment(
+            self.current_out += segments[-1].text
-            self.timestamp_offset + segments[-1].start,
+            last_segment = self.format_segment(
-            self.timestamp_offset + min(duration, segments[-1].end),
+                self.timestamp_offset + segments[-1].start,
-            self.current_out
+                self.timestamp_offset + min(duration, segments[-1].end),
-        )
+                self.current_out
+            )
 
         # if same incomplete segment is seen multiple times then update the offset
         # and append the segment to the list
@@ -1031,7 +1063,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         else:
             self.same_output_threshold = 0
 
-        if self.same_output_threshold > 5:
+        if self.same_output_threshold > 2:
             if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
                 self.text.append(self.current_out)
                 self.transcript.append(self.format_segment(

whisper_live/vad.py

@@ -94,7 +94,7 @@ class VoiceActivityDetection():
         return stacked.cpu()
 
     @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/")
 
         # Ensure the target directory exists