python/FunASR-XL.git

			@@ -14,12 +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.download.download_from_hub import download_model
			from funasr.utils.timestamp_tools import timestamp_sentence_en
			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
			@@ -91,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
			@@ -109,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)

			@@ -140,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.")
			@@ -162,7 +170,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")))
			@@ -180,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)
			@@ -208,11 +256,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.to(device)
			model = model_class(**model_conf)

			# init_param
			init_param = kwargs.get("init_param", None)
			@@ -233,6 +281,13 @@
			# fp16
			if kwargs.get("fp16", False):
			model.to(torch.float16)
			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):
			@@ -250,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()
			@@ -301,7 +358,7 @@
			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)
			time_speech_total += batch_data_time
			time_escape_total += time_escape
			@@ -309,7 +366,8 @@
			if pbar:
			# pbar.update(1)
			pbar.set_description(f"rtf_avg: {time_escape_total/time_speech_total:0.3f}")
			torch.cuda.empty_cache()
			with torch.cuda.device(next(model.parameters()).device):
			torch.cuda.empty_cache()
			return asr_result_list

			def inference_with_vad(self, input, input_len=None, **cfg):
			@@ -323,9 +381,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
			@@ -365,6 +425,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()
			@@ -427,6 +490,10 @@
			# f"time_speech_total_per_sample: {time_speech_total_per_sample: 0.3f}, "
			# f"time_escape_total_per_sample: {time_escape_total_per_sample:0.3f}")

			if len(results_sorted) != n:
			results_ret_list.append({"key": key, "text": "", "timestamp": []})
			logging.info("decoding, utt: {}, empty result".format(key))
			continue
			restored_data = [0] * n
			for j in range(n):
			index = sorted_data[j][1]
			@@ -460,23 +527,20 @@
			else:
			result[k] += restored_data[j][k]

			if not len(result["text"].strip()):
			continue
			return_raw_text = kwargs.get("return_raw_text", False)
			# step.3 compute punc model
			raw_text = None
			if self.punc_model is not None:
			if not len(result["text"].strip()):
			if return_raw_text:
			result["raw_text"] = ""
			else:
			deep_update(self.punc_kwargs, cfg)
			punc_res = self.inference(
			result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg
			)
			raw_text = copy.copy(result["text"])
			if return_raw_text:
			result["raw_text"] = raw_text
			result["text"] = punc_res[0]["text"]
			else:
			raw_text = None
			deep_update(self.punc_kwargs, cfg)
			punc_res = self.inference(
			result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg
			)
			raw_text = copy.copy(result["text"])
			if return_raw_text:
			result["raw_text"] = raw_text
			result["text"] = punc_res[0]["text"]

			# speaker embedding cluster after resorted
			if self.spk_model is not None and kwargs.get("return_spk_res", True):
			@@ -491,8 +555,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'\
			@@ -502,8 +566,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":
			@@ -513,24 +577,40 @@
			and 'iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'\
			can predict timestamp, and speaker diarization relies on timestamps."
			)
			sentence_list = timestamp_sentence(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			if kwargs.get("en_post_proc", False):
			sentence_list = timestamp_sentence_en(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			else:
			sentence_list = timestamp_sentence(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			distribute_spk(sentence_list, sv_output)
			result["sentence_info"] = sentence_list
			elif kwargs.get("sentence_timestamp", False):
			if not len(result["text"].strip()):
			sentence_list = []
			else:
			sentence_list = timestamp_sentence(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			if kwargs.get("en_post_proc", False):
			sentence_list = timestamp_sentence_en(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			else:
			sentence_list = timestamp_sentence(
			punc_res[0]["punc_array"],
			result["timestamp"],
			raw_text,
			return_raw_text=return_raw_text,
			)
			result["sentence_info"] = sentence_list
			if "spk_embedding" in result:
			del result["spk_embedding"]
			@@ -582,12 +662,6 @@
			)

			with torch.no_grad():

			if type == "onnx":
			export_dir = export_utils.export_onnx(model=model, data_in=data_list, **kwargs)
			else:
			export_dir = export_utils.export_torchscripts(
			model=model, data_in=data_list, **kwargs
			)
			export_dir = export_utils.export(model=model, data_in=data_list, **kwargs)

			return export_dir