From 0efc87352ce7d3903dbdedbfa5d01ca5e1cb19e7 Mon Sep 17 00:00:00 2001
From: Shi Xian <40013335+R1ckShi@users.noreply.github.com>
Date: 星期四, 05 十二月 2024 15:15:38 +0800
Subject: [PATCH] Merge pull request #2267 from modelscope/dev_sx2

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

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index e08cb2b..08308a2 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -147,13 +147,14 @@
         # if spk_model is not None, build spk model else None
         spk_model = kwargs.get("spk_model", None)
         spk_kwargs = {} if kwargs.get("spk_kwargs", {}) is None else kwargs.get("spk_kwargs", {})
+        cb_kwargs = {} if spk_kwargs.get("cb_kwargs", {}) is None else spk_kwargs.get("cb_kwargs", {})
         if spk_model is not None:
             logging.info("Building SPK model.")
             spk_kwargs["model"] = spk_model
             spk_kwargs["model_revision"] = kwargs.get("spk_model_revision", "master")
             spk_kwargs["device"] = kwargs["device"]
             spk_model, spk_kwargs = self.build_model(**spk_kwargs)
-            self.cb_model = ClusterBackend().to(kwargs["device"])
+            self.cb_model = ClusterBackend(**cb_kwargs).to(kwargs["device"])
             spk_mode = kwargs.get("spk_mode", "punc_segment")
             if spk_mode not in ["default", "vad_segment", "punc_segment"]:
                 logging.error("spk_mode should be one of default, vad_segment and punc_segment.")
@@ -199,6 +200,7 @@
             tokenizers_build = []
             vocab_sizes = []
             token_lists = []
+
             ### === only for kws ===
             token_list_files = kwargs.get("token_lists", [])
             seg_dicts = kwargs.get("seg_dicts", [])
@@ -213,9 +215,9 @@
 
                 ### === only for kws ===
                 if len(token_list_files) > 1:
-                    tokenizer_conf.token_list = token_list_files[i]
+                    tokenizer_conf["token_list"] = token_list_files[i]
                 if len(seg_dicts) > 1:
-                    tokenizer_conf.seg_dict = seg_dicts[i]
+                    tokenizer_conf["seg_dict"] = seg_dicts[i]
                 ### === only for kws ===
 
                 tokenizer = tokenizer_class(**tokenizer_conf)

--
Gitblit v1.9.1