From c574f5d6a600044e0f9114991e8f216f77fb6c03 Mon Sep 17 00:00:00 2001 From: makaveli10 Date: Fri, 15 Sep 2023 11:34:18 +0530 Subject: [PATCH] document server --- whisper_live/server.py | 212 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 186 insertions(+), 26 deletions(-) diff --git a/whisper_live/server.py b/whisper_live/server.py index aa4439a..7a9d4ca 100644 --- a/whisper_live/server.py +++ b/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', - model='silero_vad', - force_reload=True, - onnx=True - ) + """ + Initialize the TranscriptionServer. + """ + # Load the voice activity detection model + 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, } @@ -74,7 +111,7 @@ class TranscriptionServer: websocket.close() del websocket return - + client = ServeClient( websocket, multilingual=options["multilingual"], @@ -82,9 +119,9 @@ class TranscriptionServer: task=options["task"], client_uid=options["uid"] ) - + self.clients[websocket] = client - self.clients_start_time[websocket] = time.time() + self.clients_start_time[websocket] = time.time() while True: try: @@ -95,10 +132,11 @@ class TranscriptionServer: speech_prob = self.vad_model(torch.from_numpy(frame_np.copy()), self.RATE).item() if speech_prob < self.vad_threshold: continue - + 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 - into two lines of 50 characters. + Format the current incomplete transcription output by combining it with previous complete segments. + 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()