document server

whisper_live/server.py

@@ -24,39 +24,76 @@ class TranscriptionServer:
     Represents a transcription server that handles incoming audio from clients.
 
     Attributes:
+        RATE (int): The audio sampling rate (constant) set to 16000.
+        vad_model (torch.Module): The voice activity detection model.
+        vad_threshold (float): The voice activity detection threshold.
         clients (dict): A dictionary to store connected clients.
+        websockets (dict): A dictionary to store WebSocket connections.
+        clients_start_time (dict): A dictionary to track client start times.
+        max_clients (int): Maximum allowed connected clients.
+        max_connection_time (int): Maximum allowed connection time in seconds.
     """
+
     RATE = 16000
+
     def __init__(self):
-        # voice activity detection model
+        """
-        self.vad_model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
+        Initialize the TranscriptionServer.
-                                           model='silero_vad',
+        """
-                                           force_reload=True,
+        # Load the voice activity detection model
-                                           onnx=True
+        self.vad_model, _ = torch.hub.load(
-                                           )
+            repo_or_dir='snakers4/silero-vad',
+            model='silero_vad',
+            force_reload=True,
+            onnx=True
+        )
+        
+        # Voice activity detection threshold
         self.vad_threshold = 0.4
+
         self.clients = {}
         self.websockets = {}
         self.clients_start_time = {}
         self.max_clients = 4
-        self.max_connection_time = 600 # in seconds
+        self.max_connection_time = 600
 
     def get_wait_time(self):
+        """
+        Calculate and return the estimated wait time for clients.
+
+        Returns:
+            float: The estimated wait time in minutes.
+        """
         wait_time = None
-        for k,v in self.clients_start_time.items():
+
+        for k, v in self.clients_start_time.items():
             current_client_time_remaining = self.max_connection_time - (time.time() - v)
-            if wait_time is None:
+
+            if wait_time is None or current_client_time_remaining < wait_time:
                 wait_time = current_client_time_remaining
-            elif current_client_time_remaining < wait_time:
+
-                wait_time = current_client_time_remaining
+        return wait_time / 60
-        return wait_time/60
 
     def recv_audio(self, websocket):
         """
         Receive audio chunks from a client in an infinite loop.
         
+        Continuously receives audio frames from a connected client
+        over a WebSocket connection. It processes the audio frames using a
+        voice activity detection (VAD) model to determine if they contain speech
+        or not. If the audio frame contains speech, it is added to the client's
+        audio data for ASR.
+        If the maximum number of clients is reached, the method sends a
+        "WAIT" status to the client, indicating that they should wait
+        until a slot is available.
+        If a client's connection exceeds the maximum allowed time, it will
+        be disconnected, and the client's resources will be cleaned up.
+
         Args:
             websocket (WebSocket): The WebSocket connection for the client.
+        
+        Raises:
+            Exception: If there is an error during the audio frame processing.
         """
         logging.info("New client connected")
         options = websocket.recv()
@@ -66,7 +103,7 @@ class TranscriptionServer:
             logging.warning("Client Queue Full. Asking client to wait ...")
             wait_time = self.get_wait_time()
             response = {
-                "uid" : options["uid"],
+                "uid": options["uid"],
                 "status": "WAIT",
                 "message": wait_time,
             }
@@ -99,6 +136,7 @@ class TranscriptionServer:
                 except Exception as e:
                     logging.error(e)
                     return
+
                 self.clients[websocket].add_frames(frame_np)
 
                 elapsed_time = time.time() - self.clients_start_time[websocket]
@@ -112,7 +150,6 @@ class TranscriptionServer:
                     del websocket
                     break
 
-
             except Exception as e:
                 logging.error(e)
                 self.clients[websocket].cleanup()
@@ -120,7 +157,6 @@ class TranscriptionServer:
                 self.clients_start_time.pop(websocket)
                 logging.info("Connection Closed.")
                 logging.info(self.clients)
-
                 del websocket
                 break
 
@@ -137,11 +173,54 @@ class TranscriptionServer:
 
 
 class ServeClient:
+    """
+    Attributes:
+        RATE (int): The audio sampling rate (constant) set to 16000.
+        SERVER_READY (str): A constant message indicating that the server is ready.
+        DISCONNECT (str): A constant message indicating that the client should disconnect.
+        client_uid (str): A unique identifier for the client.
+        data (bytes): Accumulated audio data.
+        frames (bytes): Accumulated audio frames.
+        language (str): The language for transcription.
+        task (str): The task type, e.g., "transcribe."
+        transcriber (WhisperModel): The Whisper model for speech-to-text.
+        timestamp_offset (float): The offset in audio timestamps.
+        frames_np (numpy.ndarray): NumPy array to store audio frames.
+        frames_offset (float): The offset in audio frames.
+        text (list): List of transcribed text segments.
+        current_out (str): The current incomplete transcription.
+        prev_out (str): The previous incomplete transcription.
+        t_start (float): Timestamp for the start of transcription.
+        exit (bool): A flag to exit the transcription thread.
+        same_output_threshold (int): Threshold for consecutive same output segments.
+        show_prev_out_thresh (int): Threshold for showing previous output segments.
+        add_pause_thresh (int): Threshold for adding a pause (blank) segment.
+        transcript (list): List of transcribed segments.
+        send_last_n_segments (int): Number of last segments to send to the client.
+        wrapper (textwrap.TextWrapper): Text wrapper for formatting text.
+        pick_previous_segments (int): Number of previous segments to include in the output.
+        websocket: The WebSocket connection for the client.
+    """
     RATE = 16000
     SERVER_READY = "SERVER_READY"
     DISCONNECT = "DISCONNECT"
 
     def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None, client_uid=None):
+        """
+        Initialize a ServeClient instance.
+        The Whisper model is initialized based on the client's language and device availability.
+        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
+        to the client to indicate that the server is ready.
+
+        Args:
+            websocket (WebSocket): The WebSocket connection for the client.
+            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
+            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
+            multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
+            language (str, optional): The language for transcription. Defaults to None.
+            client_uid (str, optional): A unique identifier for the client. Defaults to None.
+
+        """
         self.client_uid = client_uid
         self.data = b""
         self.frames = b""
@@ -188,14 +267,24 @@ class ServeClient:
     
     def fill_output(self, output):
         """
-        Format output with current and previous complete segments
+        Format the current incomplete transcription output by combining it with previous complete segments.
-        into two lines of 50 characters.
+        The resulting transcription is wrapped into two lines, each containing a maximum of 50 characters.
+
+        Details:
+        - This method is responsible for combining the current incomplete segment with a history of
+        previous complete segments to provide a coherent and visually organized transcription. 
+        - It ensures that the combined transcription fits within two lines, with a maximum of 50 characters per line.
+        - Segments are concatenated in the order they exist in the list of previous segments, with the most
+        recent complete segment first and older segments appended as needed to maintain the character limit.
+        - If a 3-second pause is detected in the previous segments, any text preceding it is discarded to ensure
+        the transcription starts with the most recent complete content. The resulting transcription is returned
+        as a single string.
 
         Args:
-            output(str): current incomplete segment
+            output(str): The current incomplete transcription segment.
         
         Returns:
-            transcription wrapped in two lines
+            str: A formatted transcription wrapped in two lines.
         """
         text = ''
         pick_prev = min(len(self.text), self.pick_previous_segments)
@@ -209,6 +298,26 @@ class ServeClient:
         return wrapped
     
     def add_frames(self, frame_np):
+        """
+        Add audio frames to the ongoing audio stream buffer.
+
+        This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
+        of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
+        to prevent excessive memory usage.
+
+        Details:
+        - The method appends incoming audio frames to the ongoing audio stream buffer.
+        - If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
+          of audio data to maintain a reasonable buffer size.
+        - If the buffer is empty, it initializes it with the provided audio frame.
+        - The audio stream buffer is used for real-time processing of audio data for transcription.
+
+        Args:
+            frame_np (numpy.ndarray): The audio frame data as a NumPy array.
+
+        Returns:
+            None
+        """
         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):]
@@ -219,7 +328,24 @@ class ServeClient:
 
     def speech_to_text(self):
         """
-        Process audio stream in an infinite loop.
+        Process an audio stream in an infinite loop, continuously transcribing the speech.
+
+        This method continuously receives audio frames, performs real-time transcription, and sends
+        transcribed segments to the client via a WebSocket connection.
+
+        Details:
+        - If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
+        - It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results.
+        - Segments are sent to the client in real-time, and a history of segments is maintained to provide context.
+        - Pauses in speech (no output from Whisper) are handled by showing the previous output for a set duration.
+        - A blank segment is added if there is no speech for a specified duration to indicate a pause.
+
+        Returns:
+            None
+
+        Raises:
+            Exception: If there is an issue with audio processing or WebSocket communication.
+
         """
         # detect language
         if self.language is None:
@@ -327,12 +453,27 @@ class ServeClient:
         Processes the segments from whisper. Appends all the segments to the list
         except for the last segment assuming that it is incomplete.
 
+        This method takes segments obtained from the Whisper, processes them, and updates the
+        ongoing transcript with the transcribed text. It handles complete segments, incomplete
+        segments, and repeated segments while maintaining chronological order.
+
+        Details:
+        - The method updates the ongoing transcript with transcribed segments, including their start and end times.
+        - Complete segments are appended to the transcript in chronological order.
+        - Incomplete segments (assumed to be the last one) are processed to identify repeated content. If
+          the same incomplete segment is seen multiple times, it updates the offset and appends the segment
+          to the transcript.
+        - A threshold is used to detect repeated content and ensure it is only included once in the transcript.
+        - The timestamp offset is updated based on the duration of processed segments.
+        - The method returns the last processed segment, allowing it to be sent to the client for real-time updates.
+
         Args:
             segments(dict) : dictionary of segments as returned by whisper
             duration(float): duration of the current chunk
         
         Returns:
-            transcription for the current chunk
+            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.
         """
         offset = None
         self.current_out = ''
@@ -391,6 +532,15 @@ class ServeClient:
         return last_segment
     
     def disconnect(self):
+        """
+        Notify the client of disconnection and send a disconnect message.
+
+        This method sends a disconnect message to the client via the WebSocket connection to notify them
+        that the transcription service is disconnecting gracefully.
+
+        Returns:
+            None
+        """
         self.websocket.send(
             json.dumps(
                 {
@@ -401,6 +551,16 @@ class ServeClient:
         )
     
     def cleanup(self):
+        """
+        Perform cleanup tasks before exiting the transcription service.
+
+        This method performs necessary cleanup tasks, including stopping the transcription thread, marking
+        the exit flag to indicate the transcription thread should exit gracefully, and destroying resources
+        associated with the transcription process.
+
+        Returns:
+            None
+        """
         logging.info("Cleaning up.")
         self.exit = True
         self.transcriber.destroy()