import os
|
import time
|
import sys
|
import librosa
|
from funasr.utils.types import str2bool
|
|
import argparse
|
|
parser = argparse.ArgumentParser()
|
parser.add_argument("--model_dir", type=str, required=True)
|
parser.add_argument("--backend", type=str, default="onnx", help='["onnx", "torch"]')
|
parser.add_argument("--wav_file", type=str, default=None, help="amp fallback number")
|
parser.add_argument("--quantize", type=str2bool, default=False, help="quantized model")
|
parser.add_argument(
|
"--intra_op_num_threads", type=int, default=1, help="intra_op_num_threads for onnx"
|
)
|
parser.add_argument("--output_dir", type=str, default=None, help="amp fallback number")
|
args = parser.parse_args()
|
|
|
from funasr.runtime.python.libtorch.funasr_torch import Paraformer
|
|
if args.backend == "onnx":
|
from funasr.runtime.python.onnxruntime.funasr_onnx import Paraformer
|
|
model = Paraformer(
|
args.model_dir,
|
batch_size=1,
|
quantize=args.quantize,
|
intra_op_num_threads=args.intra_op_num_threads,
|
)
|
|
wav_file_f = open(args.wav_file, "r")
|
wav_files = wav_file_f.readlines()
|
|
output_dir = args.output_dir
|
if not os.path.exists(output_dir):
|
os.makedirs(output_dir)
|
if os.name == "nt": # Windows
|
newline = "\r\n"
|
else: # Linux Mac
|
newline = "\n"
|
text_f = open(os.path.join(output_dir, "text"), "w", newline=newline)
|
token_f = open(os.path.join(output_dir, "token"), "w", newline=newline)
|
|
for i, wav_path_i in enumerate(wav_files):
|
wav_name, wav_path = wav_path_i.strip().split()
|
result = model(wav_path)
|
text_i = "{} {}\n".format(wav_name, result[0]["preds"][0])
|
token_i = "{} {}\n".format(wav_name, result[0]["preds"][1])
|
text_f.write(text_i)
|
text_f.flush()
|
token_f.write(token_i)
|
token_f.flush()
|
text_f.close()
|
token_f.close()
|
