transcribe/helpers/funasr.py

# from logging import getLogger

import numpy as np
from funasr import AutoModel

import config

# logger = getLogger(__name__)


class FunASR:
    def __init__(self, source_lange: str = 'en', warmup=True) -> None:
        self.source_lange = source_lange

        model_dir = config.MODEL_DIR
        asr_model_path = model_dir / 'speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'
        vad_model_path = model_dir / 'speech_fsmn_vad_zh-cn-16k-common-pytorch'
        punc_model_path = model_dir / 'punc_ct-transformer_cn-en-common-vocab471067-large'
        self.model = AutoModel(
            model=asr_model_path.as_posix(),
            vad_model=vad_model_path.as_posix(),
            punc_model=punc_model_path.as_posix(),
            log_level="ERROR",
            disable_update=True
        )
        if warmup:
            self.warmup()

    def warmup(self, warmup_steps=1):
        warmup_soundfile = f"{config.ASSERT_DIR}/jfk.flac"
        for _ in range(warmup_steps):
            self.model.generate(input=warmup_soundfile, disable_pbar=True)

    def transcribe(self, audio_buffer: bytes, language):
        audio_frames = np.frombuffer(audio_buffer, dtype=np.float32)
        # sf.write(f'{config.ASSERT_DIR}/{time.time()}.wav', audio_frames, samplerate=16000)
        try:
            output = self.model.generate(input=audio_frames, disable_pbar=True, hotword=config.hotwords_file.as_posix())
            return output
        except Exception as e:
            print(f"Error during transcription: {e}")
            return []