change to pipelines

transcribe/pipelines/__init__.py

@@ -0,0 +1,4 @@
+
+from .pipe_translate import TranslatePipe
+from .pipe_whisper import WhisperPipe
+from .base import MetaItem

transcribe/pipelines/base.py

@@ -0,0 +1,50 @@
+
+from dataclasses import dataclass, field
+from multiprocessing import Process
+
+@dataclass
+class Segment:
+    t0: int
+    t1: int
+    text: str
+
+@dataclass
+class MetaItem:
+    segments: list[Segment] = field(default_factory=list)
+    audio: bytes = b''
+    transcribe_content: str = ''
+    translate_content: str = ''
+    source_language: str = 'zh'
+    destination_language: str = 'en'
+    
+
+class BasePipe(Process):
+    def __init__(self, in_queue, out_queue) -> None:
+        super().__init__()  # Initialize the Process class
+        self._in_queue = in_queue
+        self._out_queue = out_queue
+
+    @property
+    def output_queue(self):
+        return self._out_queue
+    
+    @property
+    def input_queue(self):
+        return self._in_queue
+    
+    def process(self, in_data: MetaItem) -> MetaItem:
+        raise NotImplementedError("Subclasses should implement this method.")
+
+
+    @classmethod
+    def init(cls):
+        raise NotImplementedError
+
+    def run(self):
+        self.init()
+        while True:
+            item = self._in_queue.get()
+            if item is None:  # Check for termination signal
+                break
+            out_item = self.process(item)
+            self._out_queue.put(out_item)

transcribe/pipelines/pipe_translate.py

@@ -0,0 +1,22 @@
+
+from .base import MetaItem, BasePipe, Segment
+from llama_cpp import Llama
+from ..translator import QwenTranslator
+from config import LLM_MODEL_PATH, LLM_SYS_PROMPT
+
+class TranslatePipe(BasePipe):
+    translator = None
+     
+    @classmethod
+    def init(cls):
+        if cls.translator is None:
+            cls.translator =  QwenTranslator(LLM_MODEL_PATH, LLM_SYS_PROMPT)
+
+    
+    def process(self, in_data: MetaItem) -> MetaItem:
+        context = in_data.transcribe_content
+        result = self.translator.translate(
+            context, src_lang=in_data.source_language, dst_lang=in_data.destination_language)
+        in_data.translate_content = result
+        return in_data
+

transcribe/pipelines/pipe_vad.py

@@ -0,0 +1,3 @@
+
+
+from .base import MetaItem, BasePipe

transcribe/pipelines/pipe_whisper.py

@@ -0,0 +1,28 @@
+
+
+from .base import MetaItem, BasePipe, Segment
+from ..whisper import WhisperCPP
+
+
+class WhisperPipe(BasePipe):
+    whisper = None
+
+    def __init__(self, in_queue, out_queue) -> None:
+        super().__init__(in_queue, out_queue)
+
+
+    @classmethod
+    def init(cls):
+        if cls.whisper is None:
+            cls.whisper = WhisperCPP()
+
+
+    def process(self, in_data: MetaItem) -> MetaItem:
+        audio_data = in_data.audio
+        source_language = in_data.source_language
+        segments = self.whisper.transcribe(audio_data, source_language)
+        texts = "".join([s.text for s in segments])
+        in_data.segments = [Segment(t0=s.t0, t1=s.t1, text=s.text) for s in segments]
+        in_data.transcribe_content = texts
+        in_data.audio = b""
+        return in_data

transcribe/translatepipes.py

@@ -0,0 +1,56 @@
+from transcribe.pipelines import WhisperPipe, TranslatePipe, MetaItem
+import multiprocessing  as mp
+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()
+        
+        # whisper 转录
+        self._whisper_pipe = WhisperPipe(
+            in_queue=self.whisper_input_q,
+            out_queue=self.translate_input_q
+            )
+        
+        # llm 翻译
+        self._translate_pipe = TranslatePipe(
+            in_queue=self.translate_input_q,
+            out_queue=self.result_queue,
+        )
+        
+        self._whisper_pipe.daemon = True
+        self._whisper_pipe.start()
+
+        self._translate_pipe.daemon = True
+        self._translate_pipe.start()
+
+    
+    def translate(self, text, src_lang, dst_lang) -> MetaItem:
+        item = MetaItem(
+            transcribe_content=text,
+              source_language=src_lang, 
+              destination_language=dst_lang)
+        self._translate_pipe.input_queue.put(item)
+        return self._translate_pipe.output_queue.get()
+    
+
+    def transcrible(self, audio_buffer:bytes, src_lang: str) -> MetaItem:
+        item = MetaItem(audio=audio_buffer, source_language=src_lang)
+        self._whisper_pipe.input_queue.put(item)
+        return self._whisper_pipe.output_queue.get()
+    
+
+if __name__ == "__main__":
+    import soundfile
+    tp = TranslatePipes()
+    # result = tp.translate("你好，今天天气怎么样?", src_lang="zh", dst_lang="en")
+    mel, _, = soundfile.read("assets/jfk.flac")
+    result = tp.transcrible(mel, 'en')
+    print(result)
+
+
+
+    

transcribe/whisper.py

@@ -0,0 +1,47 @@
+from pywhispercpp.model import Model
+import soundfile
+import config
+import numpy as np
+
+
+class WhisperCPP:
+
+    def __init__(self, warmup=True) -> None:
+        models_dir = config.MODEL_DIR.as_posix()
+        self.model = Model(
+            model=config.WHISPER_MODEL, 
+            models_dir=models_dir,
+            print_realtime=False,
+            print_progress=False,
+            print_timestamps=False,
+        )
+        if warmup:
+            self.warmup()
+
+    
+    def warmup(cls, warmup_steps=1):
+        mel, _, = soundfile.read("assets/jfk.flac")
+        for _ in range(warmup_steps):
+            cls.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}")
+
+    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(
+            audio_buffer,
+            initial_prompt=prompt,
+            language=language,  
+            token_timestamps=True, 
+            max_len=max_len
+        )
+        return output
+