change to pipelines multiple processing

config.py

@@ -13,6 +13,7 @@ logging.getLogger("pywhispercpp").setLevel(logging.WARNING)
 
 BASE_DIR = pathlib.Path(__file__).parent
 MODEL_DIR = BASE_DIR / "moyoyo_asr_models"
+ASSERT_DIR = BASE_DIR / "assets"
 # 标点
 SENTENCE_END_MARKERS =  ['.', '!', '?', '。', '！', '？', ';', '；', ':', '：']
 PAUSE_END_MARKERS = [',', '，', '、']

transcribe/pipelines/base.py

@@ -1,6 +1,7 @@
 
 from dataclasses import dataclass, field
-from multiprocessing import Process
+from multiprocessing import Process, Queue
+from multiprocessing import Event
 
 @dataclass
 class Segment:
@@ -19,10 +20,20 @@ class MetaItem:
     
 
 class BasePipe(Process):
-    def __init__(self, in_queue, out_queue) -> None:
+    def __init__(self, in_queue=None, out_queue=None) -> None:
         super().__init__()  # Initialize the Process class
-        self._in_queue = in_queue
-        self._out_queue = out_queue
+        self._in_queue = in_queue if in_queue else Queue()
+        self._out_queue = out_queue if out_queue else  Queue()
+        self._ready = Event()
+
+    def set_ready(self):
+        self._ready.set()
+    
+    def is_ready(self):
+        return self._ready.is_set()
+    
+    def wait(self):
+        self._ready.wait()
 
     @property
     def output_queue(self):
@@ -42,9 +53,10 @@ class BasePipe(Process):
 
     def run(self):
         self.init()
+        self.set_ready()
         while True:
-            item = self._in_queue.get()
+            item = self.input_queue.get()
             if item is None:  # Check for termination signal
                 break
             out_item = self.process(item)
-            self._out_queue.put(out_item)
+            self.output_queue.put(out_item)

transcribe/pipelines/pipe_translate.py

@@ -4,6 +4,7 @@ from llama_cpp import Llama
 from ..translator import QwenTranslator
 from config import LLM_MODEL_PATH, LLM_SYS_PROMPT
 
+
 class TranslatePipe(BasePipe):
     translator = None
      

transcribe/pipelines/pipe_whisper.py

@@ -7,7 +7,7 @@ from ..whisper import WhisperCPP
 class WhisperPipe(BasePipe):
     whisper = None
 
-    def __init__(self, in_queue, out_queue) -> None:
+    def __init__(self, in_queue=None, out_queue=None) -> None:
         super().__init__(in_queue, out_queue)
 
 

transcribe/transcription.py

@@ -176,8 +176,8 @@ class TranscriptionServer:
         frame_data = websocket.recv()
         if frame_data == b"END_OF_AUDIO":
             return False
-        return np.frombuffer(frame_data, dtype=np.int16).astype(np.float32) / 32768.0
-        # return np.frombuffer(frame_data, dtype=np.float32)
+        # return np.frombuffer(frame_data, dtype=np.int16).astype(np.float32) / 32768.0
+        return np.frombuffer(frame_data, dtype=np.float32)
 
 
     def handle_new_connection(self, websocket):

transcribe/translatepipes.py

@@ -5,21 +5,16 @@ import config
 class TranslatePipes:
     def __init__(self) -> None:
         
-        self.whisper_input_q = mp.Queue() 
-        self.translate_input_q = mp.Queue()
-        self.result_queue = mp.Queue()
+        # self.whisper_input_q = mp.Queue() 
+        # self.translate_input_q = mp.Queue()
+        # self.result_queue = mp.Queue()
         
         # whisper 转录
-        self._whisper_pipe = WhisperPipe(
-            in_queue=self.whisper_input_q,
-            out_queue=self.translate_input_q
-            )
+        self._whisper_pipe = WhisperPipe()
         
         # llm 翻译
-        self._translate_pipe = TranslatePipe(
-            in_queue=self.translate_input_q,
-            out_queue=self.result_queue,
-        )
+        self._translate_pipe = TranslatePipe()
+
         
         self._whisper_pipe.daemon = True
         self._whisper_pipe.start()
@@ -27,12 +22,15 @@ class TranslatePipes:
         self._translate_pipe.daemon = True
         self._translate_pipe.start()
 
+    def wait_ready(self):
+        self._whisper_pipe.wait()
+        self._translate_pipe.wait()
     
     def translate(self, text, src_lang, dst_lang) -> MetaItem:
         item = MetaItem(
             transcribe_content=text,
-              source_language=src_lang, 
-              destination_language=dst_lang)
+            source_language=src_lang, 
+            destination_language=dst_lang)
         self._translate_pipe.input_queue.put(item)
         return self._translate_pipe.output_queue.get()
     

transcribe/translator.py

@@ -8,12 +8,8 @@ class QwenTranslator:
     def __init__(self, model_path, system_prompt="") -> None:
         self.llm = Llama(
         model_path=model_path,
-        # n_gpu_layers=-1, # Uncomment to use GPU acceleration
-        # seed=1337, # Uncomment to set a specific seed
-        # n_ctx=2048, # Uncomment to increase the context window
         chat_format="chatml",
-        verbose=False
-)
+        verbose=False)
         self.sys_prompt = system_prompt
 
     def to_message(self, prompt, src_lang, dst_lang):
@@ -26,9 +22,7 @@ class QwenTranslator:
 
     def translate(self, prompt, src_lang, dst_lang) -> str:
         message = self.to_message(prompt, src_lang, dst_lang)
-        start_time = time.monotonic()
-        output = self.llm.create_chat_completion(messages=message, temperature=0.9)
-        
+        output = self.llm.create_chat_completion(messages=message, temperature=0)
         return output['choices'][0]['message']['content']
     
     def __call__(self, prompt,*args, **kwargs):

transcribe/whisper.py

@@ -20,7 +20,7 @@ class WhisperCPP:
 
     
     def warmup(cls, warmup_steps=1):
-        mel, _, = soundfile.read("assets/jfk.flac")
+        mel, _, = soundfile.read(f"{config.ASSERT_DIR}/jfk.flac")
         for _ in range(warmup_steps):
             cls.model.transcribe(mel, print_progress=False)
 
@@ -35,8 +35,7 @@ class WhisperCPP:
     def transcribe(self, audio_buffer:bytes, language):
         max_len, prompt = self.config_language(language)
         audio_buffer = np.frombuffer(audio_buffer, dtype=np.float32)
-        print("audio buffer got:", len(audio_buffer))
-        output =  self.model.transcribe(
+        output = self.model.transcribe(
             audio_buffer,
             initial_prompt=prompt,
             language=language,  

transcribe/whisper_llm_serve.py

@@ -20,9 +20,15 @@ from queue import Queue
 from scipy.io.wavfile import write
 from api_model import TransResult, Message
 from .utils import log_block
+from .translatepipes import TranslatePipes
 
 logger = getLogger("TranslatorApp")
 
+translate_pipes = TranslatePipes()
+translate_pipes.wait_ready()
+
+logger.info("Pipeline is ready.")
+
 def save_to_wave(filename, data:np.ndarray, sample_rate=16000):
     write(filename, sample_rate, data)
 
@@ -30,9 +36,6 @@ class TripleTextBuffer:
     def __init__(self, size=2):
         self.history = collections.deque(maxlen=size)
 
-    def _clean(self):
-        self.history.clear()
-    
     def add_entry(self, text, index):
         """
         text: 文本
@@ -56,7 +59,7 @@ class TripleTextBuffer:
         # print("比较： ", text1, text2," => ", sim_12)
         # sim_23 = self.text_similarity(text2, text3)
         if sim_12 >= similarity_threshold:
-            self._clean()
+            self.history.clear()
             return idx2
         return None
 
@@ -112,71 +115,6 @@ class SegmentManager:
             self.commit_segment()
         
 
-class PywhisperInference:
-    whisper_model = None
-    llm_model = None
-    # vad_model = None
- 
-    @classmethod
-    def initializer(cls, event:mp.Event, warmup=True):
-        models_dir = config.MODEL_DIR.as_posix()
-        cls.whisper_model = Model(
-            model=config.WHISPER_MODEL, 
-            models_dir=models_dir,
-            print_realtime=False,
-            print_progress=False,
-            print_timestamps=False,
-        )
-        if warmup:
-            cls.warmup()
-        
-        # init llamacpp
-        cls.llm_model = QwenTranslator(config.LLM_MODEL_PATH, config.LLM_SYS_PROMPT)
-        # cls.vad_model = VoiceActivityDetector()
-        event.set()
-     
-    
-    @classmethod
-    def init(cls):
-        pass
-        
-    
-    @classmethod
-    def warmup(cls, warmup_steps=1):
-        mel, _, = soundfile.read("assets/jfk.flac")
-        for _ in range(warmup_steps):
-            cls.whisper_model.transcribe(mel, print_progress=False)
-
-    @staticmethod
-    def config_language(language):
-        if language == "zh":
-            return config.MAX_LENTH_ZH, config.WHISPER_PROMPT_ZH
-        elif language == "en":
-            return config.MAX_LENGTH_EN, config.WHISPER_PROMPT_EN
-        raise ValueError(f"Unsupported language : {language}")
-
-    @classmethod
-    def transcribe(cls, audio_buffer, language):
-        max_len, prompt = cls.config_language(language)
-        audio_buffer = np.frombuffer(audio_buffer, dtype=np.float32)
-        return cls.whisper_model.transcribe(
-            audio_buffer,
-            initial_prompt=prompt,
-            language=language,  
-            token_timestamps=True, 
-            max_len=max_len
-        )
-    
-    @classmethod
-    def translate(cls, context: str, src_lang, dst_lang):
-        return cls.llm_model.translate(context, src_lang, dst_lang)
-    
-    # @classmethod
-    # def voice_detect(cls, audio_buffer):
-    #     audio_buffer = np.frombuffer(audio_buffer, dtype=np.float32)
-    #     return cls.vad_model(audio_buffer)
-
-
 
 class PyWhiperCppServe(ServeClientBase):
    
@@ -189,30 +127,23 @@ class PyWhiperCppServe(ServeClientBase):
         self._text_buffer = TripleTextBuffer()
         # 存储转录数据
         self._segment_manager = SegmentManager()
+        self._ready_state = mp.Event()
+        
         self.lock = threading.Lock()
         self.frames_np = None
         self.sample_rate = 16000
         # self._audio_queue = Queue()
         # 进程初始化后再开始收音频
-        self._ready_state = mp.Event()
-        self._pool = PPool(
-            max_workers=1, 
-            initializer=PywhisperInference.initializer, 
-            initargs=(self._ready_state,)
-        )
-        self._pool.submit(PywhisperInference.init)
+       
         logger.info('Create a process to process audio.')
         self.send_ready_state()
-        # self.load_frame_thread = threading.Thread(target=self.load_frame_from_queue)
-        # self.load_frame_thread.daemon = True
-        # self.load_frame_thread.start()
         
         self.trans_thread = threading.Thread(target=self.speech_to_text)
         self.trans_thread.daemon = True
         self.trans_thread.start()
 
     def send_ready_state(self):
-        self._ready_state.wait()
+        # self._ready_state.wait()
         self.websocket.send(json.dumps({
             "uid": self.client_uid,
             "message": self.SERVER_READY,
@@ -223,20 +154,6 @@ class PyWhiperCppServe(ServeClientBase):
         self.language = src_lang
         self.dst_lang = dst_lang
 
-    # def load_frame_from_queue(self):
-    #     while True:
-    #         frame_np = self._audio_queue.get()
-    #         fut = self._pool.submit(PywhisperInference.voice_detect, frame_np.tobytes())
-    #         output = fut.result()
-    #         logger.info(f"VAD: {output}")
-    #         if output == True:
-    #             with self.lock:
-    #                 if self.frames_np is None:
-    #                     self.frames_np = frame_np.copy()
-    #                 else:
-    #                     self.frames_np = np.append(self.frames_np,frame_np)
-    
-
     def add_frames(self, frame_np):
         # self._audio_queue.put(frame_np)
         with self.lock:
@@ -261,10 +178,8 @@ class PyWhiperCppServe(ServeClientBase):
         log_block("Audio buffer length", f"{audio_buffer.shape[0]/self.sample_rate:.2f}", "s")
         start_time = time.perf_counter()
 
-        transcribe_fut = self._pool.submit(
-            PywhisperInference.transcribe, audio_buffer.tobytes(), self.language)
-
-        segments = transcribe_fut.result()
+        item = translate_pipes.transcrible(audio_buffer.tobytes(), self.language)
+        segments = item.segments
         log_block("Whisper transcrible time", f"{(time.perf_counter() - start_time):.3f}", "s")
        
         return segments
@@ -275,12 +190,9 @@ class PyWhiperCppServe(ServeClientBase):
         # return "sample english"
         log_block("LLM translate input", f"{text}")
         start_time = time.perf_counter()
-        translate_fut = self._pool.submit(
-            PywhisperInference.translate, text, self.language, self.dst_lang)
-     
-        ret =   translate_fut.result()
+        ret =  translate_pipes.translate(text, self.language, self.dst_lang)
         log_block("LLM translate time", f"{(time.perf_counter() - start_time):.3f}", "s")
-        return ret
+        return ret.translate_content
      
     def _segments_split(self, segments, audio_buffer: np.ndarray):
         """根据左边第一个标点符号来将序列拆分成 观察段 和 剩余部分"""
@@ -327,23 +239,22 @@ class PyWhiperCppServe(ServeClientBase):
         return None, left_watch_string, right_watch_string, is_end_sentence
 
     def speech_to_text(self):
-        # c = 0
+        c = 0
         while True:
             if self.exit:
                 logger.info("Exiting speech to text thread")
-                self._pool.shutdown(wait=False, cancel_futures=True)
                 break
     
             if self.frames_np is None:
                 time.sleep(0.02)  # wait for any audio to arrive
                 continue
-
+            
             audio_buffer = self.get_audio_chunk_for_processing()
             # c+= 1
             # name = f"dev-{c}.wav"
             # save_to_wave(name, audio_buffer)
             try:
-                # logger.info(f"Audio buffer length: {len(audio_buffer) / self.sample_rate:.2f}s")
+                logger.info(f"Audio buffer length: {len(audio_buffer) / self.sample_rate:.2f}s")
                 segments = self.transcribe_audio(audio_buffer)
                 for tran_result in self.handle_transcription_output(segments, audio_buffer):
                     self.send_to_client(tran_result)
@@ -352,8 +263,6 @@ class PyWhiperCppServe(ServeClientBase):
             except Exception as e:
                 logger.error(f"{e}")
     
-           
-    
     def handle_transcription_output(self, segments, audio_buffer):
         texts  = "".join(i.text for i in segments)
         if not len(texts):
@@ -423,7 +332,4 @@ class PyWhiperCppServe(ServeClientBase):
         return padded_audio.copy()
 
     def cleanup(self):
-        logger.info("start shut down worker pool.")
-        self._pool.shutdown(wait=False, cancel_futures=True)
-        logger.info("shut down worker pool success.")
         return super().cleanup()#