From 0281b70896bd8069a9d26dad1ef2b5090ae5fc4a Mon Sep 17 00:00:00 2001
From: shixian.shi <shixian.shi@alibaba-inc.com>
Date: 星期四, 25 一月 2024 10:20:31 +0800
Subject: [PATCH] update umap import

---
 funasr/auto/auto_model.py |   10 +++++-----
 1 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index 107c78e..5f89dd6 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -146,7 +146,7 @@
         set_all_random_seed(kwargs.get("seed", 0))
         
         device = kwargs.get("device", "cuda")
-        if not torch.cuda.is_available() or kwargs.get("ngpu", 0):
+        if not torch.cuda.is_available() or kwargs.get("ngpu", 0) == 0:
             device = "cpu"
             kwargs["batch_size"] = 1
         kwargs["device"] = device
@@ -377,7 +377,7 @@
                             result[k] = restored_data[j][k]
                         else:
                             result[k] = torch.cat([result[k], restored_data[j][k]], dim=0)
-                    elif k == 'text':
+                    elif k == 'raw_text':
                         if k not in result:
                             result[k] = restored_data[j][k]
                         else:
@@ -398,7 +398,7 @@
             if self.spk_model is not None:
                 all_segments = sorted(all_segments, key=lambda x: x[0])
                 spk_embedding = result['spk_embedding']
-                labels = self.cb_model(spk_embedding, oracle_num=self.preset_spk_num)
+                labels = self.cb_model(spk_embedding.cpu(), oracle_num=self.preset_spk_num)
                 del result['spk_embedding']
                 sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu())
                 if self.spk_mode == 'vad_segment':
@@ -406,12 +406,12 @@
                     for res, vadsegment in zip(restored_data, vadsegments):
                         sentence_list.append({"start": vadsegment[0],\
                                                 "end": vadsegment[1],
-                                                "sentence": res['text'],
+                                                "sentence": res['raw_text'],
                                                 "timestamp": res['timestamp']})
                 else: # punc_segment
                     sentence_list = timestamp_sentence(punc_res[0]['punc_array'], \
                                                         result['timestamp'], \
-                                                        result['text'])
+                                                        result['raw_text'])
                 distribute_spk(sentence_list, sv_output)
                 result['sentence_info'] = sentence_list
                     

--
Gitblit v1.9.1