diff --git a/STT/paraformer_handler.py b/STT/paraformer_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..0a2a9c002716b9acd1b3ddf1c52e847aa571ef2e --- /dev/null +++ b/STT/paraformer_handler.py @@ -0,0 +1,61 @@ +import logging +from time import perf_counter + +from tensorstore import dtype + +from baseHandler import BaseHandler +from funasr import AutoModel +import numpy as np +from rich.console import Console +import torch + +logging.basicConfig( + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger(__name__) + +console = Console() + + +class ParaformerSTTHandler(BaseHandler): + """ + Handles the Speech To Text generation using a Whisper model. + """ + + def setup( + self, + model_name="paraformer-zh", + device="cuda", + torch_dtype="float32", + compile_mode=None, + gen_kwargs={}, + ): + print(model_name) + if len(model_name.split("/")) > 1: + model_name = model_name.split("/")[-1] + self.device = device + self.model = AutoModel(model=model_name) + self.warmup() + + def warmup(self): + logger.info(f"Warming up {self.__class__.__name__}") + + # 2 warmup steps for no compile or compile mode with CUDA graphs capture + n_steps = 1 + dummy_input = np.array([0] * 512,dtype=np.float32) + for _ in range(n_steps): + _ = self.model.generate(dummy_input)[0]["text"].strip().replace(" ","") + + def process(self, spoken_prompt): + logger.debug("infering paraformer...") + + global pipeline_start + pipeline_start = perf_counter() + + pred_text = self.model.generate(spoken_prompt)[0]["text"].strip().replace(" ","") + torch.mps.empty_cache() + + logger.debug("finished paraformer inference") + console.print(f"[yellow]USER: {pred_text}") + + yield pred_text diff --git a/s2s_pipeline.py b/s2s_pipeline.py index 7a9e204ae16bdcde3992ecfa09bda7e73dabae95..002cde968ed8134cebe6a86b8ab62ce195efc8ff 100644 --- a/s2s_pipeline.py +++ b/s2s_pipeline.py @@ -243,6 +243,14 @@ def main(): queue_out=text_prompt_queue, setup_kwargs=vars(whisper_stt_handler_kwargs), ) + elif module_kwargs.stt == "paraformer": + from STT.paraformer_handler import ParaformerSTTHandler + stt = ParaformerSTTHandler( + stop_event, + queue_in=spoken_prompt_queue, + queue_out=text_prompt_queue, + # setup_kwargs=vars(whisper_stt_handler_kwargs), + ) else: raise ValueError("The STT should be either whisper or whisper-mlx") if module_kwargs.llm == "transformers":