From 6427c834dfd97b1f05c6659cdc7ccf010bf82fe1 Mon Sep 17 00:00:00 2001
From: 嘉渊 <wangjiaming.wjm@alibaba-inc.com>
Date: 星期一, 24 四月 2023 19:50:07 +0800
Subject: [PATCH] update
---
funasr/export/export_model.py | 60 +++++++++++++++++++++++++++++++++++++++++++++++-------------
1 files changed, 47 insertions(+), 13 deletions(-)
diff --git a/funasr/export/export_model.py b/funasr/export/export_model.py
index b1161cb..b69eeee 100644
--- a/funasr/export/export_model.py
+++ b/funasr/export/export_model.py
@@ -19,6 +19,7 @@
self,
cache_dir: Union[Path, str] = None,
onnx: bool = True,
+ device: str = "cpu",
quant: bool = True,
fallback_num: int = 0,
audio_in: str = None,
@@ -36,6 +37,7 @@
)
print("output dir: {}".format(self.cache_dir))
self.onnx = onnx
+ self.device = device
self.quant = quant
self.fallback_num = fallback_num
self.frontend = None
@@ -112,6 +114,10 @@
else:
dummy_input = model.get_dummy_inputs()
+ if self.device == 'cuda':
+ model = model.cuda()
+ dummy_input = tuple([i.cuda() for i in dummy_input])
+
# model_script = torch.jit.script(model)
model_script = torch.jit.trace(model, dummy_input)
model_script.save(os.path.join(path, f'{model.model_name}.torchscripts'))
@@ -161,31 +167,57 @@
def export(self,
tag_name: str = 'damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch',
- mode: str = 'paraformer',
+ mode: str = None,
):
model_dir = tag_name
- if model_dir.startswith('damo/'):
+ if model_dir.startswith('damo'):
from modelscope.hub.snapshot_download import snapshot_download
model_dir = snapshot_download(model_dir, cache_dir=self.cache_dir)
- asr_train_config = os.path.join(model_dir, 'config.yaml')
- asr_model_file = os.path.join(model_dir, 'model.pb')
- cmvn_file = os.path.join(model_dir, 'am.mvn')
- json_file = os.path.join(model_dir, 'configuration.json')
+
if mode is None:
import json
+ json_file = os.path.join(model_dir, 'configuration.json')
with open(json_file, 'r') as f:
config_data = json.load(f)
- mode = config_data['model']['model_config']['mode']
+ if config_data['task'] == "punctuation":
+ mode = config_data['model']['punc_model_config']['mode']
+ else:
+ mode = config_data['model']['model_config']['mode']
if mode.startswith('paraformer'):
from funasr.tasks.asr import ASRTaskParaformer as ASRTask
- elif mode.startswith('uniasr'):
- from funasr.tasks.asr import ASRTaskUniASR as ASRTask
+ config = os.path.join(model_dir, 'config.yaml')
+ model_file = os.path.join(model_dir, 'model.pb')
+ cmvn_file = os.path.join(model_dir, 'am.mvn')
+ model, asr_train_args = ASRTask.build_model_from_file(
+ config, model_file, cmvn_file, 'cpu'
+ )
+ self.frontend = model.frontend
+ elif mode.startswith('offline'):
+ from funasr.tasks.vad import VADTask
+ config = os.path.join(model_dir, 'vad.yaml')
+ model_file = os.path.join(model_dir, 'vad.pb')
+ cmvn_file = os.path.join(model_dir, 'vad.mvn')
- model, asr_train_args = ASRTask.build_model_from_file(
- asr_train_config, asr_model_file, cmvn_file, 'cpu'
- )
- self.frontend = model.frontend
+ model, vad_infer_args = VADTask.build_model_from_file(
+ config, model_file, cmvn_file=cmvn_file, device='cpu'
+ )
+ self.export_config["feats_dim"] = 400
+ self.frontend = model.frontend
+ elif mode.startswith('punc'):
+ from funasr.tasks.punctuation import PunctuationTask as PUNCTask
+ punc_train_config = os.path.join(model_dir, 'config.yaml')
+ punc_model_file = os.path.join(model_dir, 'punc.pb')
+ model, punc_train_args = PUNCTask.build_model_from_file(
+ punc_train_config, punc_model_file, 'cpu'
+ )
+ elif mode.startswith('punc_VadRealtime'):
+ from funasr.tasks.punctuation import PunctuationTask as PUNCTask
+ punc_train_config = os.path.join(model_dir, 'config.yaml')
+ punc_model_file = os.path.join(model_dir, 'punc.pb')
+ model, punc_train_args = PUNCTask.build_model_from_file(
+ punc_train_config, punc_model_file, 'cpu'
+ )
self._export(model, tag_name)
@@ -234,6 +266,7 @@
parser.add_argument('--model-name', type=str, required=True)
parser.add_argument('--export-dir', type=str, required=True)
parser.add_argument('--type', type=str, default='onnx', help='["onnx", "torch"]')
+ parser.add_argument('--device', type=str, default='cpu', help='["cpu", "cuda"]')
parser.add_argument('--quantize', type=str2bool, default=False, help='export quantized model')
parser.add_argument('--fallback-num', type=int, default=0, help='amp fallback number')
parser.add_argument('--audio_in', type=str, default=None, help='["wav", "wav.scp"]')
@@ -243,6 +276,7 @@
export_model = ModelExport(
cache_dir=args.export_dir,
onnx=args.type == 'onnx',
+ device=args.device,
quant=args.quantize,
fallback_num=args.fallback_num,
audio_in=args.audio_in,
--
Gitblit v1.9.1