Dev gzf exp (#1647)

zhifu gao

2024-04-23 2ac38adbe5f4e1374a079e032ed4b504351a207c

 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)
@@ -347,8 +347,10 @@
                punc_array = punctuations
            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:
            punc_array = punc_array.reshape(-1)
            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 +361,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