transcribe/translatepipes.py

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()
        
        # llm 翻译
        self._translate_pipe = TranslatePipe()

        
        self._whisper_pipe.daemon = True
        self._whisper_pipe.start()

        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)
        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)