From a8591060d3889cd7a72841fa32a7ee64b49db1d2 Mon Sep 17 00:00:00 2001
From: msgk <zxr935867802@outlook.com>
Date: 星期五, 14 二月 2025 14:16:51 +0800
Subject: [PATCH] fix(spk): 修复 speaker embedding 集群后的重新排序问题
---
funasr/auto/auto_model.py | 124 ++++++++++++++++++++++++++++++++++-------
1 files changed, 103 insertions(+), 21 deletions(-)
diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py
index 75324dc..f5cbe01 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,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.")
@@ -187,21 +191,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 +262,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 +307,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()
@@ -323,7 +368,11 @@
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):
@@ -337,7 +386,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
@@ -500,8 +549,41 @@
# speaker embedding cluster after resorted
if self.spk_model is not None and kwargs.get("return_spk_res", True):
- if raw_text is None:
- logging.error("Missing punc_model, which is required by spk_model.")
+ # 1. 鍏堟鏌ユ椂闂存埑
+ has_timestamp = (
+ hasattr(self.model, "internal_punc") or
+ self.punc_model is not None or
+ "timestamp" in result
+ )
+
+ if not has_timestamp:
+ logging.error("Need timestamp support...")
+ return results_ret_list
+
+ # 2. 鍒濆鍖� punc_res
+ punc_res = None
+
+ # 3. 鏍规嵁涓嶅悓鎯呭喌璁剧疆 punc_res
+ if hasattr(self.model, "internal_punc"):
+ punc_res = [{
+ "text": result["text"],
+ "punc_array": result.get("punc_array", []),
+ "timestamp": result.get("timestamp", [])
+ }]
+ elif self.punc_model is not None:
+ punc_res = self.inference(
+ result["text"],
+ model=self.punc_model,
+ kwargs=self.punc_kwargs,
+ **cfg
+ )
+ else:
+ # 濡傛灉鍙湁鏃堕棿鎴筹紝鍒涘缓涓�涓熀鏈殑 punc_res
+ punc_res = [{
+ "text": result["text"],
+ "punc_array": [], # 绌虹殑鏍囩偣鏁扮粍
+ "timestamp": result["timestamp"]
+ }]
all_segments = sorted(all_segments, key=lambda x: x[0])
spk_embedding = result["spk_embedding"]
labels = self.cb_model(
--
Gitblit v1.9.1