From 28ccfbfc51068a663a80764e14074df5edf2b5ba Mon Sep 17 00:00:00 2001
From: kongdeqiang <kongdeqiang960204@163.com>
Date: 星期五, 13 三月 2026 17:41:41 +0800
Subject: [PATCH] 提交

---
 funasr/auto/auto_model.py |  156 ++++++++++++++++++++++++++++++++++++++++-----------
 1 files changed, 121 insertions(+), 35 deletions(-)

diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index a6cd3a6..a864dad 100644
--- a/funasr/auto/auto_model.py
+++ b/funasr/auto/auto_model.py
@@ -14,13 +14,14 @@
 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
 from funasr.download.file import download_from_url
 from funasr.utils.timestamp_tools import timestamp_sentence
 from funasr.utils.timestamp_tools import timestamp_sentence_en
-from funasr.download.download_from_hub import download_model
+from funasr.download.download_model_from_hub import download_model
 from funasr.utils.vad_utils import slice_padding_audio_samples
 from funasr.utils.vad_utils import merge_vad
 from funasr.utils.load_utils import load_audio_text_image_video
@@ -92,7 +93,8 @@
                 if isinstance(data_i, str) and os.path.exists(data_i):
                     key = misc.extract_filename_without_extension(data_i)
                 else:
-                    key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
+                    if key is None:
+                        key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
                 key_list.append(key)
 
     else:  # raw text; audio sample point, fbank; bytes
@@ -110,11 +112,15 @@
 
     def __init__(self, **kwargs):
 
+        try:
+            from funasr.utils.version_checker import check_for_update
+
+            check_for_update(disable=kwargs.get("disable_update", False))
+        except:
+            pass
+
         log_level = getattr(logging, kwargs.get("log_level", "INFO").upper())
         logging.basicConfig(level=log_level)
-
-        if not kwargs.get("disable_log", True):
-            tables.print()
 
         model, kwargs = self.build_model(**kwargs)
 
@@ -141,13 +147,16 @@
         # 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.")
@@ -163,7 +172,8 @@
         self.spk_kwargs = spk_kwargs
         self.model_path = kwargs.get("model_path")
 
-    def build_model(self, **kwargs):
+    @staticmethod
+    def build_model(**kwargs):
         assert "model" in kwargs
         if "model_conf" not in kwargs:
             logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
@@ -172,7 +182,10 @@
         set_all_random_seed(kwargs.get("seed", 0))
 
         device = kwargs.get("device", "cuda")
-        if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
+        if ((device =="cuda" and not torch.cuda.is_available())
+            or (device == "xpu" and not torch.xpu.is_available())
+            or (device == "mps" and not torch.backends.mps.is_available())
+            or kwargs.get("ngpu", 1) == 0):
             device = "cpu"
             kwargs["batch_size"] = 1
         kwargs["device"] = device
@@ -181,21 +194,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)
@@ -209,10 +261,11 @@
         kwargs["frontend"] = frontend
         # build model
         model_class = tables.model_classes.get(kwargs["model"])
+        assert model_class is not None, f'{kwargs["model"]} is not registered'
         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)
@@ -236,6 +289,10 @@
         elif kwargs.get("bf16", False):
             model.to(torch.bfloat16)
         model.to(device)
+
+        if not kwargs.get("disable_log", True):
+            tables.print()
+
         return model, kwargs
 
     def __call__(self, *args, **cfg):
@@ -244,15 +301,30 @@
         res = self.model(*args, kwargs)
         return res
 
-    def generate(self, input, input_len=None, **cfg):
+    def generate(self, input, input_len=None, progress_callback=None, **cfg):
         if self.vad_model is None:
-            return self.inference(input, input_len=input_len, **cfg)
+            return self.inference(
+                input, input_len=input_len, progress_callback=progress_callback, **cfg
+            )
 
         else:
-            return self.inference_with_vad(input, input_len=input_len, **cfg)
+            return self.inference_with_vad(
+                input, input_len=input_len, progress_callback=progress_callback, **cfg
+            )
 
-    def inference(self, input, input_len=None, model=None, kwargs=None, key=None, **cfg):
+    def inference(
+        self,
+        input,
+        input_len=None,
+        model=None,
+        kwargs=None,
+        key=None,
+        progress_callback=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()
@@ -304,15 +376,24 @@
             speed_stats["rtf"] = f"{(time_escape) / batch_data_time:0.3f}"
             description = f"{speed_stats}, "
             if pbar:
-                pbar.update(1)
+                pbar.update(end_idx - beg_idx)
                 pbar.set_description(description)
+            if progress_callback:
+                try:
+                    progress_callback(end_idx, num_samples)
+                except Exception as e:
+                    logging.error(f"progress_callback error: {e}")
             time_speech_total += batch_data_time
             time_escape_total += time_escape
 
         if pbar:
             # pbar.update(1)
             pbar.set_description(f"rtf_avg: {time_escape_total/time_speech_total:0.3f}")
-        torch.cuda.empty_cache()
+
+        device = next(model.parameters()).device
+        if device.type == "cuda":
+            with torch.cuda.device(device):
+                torch.cuda.empty_cache()
         return asr_result_list
 
     def inference_with_vad(self, input, input_len=None, **cfg):
@@ -326,9 +407,11 @@
         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", 15000))
+                res[i]["value"] = merge_vad(
+                    res[i]["value"], kwargs.get("merge_length_s", 15) * 1000
+                )
 
         # step.2 compute asr model
         model = self.model
@@ -368,6 +451,9 @@
 
             if len(sorted_data) > 0 and len(sorted_data[0]) > 0:
                 batch_size = max(batch_size, sorted_data[0][0][1] - sorted_data[0][0][0])
+
+            if kwargs["device"] == "cpu":
+                batch_size = 0
 
             beg_idx = 0
             beg_asr_total = time.time()
@@ -495,8 +581,8 @@
                 sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu())
                 if self.spk_mode == "vad_segment":  # recover sentence_list
                     sentence_list = []
-                    for res, vadsegment in zip(restored_data, vadsegments):
-                        if "timestamp" not in res:
+                    for rest, vadsegment in zip(restored_data, vadsegments):
+                        if "timestamp" not in rest:
                             logging.error(
                                 "Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \
                                            and 'iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'\
@@ -506,8 +592,8 @@
                             {
                                 "start": vadsegment[0],
                                 "end": vadsegment[1],
-                                "sentence": res["text"],
-                                "timestamp": res["timestamp"],
+                                "sentence": rest["text"],
+                                "timestamp": rest["timestamp"],
                             }
                         )
                 elif self.spk_mode == "punc_segment":

--
Gitblit v1.9.1