From 6c467e6f0abfc6d20d0621fbbf67b4dbd81776cc Mon Sep 17 00:00:00 2001
From: Shi Xian <40013335+R1ckShi@users.noreply.github.com>
Date: 星期二, 18 六月 2024 10:01:56 +0800
Subject: [PATCH] Merge pull request #1825 from modelscope/dev_libt

---
 funasr/auto/auto_model.py |   43 ++++++++++++++++++++-----------------------
 1 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index fb81608..91e80d8 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -234,6 +234,8 @@
         # fp16
         if kwargs.get("fp16", False):
             model.to(torch.float16)
+        elif kwargs.get("bf16", False):
+            model.to(torch.bfloat16)
         return model, kwargs
 
     def __call__(self, *args, **cfg):
@@ -428,6 +430,10 @@
             #                      f"time_speech_total_per_sample: {time_speech_total_per_sample: 0.3f}, "
             #                      f"time_escape_total_per_sample: {time_escape_total_per_sample:0.3f}")
 
+            if len(results_sorted) != n:
+                results_ret_list.append({"key": key, "text": "", "timestamp": []})
+                logging.info("decoding, utt: {}, empty result".format(key))
+                continue
             restored_data = [0] * n
             for j in range(n):
                 index = sorted_data[j][1]
@@ -460,25 +466,22 @@
                             result[k] = restored_data[j][k]
                         else:
                             result[k] += restored_data[j][k]
-
+                            
+            if not len(result["text"].strip()):
+                continue
             return_raw_text = kwargs.get("return_raw_text", False)
             # step.3 compute punc model
+            raw_text = None
             if self.punc_model is not None:
-                if not len(result["text"].strip()):
-                    if return_raw_text:
-                        result["raw_text"] = ""
-                else:
-                    deep_update(self.punc_kwargs, cfg)
-                    punc_res = self.inference(
-                        result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg
-                    )
-                    raw_text = copy.copy(result["text"])
-                    if return_raw_text:
-                        result["raw_text"] = raw_text
-                    result["text"] = punc_res[0]["text"]
-            else:
-                raw_text = None
-
+                deep_update(self.punc_kwargs, cfg)
+                punc_res = self.inference(
+                    result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg
+                )
+                raw_text = copy.copy(result["text"])
+                if return_raw_text:
+                    result["raw_text"] = raw_text
+                result["text"] = punc_res[0]["text"]
+                
             # speaker embedding cluster after resorted
             if self.spk_model is not None and kwargs.get("return_spk_res", True):
                 if raw_text is None:
@@ -599,12 +602,6 @@
         )
 
         with torch.no_grad():
-
-            if type == "onnx":
-                export_dir = export_utils.export_onnx(model=model, data_in=data_list, **kwargs)
-            else:
-                export_dir = export_utils.export_torchscripts(
-                    model=model, data_in=data_list, **kwargs
-                )
+            export_dir = export_utils.export(model=model, data_in=data_list,  **kwargs)
 
         return export_dir

--
Gitblit v1.9.1