From 94f7bd079c78d6b7f9ab1961e991caa2a929b858 Mon Sep 17 00:00:00 2001
From: shixian.shi <shixian.shi@alibaba-inc.com>
Date: 星期一, 05 二月 2024 11:15:06 +0800
Subject: [PATCH] update preset_spk_num

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

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index 3986a11..8007d6e 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -121,9 +121,6 @@
             if spk_mode not in ["default", "vad_segment", "punc_segment"]:
                 logging.error("spk_mode should be one of default, vad_segment and punc_segment.")
             self.spk_mode = spk_mode
-            self.preset_spk_num = kwargs.get("preset_spk_num", None)
-            if self.preset_spk_num:
-                logging.warning("Using preset speaker number: {}".format(self.preset_spk_num))
             
         self.kwargs = kwargs
         self.model = model
@@ -228,7 +225,7 @@
             data_batch = data_list[beg_idx:end_idx]
             key_batch = key_list[beg_idx:end_idx]
             batch = {"data_in": data_batch, "key": key_batch}
-            if (end_idx - beg_idx) == 1 and isinstance(data_batch[0], torch.Tensor): # fbank
+            if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank
                 batch["data_in"] = data_batch[0]
                 batch["data_lengths"] = input_len
         
@@ -391,7 +388,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.cpu(), oracle_num=self.preset_spk_num)
+                labels = self.cb_model(spk_embedding.cpu(), oracle_num=kwargs['preset_spk_num'])
                 del result['spk_embedding']
                 sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu())
                 if self.spk_mode == 'vad_segment':  # recover sentence_list

--
Gitblit v1.9.1