520 lines
22 KiB
Python
520 lines
22 KiB
Python
import os
|
||
import sys
|
||
import time
|
||
import uuid
|
||
import datetime
|
||
import tempfile
|
||
import soundfile as sf
|
||
import sherpa_onnx
|
||
import traceback
|
||
|
||
from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Form
|
||
|
||
os.makedirs("./input", exist_ok=True)
|
||
status = "Running"
|
||
recognizer = None
|
||
device = ""
|
||
model_type = ""
|
||
app = FastAPI()
|
||
|
||
CUSTOM_DEVICE = os.getenv("CUSTOM_DEVICE", "")
|
||
|
||
# 根据名称判断模型类型，比较杂，一共种类的自定义类型包括(针对OfflineRecognizer)
|
||
# moonshine
|
||
# fire_red
|
||
# dolphin_ctc
|
||
# paraformer
|
||
# telespeech_ctc
|
||
# whisper
|
||
# sensevoice
|
||
# zipformer_ctc
|
||
# transducer
|
||
# nemo_ctc
|
||
# nemo_canary
|
||
# wenet_ctc
|
||
# 针对OnlineRecognizer只有 zipformer_ctc transducer paraformer nemo_ctc wenet_ctc 四种
|
||
def get_asr_model_type(model_name):
|
||
# 根据名称判断模型类型以及需要检测的语种任务
|
||
# nemo_ctc, nemo_canary, moonshine 目前sherpa-onnx没有中文模型，执行英文ASR任务，其余模型执行中文ASR
|
||
# 所有nemo模型(nemo_ctc, nemo_canary以及transuducer中的nemo模型)均无中文模型
|
||
# 英文模型也并非全部大类都支持
|
||
|
||
# 特殊规则
|
||
# zipformer带ctc的才属于zipformer_ctc那一类，否则属于transducer类
|
||
# nemo也是带上ctc或者canary才属于单独类别，否则属于transducer类
|
||
# conformer均为transducer类,但是得在nemo之后判断
|
||
# wenet 由于同时wenetspeech为数据集名称，各种类型都有可能，这个逻辑需放在后面
|
||
model_type = "unknown"
|
||
model_name_lower = model_name.lower()
|
||
if "tdnn" in model_name_lower:
|
||
model_type = "tdnn" # tdnn类别不适用，目前仅有一个模型只能识别希伯来语中的yes/no两种词语
|
||
elif "moonshine" in model_name_lower:
|
||
model_type = "moonshine"
|
||
elif "fire-red" in model_name_lower:
|
||
model_type = "fire_red"
|
||
elif "dolphin" in model_name_lower:
|
||
model_type = "dolphin_ctc"
|
||
elif "paraformer" in model_name_lower:
|
||
model_type = "paraformer"
|
||
elif "telespeech" in model_name_lower:
|
||
model_type = "telespeech_ctc"
|
||
elif "whisper" in model_name_lower:
|
||
model_type = "whisper"
|
||
elif "sense-voice" in model_name_lower:
|
||
model_type = "sensevoice"
|
||
elif "zipformer" in model_name_lower:
|
||
if "ctc" in model_name_lower:
|
||
model_type = "zipformer_ctc"
|
||
else:
|
||
model_type = "transducer"
|
||
elif "nemo" in model_name_lower:
|
||
if "ctc" in model_name_lower:
|
||
model_type = "nemo_ctc"
|
||
elif "canary" in model_name_lower:
|
||
model_type = "nemo_canary"
|
||
else:
|
||
model_type = "transducer"
|
||
elif "conformer" in model_name_lower or "lstm" in model_name_lower:
|
||
model_type = "transducer"
|
||
elif "wenet" in model_name_lower:
|
||
model_type = "wenet_ctc"
|
||
else:
|
||
model_type = "unknown"
|
||
return model_type
|
||
|
||
@app.on_event("startup")
|
||
def load_model():
|
||
global status, recognizer, device, model_type
|
||
config = app.state.config
|
||
use_gpu = config.get("use_gpu", True)
|
||
model_dir = config.get("model_dir", "/model")
|
||
_model_type = config.get("model_type", None)
|
||
_model_name = config.get("model_name", None)
|
||
warmup = config.get("warmup", False)
|
||
isOffline = config.get("offline_model", True)
|
||
num_threads = config.get("num_threads", 2)
|
||
|
||
device = "cpu"
|
||
if use_gpu:
|
||
if CUSTOM_DEVICE.startswith("mlu"):
|
||
device = "mlu:0"
|
||
elif CUSTOM_DEVICE.startswith("ascend"):
|
||
device = "npu:0"
|
||
else:
|
||
device = "cuda:0"
|
||
|
||
# sherpa-onnx类型繁杂，当用户清楚的时候可提供model_type参数，抑或是提供完整的模型名称也行
|
||
# 因为挂载进入镜像的时候镜像内的文件路径不一定包含了模型名称
|
||
if _model_type:
|
||
model_type = _model_type
|
||
elif _model_name:
|
||
model_type = get_asr_model_type(_model_name)
|
||
else:
|
||
print("model_name and model_type both not provided, start guessing using model_dir", flush=True)
|
||
model_name = os.path.basename(model_dir)
|
||
model_type = get_asr_model_type(model_name)
|
||
|
||
print(">> Startup config:")
|
||
print(" model_dir =", model_dir, flush=True)
|
||
print(" model_type =", model_type, flush=True)
|
||
print(" use_gpu =", use_gpu, flush=True)
|
||
print(" warmup =", warmup, flush=True)
|
||
print(" isOffline =", isOffline, flush=True)
|
||
print(" num_threads =", num_threads, flush=True)
|
||
|
||
try:
|
||
recognizer = None
|
||
provider = "cuda" if use_gpu else "cpu"
|
||
file_list = os.listdir(model_dir)
|
||
# 目录内的模型文件可能会有多套（例如量化和不带量化版），选取大小最大的那一套
|
||
if model_type == "whisper":
|
||
encoder_list, decoder_list = [], []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if "encode" in file and file.endswith(".onnx"):
|
||
encoder_list.append(file)
|
||
elif "decode" in file and file.endswith(".onnx"):
|
||
decoder_list.append(file)
|
||
elif "token" in file and file.endswith(".txt"):
|
||
tokens = file
|
||
encoder = sorted(encoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
decoder = sorted(decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_whisper(
|
||
encoder=model_dir + "/" + encoder,
|
||
decoder=model_dir + "/" + decoder,
|
||
tokens=model_dir + "/" + tokens,
|
||
language="zh",
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
|
||
elif model_type == "sensevoice":
|
||
model_list = []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_sense_voice(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
use_itn=True,
|
||
language="zh",
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "paraformer":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
if isOffline:
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_paraformer(
|
||
paraformer=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
recognizer = sherpa_onnx.OnlineRecognizer.from_paraformer(
|
||
paraformer=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "zipformer_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
if isOffline:
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_zipformer_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
recognizer = sherpa_onnx.OnlineRecognizer.from_zipformer2_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "telespeech_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_telespeech_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "fire_red":
|
||
encoder_list, decoder_list = [], []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if "encode" in file and file.endswith(".onnx"):
|
||
encoder_list.append(file)
|
||
elif "decode" in file and file.endswith(".onnx"):
|
||
decoder_list.append(file)
|
||
elif "token" in file and file.endswith(".txt"):
|
||
tokens = file
|
||
encoder = sorted(encoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
decoder = sorted(decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_fire_red_asr(
|
||
encoder=model_dir + "/" + encoder,
|
||
decoder=model_dir + "/" + decoder,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "wenet_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
if isOffline:
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_wenet_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
recognizer = sherpa_onnx.OnlineRecognizer.from_wenet_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "dolphin_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_dolphin_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "transducer":
|
||
encoder_list, decoder_list, joiner_list = [], [], []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if "encode" in file and file.endswith(".onnx"):
|
||
encoder_list.append(file)
|
||
elif "decode" in file and file.endswith(".onnx"):
|
||
decoder_list.append(file)
|
||
elif "joiner" in file and file.endswith(".onnx"):
|
||
joiner_list.append(file)
|
||
elif "token" in file and file.endswith(".txt"):
|
||
tokens = file
|
||
encoder = sorted(encoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
decoder = sorted(decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
joiner = sorted(joiner_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
# 特殊情况，zipformer,conformer都是icefall导出，默认类型即可，nemo-transducer需要专门区分
|
||
transducer_type = "nemo_transducer" if "nemo" in model_name.lower() else "transducer"
|
||
if isOffline:
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_transducer(
|
||
encoder=model_dir + "/" + encoder,
|
||
decoder=model_dir + "/" + decoder,
|
||
joiner=model_dir + "/" + joiner,
|
||
tokens=model_dir + "/" + tokens,
|
||
model_type=transducer_type,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
recognizer = sherpa_onnx.OnlineRecognizer.from_transducer(
|
||
encoder=model_dir + "/" + encoder,
|
||
decoder=model_dir + "/" + decoder,
|
||
joiner=model_dir + "/" + joiner,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "nemo_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
if isOffline:
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
recognizer = sherpa_onnx.OnlineRecognizer.from_nemo_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "nemo_canary":
|
||
encoder_list, decoder_list = [], []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if "encode" in file and file.endswith(".onnx"):
|
||
encoder_list.append(file)
|
||
elif "decode" in file and file.endswith(".onnx"):
|
||
decoder_list.append(file)
|
||
elif "token" in file and file.endswith(".txt"):
|
||
tokens = file
|
||
encoder = sorted(encoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
decoder = sorted(decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_nemo_canary(
|
||
encoder=model_dir + "/" + encoder,
|
||
decoder=model_dir + "/" + decoder,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "moonshine":
|
||
preprocessor_list, encoder_list, cached_decoder_list, uncached_decoder_list = [], [], [], []
|
||
tokens = ""
|
||
for file in file_list:
|
||
if "preprocess" in file and file.endswith(".onnx"):
|
||
preprocessor_list.append(file)
|
||
elif "encode" in file and file.endswith(".onnx"):
|
||
encoder_list.append(file)
|
||
elif "uncached_decode" in file and file.endswith(".onnx"):
|
||
uncached_decoder_list.append(file)
|
||
elif "cached_decode" in file and file.endswith(".onnx"):
|
||
cached_decoder_list.append(file)
|
||
elif "token" in file and file.endswith(".txt"):
|
||
tokens = file
|
||
preprocessor = sorted(preprocessor_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
encoder = sorted(encoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
cached_decoder = sorted(cached_decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
uncached_decoder = sorted(uncached_decoder_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_moonshine(
|
||
preprocessor=model_dir + "/" + preprocessor,
|
||
encoder=model_dir + "/" + encoder,
|
||
cached_decoder=model_dir + "/" + cached_decoder,
|
||
uncached_decoder=model_dir + "/" + uncached_decoder,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
elif model_type == "tdnn_ctc":
|
||
model_list = []
|
||
for file in file_list:
|
||
if file.endswith(".onnx"):
|
||
model_list.append(file)
|
||
elif file.endswith(".txt") and "token" in file:
|
||
tokens = file
|
||
model = sorted(model_list, key=lambda x: os.path.getsize(model_dir + "/" + x), reverse=True)[0]
|
||
recognizer = sherpa_onnx.OfflineRecognizer.from_tdnn_ctc(
|
||
model=model_dir + "/" + model,
|
||
tokens=model_dir + "/" + tokens,
|
||
debug=False,
|
||
provider=provider,
|
||
num_threads=num_threads
|
||
)
|
||
else:
|
||
raise RuntimeError("Cannot recognize model_type")
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to initial cuda model: {e}")
|
||
|
||
if warmup:
|
||
print("Start warmup...", flush=True)
|
||
stream = recognizer.create_stream()
|
||
audio, sample_rate = sf.read("warmup.wav", dtype="float32", always_2d=True)
|
||
stream.accept_waveform(sample_rate, audio)
|
||
recognizer.decode_stream(stream)
|
||
print("warmup complete.", flush=True)
|
||
|
||
status = "Success"
|
||
|
||
def test_sherpa(wavefile):
|
||
isOffline = app.state.config.get("offline_model", True)
|
||
audio, sample_rate = sf.read(wavefile, dtype="float32", always_2d=True)
|
||
audio = audio[:, 0]
|
||
generated_text = ""
|
||
|
||
start_t = datetime.datetime.now()
|
||
if isOffline:
|
||
# OfflineRecognizer非流式模型推理
|
||
if model_type in ["sensevoice"]:
|
||
stream = recognizer.create_stream()
|
||
stream.accept_waveform(sample_rate, audio)
|
||
recognizer.decode_stream(stream)
|
||
generated_text = stream.result.text
|
||
else:
|
||
# offline-asr model 大多对长音频支持不佳，模型训练音频不长以及导出onnx结构中对一些中间态维度可能有上限
|
||
# 哪怕原版CPU推理中间可能都会崩溃，采取小段切分形式测试
|
||
start_index = 0
|
||
internal = int(sample_rate * 29)
|
||
while start_index < len(audio):
|
||
stream = recognizer.create_stream()
|
||
stream.accept_waveform(sample_rate, audio[start_index:start_index + internal])
|
||
recognizer.decode_stream(stream)
|
||
generated_text += stream.result.text
|
||
start_index += internal
|
||
else:
|
||
# OnlineRecognizer流式模型推理,统一每一次只投喂2s音频数据
|
||
stream = recognizer.create_stream()
|
||
start_index = 0
|
||
chunk_size = int(sample_rate * 2)
|
||
|
||
while start_index < len(audio):
|
||
chunk = audio[start_index:start_index + chunk_size]
|
||
stream.accept_waveform(sample_rate, chunk)
|
||
|
||
while recognizer.is_ready(stream):
|
||
recognizer.decode_stream(stream)
|
||
# mid_text = recognizer.get_result(stream)
|
||
# print("partial result: " + mid_text, flush=True)
|
||
start_index += chunk_size
|
||
|
||
while recognizer.is_ready(stream):
|
||
recognizer.decode_stream(stream)
|
||
generated_text = recognizer.get_result(stream)
|
||
|
||
end_t = datetime.datetime.now()
|
||
elapsed_seconds = (end_t - start_t).total_seconds()
|
||
duration = audio.shape[-1] / sample_rate
|
||
rtf = elapsed_seconds / duration
|
||
|
||
print("Text:", generated_text, flush=True)
|
||
print(f"Audio duration:\t{duration:.3f} s", flush=True)
|
||
print(f"Elapsed:\t{elapsed_seconds:.3f} s", flush=True)
|
||
print(f"RTF = {elapsed_seconds:.3f}/{duration:.3f} = {rtf:.3f}", flush=True)
|
||
|
||
return generated_text
|
||
|
||
@app.get("/health")
|
||
def health():
|
||
|
||
if status=="Running":
|
||
return {
|
||
"status":"loading model"
|
||
}
|
||
ret = {
|
||
"status": "ok" if status == "Success" else "failed",
|
||
}
|
||
return ret
|
||
|
||
@app.post("/transduce")
|
||
def transduce(
|
||
audio: UploadFile = File(...),
|
||
lang: str = Form("zh"),
|
||
background_tasks: BackgroundTasks = None
|
||
):
|
||
try:
|
||
file_path = f"./input/{uuid.uuid4()}.wav"
|
||
with open(file_path, "wb") as f:
|
||
f.write(audio.file.read())
|
||
background_tasks.add_task(os.remove, file_path)
|
||
generated_text = test_sherpa(file_path)
|
||
|
||
return {"generated_text": generated_text}
|
||
except Exception:
|
||
raise HTTPException(status_code=500, detail=f"Processing failed: \n{traceback.format_exc()}")
|
||
|
||
# if __name__ == "__main__":
|
||
|
||
# uvicorn.run("fastapi_sherpa:app", host="0.0.0.0", port=1111, workers=1) |