From 9b4e9cc8a0311e5243d69b73ed073e7ea441982e Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 27 三月 2024 16:05:29 +0800
Subject: [PATCH] train update
---
funasr/models/ct_transformer/model.py | 31 +++++++++++++++++++------------
1 files changed, 19 insertions(+), 12 deletions(-)
diff --git a/funasr/models/ct_transformer/model.py b/funasr/models/ct_transformer/model.py
index 1891ac7..57a23cc 100644
--- a/funasr/models/ct_transformer/model.py
+++ b/funasr/models/ct_transformer/model.py
@@ -17,8 +17,10 @@
from funasr.utils.load_utils import load_audio_text_image_video
from funasr.models.transformer.utils.nets_utils import make_pad_mask
from funasr.models.ct_transformer.utils import split_to_mini_sentence, split_words
-
-
+try:
+ import jieba
+except:
+ pass
if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
from torch.cuda.amp import autocast
else:
@@ -70,6 +72,10 @@
self.sos = sos
self.eos = eos
self.sentence_end_id = sentence_end_id
+ self.jieba_usr_dict = None
+ if kwargs.get("jieba_usr_dict", None) is not None:
+ jieba.load_userdict(kwargs["jieba_usr_dict"])
+ self.jieba_usr_dict = jieba
@@ -238,14 +244,8 @@
# text = data_in[0]
# text_lengths = data_lengths[0] if data_lengths is not None else None
split_size = kwargs.get("split_size", 20)
-
- jieba_usr_dict = kwargs.get("jieba_usr_dict", None)
- if jieba_usr_dict and isinstance(jieba_usr_dict, str):
- import jieba
- jieba.load_userdict(jieba_usr_dict)
- jieba_usr_dict = jieba
- kwargs["jieba_usr_dict"] = "jieba_usr_dict"
- tokens = split_words(text, jieba_usr_dict=jieba_usr_dict)
+
+ tokens = split_words(text, jieba_usr_dict=self.jieba_usr_dict)
tokens_int = tokenizer.encode(tokens)
mini_sentences = split_to_mini_sentence(tokens, split_size)
@@ -348,7 +348,7 @@
else:
punc_array = torch.cat([punc_array, punctuations], dim=0)
# post processing when using word level punc model
- if jieba_usr_dict:
+ if self.jieba_usr_dict is not None:
len_tokens = len(tokens)
new_punc_array = copy.copy(punc_array).tolist()
# for i, (token, punc_id) in enumerate(zip(tokens[::-1], punc_array.tolist()[::-1])):
@@ -359,9 +359,16 @@
ind_append = len_tokens - i - 1
for _ in range(num_append):
new_punc_array.insert(ind_append, 1)
- punc_array = torch.tensor(new_punc_array)
+ punc_array = torch.tensor(new_punc_array)
result_i = {"key": key[0], "text": new_mini_sentence_out, "punc_array": punc_array}
results.append(result_i)
return results, meta_data
+ def export(self, **kwargs):
+
+ from .export_meta import export_rebuild_model
+ models = export_rebuild_model(model=self, **kwargs)
+ return models
+
+
--
Gitblit v1.9.1