From 6ebd2676480068c1cb27ffc1b9318b09e1662173 Mon Sep 17 00:00:00 2001
From: 九耳 <mengzhe.cmz@alibaba-inc.com>
Date: 星期三, 29 三月 2023 16:48:57 +0800
Subject: [PATCH] general punc model conversion onnx
---
funasr/export/export_model.py | 12 +++++++++++-
1 files changed, 11 insertions(+), 1 deletions(-)
diff --git a/funasr/export/export_model.py b/funasr/export/export_model.py
index 4aed49f..9afa7b1 100644
--- a/funasr/export/export_model.py
+++ b/funasr/export/export_model.py
@@ -174,7 +174,10 @@
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
config = os.path.join(model_dir, 'config.yaml')
@@ -195,6 +198,13 @@
)
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'
+ )
self._export(model, tag_name)
--
Gitblit v1.9.1