TestTranslator / lib /asr_models /funasr_quant.py
yujuanqin's picture
add asr test
db0d138
from pathlib import Path
from funasr_onnx import SeacoParaformer, CT_Transformer, Fsmn_vad
from lib.utils import Timer
from lib.asr_models.base_model import AbstractASRModel, ModelName
MODEL_DIR = "/Users/jeqin/work/code/Translator/python_server/moyoyo_asr_models"
class FunasrQuant(AbstractASRModel):
def __init__(self, device='mps'):
super().__init__(device=device)
self.name = ModelName.FUNASR_QUANT
def load(self, model_dir=MODEL_DIR, language=""):
quantize=True
with Timer("Loading Fun-ASR-Quant model"):
model_dir = Path(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.vad_model = Fsmn_vad(vad_model_path, quantize=quantize)
self.asr_model = SeacoParaformer(asr_model_path, quantize=quantize)
self.punc_model = CT_Transformer(punc_model_path, quantize=quantize)
def transcribe(self, wav, language="zh"):
with Timer("Transcribing audio") as t:
asr_res = self.asr_model(str(wav), hotwords="", language=language)
text = ""
if len(asr_res) > 0:
asr_text = asr_res[0]["preds"]
result = self.punc_model(asr_text)
text = result[0]
return text, t.duration
if __name__ == "__main__":
model = FunasrQuant(device='mps')
model.load()
text, cost = model.transcribe('../../test_data/recordings/1.wav', language="en")
print("inference time: ", cost)
print(text)