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

---
 funasr/auto/auto_model.py |   79 ++++++++++++++++++++++++++++++---------
 1 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index 75324dc..08308a2 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -14,6 +14,7 @@
 import numpy as np
 from tqdm import tqdm
 
+from omegaconf import DictConfig, ListConfig
 from funasr.utils.misc import deep_update
 from funasr.register import tables
 from funasr.utils.load_utils import load_bytes
@@ -114,7 +115,7 @@
         try:
             from funasr.utils.version_checker import check_for_update
 
-            check_for_update()
+            check_for_update(disable=kwargs.get("disable_update", False))
         except:
             pass
 
@@ -146,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.")
@@ -187,21 +189,60 @@
 
         # build tokenizer
         tokenizer = kwargs.get("tokenizer", None)
-        if tokenizer is not None:
-            tokenizer_class = tables.tokenizer_classes.get(tokenizer)
-            tokenizer = tokenizer_class(**kwargs.get("tokenizer_conf", {}))
-            kwargs["token_list"] = (
-                tokenizer.token_list if hasattr(tokenizer, "token_list") else None
-            )
-            kwargs["token_list"] = (
-                tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else kwargs["token_list"]
-            )
-            vocab_size = len(kwargs["token_list"]) if kwargs["token_list"] is not None else -1
-            if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"):
-                vocab_size = tokenizer.get_vocab_size()
-        else:
-            vocab_size = -1
         kwargs["tokenizer"] = tokenizer
+        kwargs["vocab_size"] = -1
+
+        if tokenizer is not None:
+            tokenizers = (
+                tokenizer.split(",") if isinstance(tokenizer, str) else tokenizer
+            )  # type of tokenizers is list!!!
+            tokenizers_conf = kwargs.get("tokenizer_conf", {})
+            tokenizers_build = []
+            vocab_sizes = []
+            token_lists = []
+
+            ### === only for kws ===
+            token_list_files = kwargs.get("token_lists", [])
+            seg_dicts = kwargs.get("seg_dicts", [])
+            ### === only for kws ===
+
+            if not isinstance(tokenizers_conf, (list, tuple, ListConfig)):
+                tokenizers_conf = [tokenizers_conf] * len(tokenizers)
+
+            for i, tokenizer in enumerate(tokenizers):
+                tokenizer_class = tables.tokenizer_classes.get(tokenizer)
+                tokenizer_conf = tokenizers_conf[i]
+
+                ### === only for kws ===
+                if len(token_list_files) > 1:
+                    tokenizer_conf["token_list"] = token_list_files[i]
+                if len(seg_dicts) > 1:
+                    tokenizer_conf["seg_dict"] = seg_dicts[i]
+                ### === only for kws ===
+
+                tokenizer = tokenizer_class(**tokenizer_conf)
+                tokenizers_build.append(tokenizer)
+                token_list = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
+                token_list = (
+                    tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else token_list
+                )
+                vocab_size = -1
+                if token_list is not None:
+                    vocab_size = len(token_list)
+
+                if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"):
+                    vocab_size = tokenizer.get_vocab_size()
+                token_lists.append(token_list)
+                vocab_sizes.append(vocab_size)
+
+            if len(tokenizers_build) <= 1:
+                tokenizers_build = tokenizers_build[0]
+                token_lists = token_lists[0]
+                vocab_sizes = vocab_sizes[0]
+
+            kwargs["tokenizer"] = tokenizers_build
+            kwargs["vocab_size"] = vocab_sizes
+            kwargs["token_list"] = token_lists
 
         # build frontend
         frontend = kwargs.get("frontend", None)
@@ -219,7 +260,7 @@
         model_conf = {}
         deep_update(model_conf, kwargs.get("model_conf", {}))
         deep_update(model_conf, kwargs)
-        model = model_class(**model_conf, vocab_size=vocab_size)
+        model = model_class(**model_conf)
 
         # init_param
         init_param = kwargs.get("init_param", None)
@@ -264,6 +305,8 @@
 
     def inference(self, input, input_len=None, model=None, kwargs=None, key=None, **cfg):
         kwargs = self.kwargs if kwargs is None else kwargs
+        if "cache" in kwargs:
+            kwargs.pop("cache")
         deep_update(kwargs, cfg)
         model = self.model if model is None else model
         model.eval()
@@ -337,7 +380,7 @@
         end_vad = time.time()
 
         #  FIX(gcf): concat the vad clips for sense vocie model for better aed
-        if kwargs.get("merge_vad", False):
+        if cfg.get("merge_vad", False):
             for i in range(len(res)):
                 res[i]["value"] = merge_vad(
                     res[i]["value"], kwargs.get("merge_length_s", 15) * 1000

--
Gitblit v1.9.1