Merge pull request #687 from alibaba-damo-academy/dev_lhn

hnluo

2023-06-29 c2dee5e3c29eba79e591d9e9caebaef15ea4e56b

 egs_modelscope/asr/paraformer/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online/infer.py

@@ -1,52 +1,33 @@
import torch
import torchaudio
import os
import shutil
import argparse
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

inference_pipeline = pipeline(
    task=Tasks.auto_speech_recognition,
    model='damo/speech_paraformer_asr_nat-zh-cn-16k-common-vocab8404-online',
    model_revision='v1.0.2')
def modelscope_infer(args):
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpuid)
    inference_pipeline = pipeline(
        task=Tasks.auto_speech_recognition,
        model=args.model,
        output_dir=args.output_dir,
        batch_size=args.batch_size,
        model_revision='v1.0.6',
        update_model=False,
        mode="paraformer_fake_streaming",
        param_dict={"decoding_model": args.decoding_mode, "hotword": args.hotword_txt}
    )
    inference_pipeline(audio_in=args.audio_in)

waveform, sample_rate = torchaudio.load("asr_example_zh.wav")
speech_length = waveform.shape[1]
speech = waveform[0]

cache_en = {"start_idx": 0, "pad_left": 0, "stride": 10, "pad_right": 5, "cif_hidden": None, "cif_alphas": None}
cache_de = {"decode_fsmn": None}
cache = {"encoder": cache_en, "decoder": cache_de}
param_dict = {}
param_dict["cache"] = cache

first_chunk = True
speech_buffer = speech
speech_cache = []
final_result = ""

while len(speech_buffer) > 0:
    if first_chunk:
        if len(speech_buffer) >= 14400:
            rec_result = inference_pipeline(audio_in=speech_buffer[0:14400], param_dict=param_dict)
            speech_buffer = speech_buffer[4800:]
        else:
            cache_en["stride"] = len(speech_buffer) // 960
            cache_en["pad_right"] = 0
            rec_result = inference_pipeline(audio_in=speech_buffer, param_dict=param_dict)
            speech_buffer = []
        cache_en["start_idx"] = -5
        first_chunk = False
    else:
        cache_en["start_idx"] += 10
        if len(speech_buffer) >= 4800:
            cache_en["pad_left"] = 5
            rec_result = inference_pipeline(audio_in=speech_buffer[:19200], param_dict=param_dict)
            speech_buffer = speech_buffer[9600:]
        else:
            cache_en["stride"] = len(speech_buffer) // 960
            cache_en["pad_right"] = 0
            rec_result = inference_pipeline(audio_in=speech_buffer, param_dict=param_dict)
            speech_buffer = []
    if rec_result['text'] != "sil":
        final_result += rec_result['text']
    print(rec_result)
print(final_result)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")
    parser.add_argument('--audio_in', type=str, default="./data/test/wav.scp")
    parser.add_argument('--output_dir', type=str, default="./results/")
    parser.add_argument('--decoding_mode', type=str, default="normal")
    parser.add_argument('--model_revision', type=str, default=None)
    parser.add_argument('--mode', type=str, default=None)
    parser.add_argument('--hotword_txt', type=str, default=None)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--gpuid', type=str, default="0")
    args = parser.parse_args()
    modelscope_infer(args)