add multiprocess communication

services/llm_service.py → llm_service.py

@@ -103,6 +103,7 @@ class MistralTensorRTLLM:
                          rank=self.runtime_rank,
                          debug_mode=False,
                          lora_ckpt_source='hf')
+        self.runner = self.runner_cls.from_dir(**self.runner_kwargs)
 
     def parse_input(
         self,
@@ -152,75 +153,83 @@ class MistralTensorRTLLM:
                 output.append(output_text)
         return output
     
-    def __call__(
-        self, 
-        input_text, 
-        max_output_len=100, 
+    def run(
+        self,
+        transcription_queue=None,
+        llm_queue=None,
+        input_text=None, 
+        max_output_len=20, 
         max_attention_window_size=4096, 
         num_beams=1, 
         streaming=True,
         streaming_interval=4,
+        debug=False,
     ):
-        import time
-        start = time.time()
-        batch_input_ids = self.parse_input(
-            input_text=input_text,
-            add_special_tokens=True,
-            max_input_length=923,
-            pad_id=None,
+        self.initialize_model(
+        "/root/TensorRT-LLM/examples/llama/tmp/mistral/7B/trt_engines/fp16/1-gpu",
+        "teknium/OpenHermes-2.5-Mistral-7B",
         )
+        print("Loaded LLM...")
+        while True:
+ 
+            # while transcription
+            transcription_output = transcription_queue.get()
+            input_text=transcription_output['prompt'].strip()
+            
+            print("Whisper: ", input_text)
+            batch_input_ids = self.parse_input(
+                input_text=input_text,
+                add_special_tokens=True,
+                max_input_length=923,
+                pad_id=None,
+            )
 
-        input_lengths = [x.size(1) for x in batch_input_ids]
-        print(self.runner_kwargs)
-        runner = self.runner_cls.from_dir(**self.runner_kwargs)
-        with torch.no_grad():
-            outputs = runner.generate(
-                batch_input_ids,
-                max_new_tokens=max_output_len,
-                max_attention_window_size=max_attention_window_size,
-                end_id=self.end_id,
-                pad_id=self.pad_id,
-                temperature=1.0,
-                top_k=1,
-                top_p=0.0,
-                num_beams=num_beams,
-                length_penalty=1.0,
-                repetition_penalty=1.0,
-                stop_words_list=None,
-                bad_words_list=None,
-                lora_uids=None,
-                prompt_table_path=None,
-                prompt_tasks=None,
-                streaming=streaming,
-                output_sequence_lengths=True,
-                return_dict=True)
-            torch.cuda.synchronize()
-        print(outputs)
-        if streaming:
-            for curr_outputs in throttle_generator(outputs, streaming_interval):
-                output_ids = curr_outputs['output_ids']
-                sequence_lengths = curr_outputs['sequence_lengths']
+            input_lengths = [x.size(1) for x in batch_input_ids]
+            with torch.no_grad():
+                outputs = self.runner.generate(
+                    batch_input_ids,
+                    max_new_tokens=max_output_len,
+                    max_attention_window_size=max_attention_window_size,
+                    end_id=self.end_id,
+                    pad_id=self.pad_id,
+                    temperature=1.0,
+                    top_k=1,
+                    top_p=0.0,
+                    num_beams=num_beams,
+                    length_penalty=1.0,
+                    repetition_penalty=1.0,
+                    stop_words_list=None,
+                    bad_words_list=None,
+                    lora_uids=None,
+                    prompt_table_path=None,
+                    prompt_tasks=None,
+                    streaming=streaming,
+                    output_sequence_lengths=True,
+                    return_dict=True)
+                torch.cuda.synchronize()
+            if streaming:
+                for curr_outputs in throttle_generator(outputs, streaming_interval):
+                    output_ids = curr_outputs['output_ids']
+                    sequence_lengths = curr_outputs['sequence_lengths']
+                    output = self.decode_tokens(
+                        output_ids,
+                        input_lengths,
+                        sequence_lengths
+                    )
+            else:
+                output_ids = outputs['output_ids']
+                sequence_lengths = outputs['sequence_lengths']
+                context_logits = None
+                generation_logits = None
+                if runner.gather_all_token_logits:
+                    context_logits = outputs['context_logits']
+                    generation_logits = outputs['generation_logits']
                 output = self.decode_tokens(
                     output_ids,
                     input_lengths,
-                    sequence_lengths
+                    sequence_lengths,
                 )
-                print(time.time() - start)
-                print(input_text[0] + " " + output[0])
-        else:
-            output_ids = outputs['output_ids']
-            sequence_lengths = outputs['sequence_lengths']
-            context_logits = None
-            generation_logits = None
-            if runner.gather_all_token_logits:
-                context_logits = outputs['context_logits']
-                generation_logits = outputs['generation_logits']
-            output = self.decode_tokens(
-                output_ids,
-                input_lengths,
-                sequence_lengths,
-            )
-        return output
+            llm_queue.put({"uid": transcription_output["uid"], "llm_output": output})
 
 
 if __name__=="__main__":

main.py

@@ -0,0 +1,35 @@
+from whisper_live.trt_server import TranscriptionServer
+from llm_service import MistralTensorRTLLM
+import multiprocessing
+import threading
+import ssl
+import time
+import sys
+import functools
+
+from multiprocessing import Process, Manager, Value, Queue
+
+
+if __name__ == "__main__":
+    # multiprocessing.set_start_method('spawn')
+    
+    lock = multiprocessing.Lock()
+    
+    manager = Manager()
+    shared_output = manager.list()
+
+    transcription_queue = Queue()
+    llm_queue = Queue()
+
+
+    whisper_server = TranscriptionServer()
+    whisper_process = multiprocessing.Process(target=whisper_server.run, args=("0.0.0.0", 6006, transcription_queue, llm_queue))
+    whisper_process.start()
+
+    # llm_provider = MistralTensorRTLLM()
+    # # llm_provider = MistralTensorRTLLMProvider()
+    # llm_process = multiprocessing.Process(target=llm_provider.run, args=(transcription_queue, llm_queue))
+    # llm_process.start()
+
+    # llm_process.join()
+    whisper_process.join()

whisper_live/trt_server.py

@@ -12,6 +12,8 @@ from websockets.sync.server import serve
 import torch
 import numpy as np
 import time
+import queue
+
 from whisper_live.vad import VoiceActivityDetection
 from whisper_live.trt_transcriber import WhisperTRTLLM
 
@@ -47,13 +49,13 @@ class TranscriptionServer:
 
     def __init__(self):
         # voice activity detection model
-        self.vad_model = VoiceActivityDetection()
-        self.vad_threshold = 0.5
+        
         self.clients = {}
         self.websockets = {}
         self.clients_start_time = {}
         self.max_clients = 4
         self.max_connection_time = 600
+        print("done loading")
 
     def get_wait_time(self):
         """
@@ -72,7 +74,7 @@ class TranscriptionServer:
 
         return wait_time / 60
 
-    def recv_audio(self, websocket):
+    def recv_audio(self, websocket, transcription_queue=None, llm_queue=None):
         """
         Receive audio chunks from a client in an infinite loop.
         
@@ -93,6 +95,9 @@ class TranscriptionServer:
         Raises:
             Exception: If there is an error during the audio frame processing.
         """
+        self.vad_model = VoiceActivityDetection()
+        self.vad_threshold = 0.5
+
         logging.info("New client connected")
         options = websocket.recv()
         options = json.loads(options)
@@ -115,22 +120,26 @@ class TranscriptionServer:
             multilingual=options["multilingual"],
             language=options["language"],
             task=options["task"],
-            client_uid=options["uid"]
+            client_uid=options["uid"],
+            transcription_queue=transcription_queue,
+            llm_queue=llm_queue
         )
 
         self.clients[websocket] = client
         self.clients_start_time[websocket] = time.time()
         no_voice_activity_chunks = 0
+        print()
         while True:
             try:
                 frame_data = websocket.recv()
                 frame_np = np.frombuffer(frame_data, dtype=np.float32)
+                print(frame_np.shape)
                 # VAD
                 try:
                     speech_prob = self.vad_model(torch.from_numpy(frame_np.copy()), self.RATE).item()
                     if speech_prob < self.vad_threshold:
                         no_voice_activity_chunks += 1
-                        print("No speech", no_voice_activity_chunks)
+                        print("No speech", no_voice_activity_chunks, self.clients[websocket].eos)
                         if no_voice_activity_chunks > 2:
                             if not self.clients[websocket].eos:
                                 self.clients[websocket].set_eos(True)
@@ -164,7 +173,7 @@ class TranscriptionServer:
                 del websocket
                 break
 
-    def run(self, host, port=9090):
+    def run(self, host, port=9090, transcription_queue=None, llm_queue=None):
         """
         Run the transcription server.
 
@@ -172,7 +181,15 @@ class TranscriptionServer:
             host (str): The host address to bind the server.
             port (int): The port number to bind the server.
         """
-        with serve(self.recv_audio, host, port) as server:
+        with serve(
+            functools.partial(
+                self.recv_audio, 
+                transcription_queue=transcription_queue, 
+                llm_queue=llm_queue,
+            ),
+            host,
+            port
+        ) as server:
             server.serve_forever()
 
 
@@ -209,7 +226,17 @@ class ServeClient:
     SERVER_READY = "SERVER_READY"
     DISCONNECT = "DISCONNECT"
 
-    def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None, client_uid=None):
+    def __init__(
+        self,
+        websocket,
+        task="transcribe",
+        device=None,
+        multilingual=False,
+        language=None, 
+        client_uid=None,
+        transcription_queue=None,
+        llm_queue=None,
+        ):
         """
         Initialize a ServeClient instance.
         The Whisper model is initialized based on the client's language and device availability.
@@ -226,6 +253,8 @@ class ServeClient:
 
         """
         self.client_uid = client_uid
+        self.transcription_queue = transcription_queue
+        self.llm_queue = llm_queue
         self.data = b""
         self.frames = b""
         self.language = language if multilingual else "en"
@@ -246,6 +275,7 @@ class ServeClient:
         self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
         self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
         self.transcript = []
+        self.prompt = None
         self.send_last_n_segments = 10
 
         # text formatting
@@ -318,6 +348,14 @@ class ServeClient:
 
         """
         while True:
+            try:
+                if self.llm_queue is not None:
+                    llm_output = self.llm_queue.get_nowait()
+                    if llm_output:
+                        self.websocket.send(json.dumps(llm_output))
+            except queue.Empty:
+                pass
+            
             if self.exit:
                 logging.info("Exiting speech to text thread")
                 break
@@ -334,28 +372,39 @@ class ServeClient:
             samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE)
             input_bytes = self.frames_np[int(samples_take):].copy()
             duration = input_bytes.shape[0] / self.RATE
-            if duration<1.0 or not self.eos:
+            if duration<1.0:
                 continue
 
             try:
                 input_sample = input_bytes.copy()
-                save_wav(input_sample)
+                # save_wav(input_sample)
                 # whisper transcribe with prompt
                 mel, duration = self.transcriber.log_mel_spectrogram(input_sample)
-                print(mel.shape, duration)
-                result = self.transcriber.transcribe(mel)
-                self.append_segment(result)
-                self.set_eos(False)
-                self.timestamp_offset += duration
-                if len(result):
-                    segments = self.transcript[-self.send_last_n_segments:]
+                last_segment = self.transcriber.transcribe(mel)
+                
+                if len(last_segment):
+                    if len(self.transcript) < self.send_last_n_segments:
+                        segments = self.transcript
+                    else:
+                        segments = self.transcript[-self.send_last_n_segments:]
+                    segments.append({"text": last_segment})
                     try:
+                        print(f"Sending... {segments}")
                         self.websocket.send(
                             json.dumps({
                                 "uid": self.client_uid,
-                                "segments": segments
+                                "segments": segments,
+                                "eos": self.eos
                             })
                         )
+                        if self.eos:
+                            self.append_segment(last_segment)
+                            self.timestamp_offset += duration
+                            self.prompt = ' '.join(segment['text'] for segment in segments)
+                            self.transcription_queue.put({"uid": self.client_uid, "prompt": self.prompt})
+                            self.transcript = []
+                            self.set_eos(False)
+
                     except Exception as e:
                         logging.error(f"[ERROR]: {e}")
 
@@ -364,6 +413,7 @@ class ServeClient:
                 time.sleep(0.01)
     
     def append_segment(self, result):
+        print("adding to trasncript: ", result)
         if not len(self.transcript):
             self.transcript.append({"text": result + " "})
         else:

whisper_live/trt_transcriber.py

@@ -339,9 +339,9 @@ def decode_wav_file(
 
 if __name__=="__main__":
     tensorrt_llm.logger.set_level("error")
-    model = WhisperTRTLLM("../whisper_small_en", False, "../assets", device="cuda")
+    model = WhisperTRTLLM("/root/TensorRT-LLM/examples/whisper/whisper_small_en", False, "../assets", device="cuda")
     mel, total_duration = model.log_mel_spectrogram(
-        "/root/Code/outputs/output3.wav",
+        "../assets/1221-135766-0002.wav",
     )
     results = model.transcribe(mel)
     print(results, total_duration)

whisper_live/vad.py

@@ -10,19 +10,23 @@ import onnxruntime
 class VoiceActivityDetection():
 
     def __init__(self, force_onnx_cpu=True):
+        print("downloading ONNX model...")
         path = self.download()
+        print("loading session")
+
         opts = onnxruntime.SessionOptions()
         opts.log_severity_level = 3
 
         opts.inter_op_num_threads = 1
         opts.intra_op_num_threads = 1
 
+        print("loading onnx model")
         if force_onnx_cpu and 'CPUExecutionProvider' in onnxruntime.get_available_providers():
             self.session = onnxruntime.InferenceSession(path, providers=['CPUExecutionProvider'], sess_options=opts)
         else:
             self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts)
 
-
+        print("reset states")
         self.reset_states()
         self.sample_rates = [8000, 16000]