python/FunASR-XL.git

			@@ -1,14 +1,13 @@
			import json
			import time
			import copy
			import torch
			import hydra
			import random
			import string
			import logging
			import os.path
			import numpy as np
			from tqdm import tqdm
			from omegaconf import DictConfig, OmegaConf, ListConfig

			from funasr.register import tables
			from funasr.utils.load_utils import load_bytes
			@@ -17,11 +16,14 @@
			from funasr.utils.vad_utils import slice_padding_audio_samples
			from funasr.train_utils.set_all_random_seed import set_all_random_seed
			from funasr.train_utils.load_pretrained_model import load_pretrained_model
			from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
			from funasr.utils.load_utils import load_audio_text_image_video
			from funasr.utils.timestamp_tools import timestamp_sentence
			from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
			from funasr.models.campplus.cluster_backend import ClusterBackend

			try:
			from funasr.models.campplus.cluster_backend import ClusterBackend
			except:
			print("If you want to use the speaker diarization, please `pip install hdbscan`")
			import pdb

			def prepare_data_iterator(data_in, input_len=None, data_type=None, key=None):
			"""
			@@ -88,7 +90,8 @@
			class AutoModel:

			def __init__(self, **kwargs):
			tables.print()
			if not kwargs.get("disable_log", False):
			tables.print()

			model, kwargs = self.build_model(**kwargs)

			@@ -120,10 +123,6 @@
			if spk_mode not in ["default", "vad_segment", "punc_segment"]:
			logging.error("spk_mode should be one of default, vad_segment and punc_segment.")
			self.spk_mode = spk_mode
			self.preset_spk_num = kwargs.get("preset_spk_num", None)
			if self.preset_spk_num:
			logging.warning("Using preset speaker number: {}".format(self.preset_spk_num))
			logging.warning("Many to print when using speaker model...")

			self.kwargs = kwargs
			self.model = model
			@@ -134,8 +133,6 @@
			self.spk_model = spk_model
			self.spk_kwargs = spk_kwargs
			self.model_path = kwargs.get("model_path")



			def build_model(self, **kwargs):
			assert "model" in kwargs
			@@ -144,9 +141,9 @@
			kwargs = download_model(**kwargs)

			set_all_random_seed(kwargs.get("seed", 0))


			device = kwargs.get("device", "cuda")
			if not torch.cuda.is_available() or kwargs.get("ngpu", 0):
			if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
			device = "cpu"
			kwargs["batch_size"] = 1
			kwargs["device"] = device
			@@ -164,19 +161,18 @@
			vocab_size = len(tokenizer.token_list)
			else:
			vocab_size = -1

			# build frontend
			frontend = kwargs.get("frontend", None)

			if frontend is not None:
			frontend_class = tables.frontend_classes.get(frontend)
			frontend = frontend_class(**kwargs["frontend_conf"])
			kwargs["frontend"] = frontend
			kwargs["input_size"] = frontend.output_size()


			# build model
			model_class = tables.model_classes.get(kwargs["model"])
			model = model_class(kwargs, kwargs["model_conf"], vocab_size=vocab_size)
			model.eval()
			model.to(device)

			# init_param
			@@ -200,8 +196,6 @@
			res = self.model(*args, kwargs)
			return res



			def generate(self, input, input_len=None, **cfg):
			if self.vad_model is None:
			return self.inference(input, input_len=input_len, **cfg)
			@@ -213,18 +207,19 @@
			kwargs = self.kwargs if kwargs is None else kwargs
			kwargs.update(cfg)
			model = self.model if model is None else model
			model.eval()

			batch_size = kwargs.get("batch_size", 1)
			# if kwargs.get("device", "cpu") == "cpu":
			# batch_size = 1

			key_list, data_list = prepare_data_iterator(input, input_len=input_len, data_type=kwargs.get("data_type", None), key=key)


			speed_stats = {}
			asr_result_list = []
			num_samples = len(data_list)
			disable_pbar = kwargs.get("disable_pbar", False)
			pbar = tqdm(colour="blue", total=num_samples+1, dynamic_ncols=True) if not disable_pbar else None
			pbar = tqdm(colour="blue", total=num_samples, dynamic_ncols=True) if not disable_pbar else None
			time_speech_total = 0.0
			time_escape_total = 0.0
			for beg_idx in range(0, num_samples, batch_size):
			@@ -232,7 +227,8 @@
			data_batch = data_list[beg_idx:end_idx]
			key_batch = key_list[beg_idx:end_idx]
			batch = {"data_in": data_batch, "key": key_batch}
			if (end_idx - beg_idx) == 1 and isinstance(data_batch[0], torch.Tensor): # fbank

			if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank
			batch["data_in"] = data_batch[0]
			batch["data_lengths"] = input_len

			@@ -261,7 +257,7 @@
			time_escape_total += time_escape

			if pbar:
			pbar.update(1)
			# pbar.update(1)
			pbar.set_description(f"rtf_avg: {time_escape_total/time_speech_total:0.3f}")
			torch.cuda.empty_cache()
			return asr_result_list
			@@ -286,10 +282,10 @@

			key_list, data_list = prepare_data_iterator(input, input_len=input_len, data_type=kwargs.get("data_type", None))
			results_ret_list = []
			time_speech_total_all_samples = 0.0
			time_speech_total_all_samples = 1e-6

			beg_total = time.time()
			pbar_total = tqdm(colour="red", total=len(res) + 1, dynamic_ncols=True)
			pbar_total = tqdm(colour="red", total=len(res), dynamic_ncols=True)
			for i in range(len(res)):
			key = res[i]["key"]
			vadsegments = res[i]["value"]
			@@ -314,11 +310,11 @@
			time_speech_total_per_sample = speech_lengths/16000
			time_speech_total_all_samples += time_speech_total_per_sample

			pbar_sample = tqdm(colour="blue", total=n + 1, dynamic_ncols=True)
			# pbar_sample = tqdm(colour="blue", total=n, dynamic_ncols=True)

			all_segments = []
			for j, _ in enumerate(range(0, n)):
			pbar_sample.update(1)
			# pbar_sample.update(1)
			batch_size_ms_cum += (sorted_data[j][0][1] - sorted_data[j][0][0])
			if j < n - 1 and (
			batch_size_ms_cum + sorted_data[j + 1][0][1] - sorted_data[j + 1][0][0]) < batch_size and (
			@@ -329,8 +325,6 @@
			speech_j, speech_lengths_j = slice_padding_audio_samples(speech, speech_lengths, sorted_data[beg_idx:end_idx])
			results = self.inference(speech_j, input_len=None, model=model, kwargs=kwargs, disable_pbar=True, **cfg)
			if self.spk_model is not None:


			# compose vad segments: [[start_time_sec, end_time_sec, speech], [...]]
			for _b in range(len(speech_j)):
			vad_segments = [[sorted_data[beg_idx:end_idx][_b][0][0]/1000.0,
			@@ -339,22 +333,20 @@
			segments = sv_chunk(vad_segments)
			all_segments.extend(segments)
			speech_b = [i[2] for i in segments]
			spk_res = self.inference(speech_b, input_len=None, model=self.spk_model, kwargs=kwargs, **cfg)
			spk_res = self.inference(speech_b, input_len=None, model=self.spk_model, kwargs=kwargs, disable_pbar=True, **cfg)
			results[_b]['spk_embedding'] = spk_res[0]['spk_embedding']
			beg_idx = end_idx
			if len(results) < 1:
			continue
			results_sorted.extend(results)



			end_asr_total = time.time()
			time_escape_total_per_sample = end_asr_total - beg_asr_total
			pbar_sample.set_description(f"rtf_avg_per_sample: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, "
			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}")
			# end_asr_total = time.time()
			# time_escape_total_per_sample = end_asr_total - beg_asr_total
			# pbar_sample.update(1)
			# pbar_sample.set_description(f"rtf_avg_per_sample: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, "
			# 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}")


			restored_data = [0] * n
			for j in range(n):
			index = sorted_data[j][1]
			@@ -377,7 +369,7 @@
			result[k] = restored_data[j][k]
			else:
			result[k] = torch.cat([result[k], restored_data[j][k]], dim=0)
			elif k == 'text':
			elif 'text' in k:
			if k not in result:
			result[k] = restored_data[j][k]
			else:
			@@ -387,43 +379,71 @@
			result[k] = restored_data[j][k]
			else:
			result[k] += restored_data[j][k]


			return_raw_text = kwargs.get('return_raw_text', False)
			# step.3 compute punc model
			if self.punc_model is not None:
			self.punc_kwargs.update(cfg)
			punc_res = self.inference(result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg)
			punc_res = self.inference(result["text"], model=self.punc_model, kwargs=self.punc_kwargs, disable_pbar=True, **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

			# speaker embedding cluster after resorted
			if self.spk_model is not None:
			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.")
			all_segments = sorted(all_segments, key=lambda x: x[0])
			spk_embedding = result['spk_embedding']
			labels = self.cb_model(spk_embedding.cpu(), oracle_num=self.preset_spk_num)
			del result['spk_embedding']
			labels = self.cb_model(spk_embedding.cpu(), oracle_num=kwargs.get('preset_spk_num', None))
			# del result['spk_embedding']
			sv_output = postprocess(all_segments, None, labels, spk_embedding.cpu())
			if self.spk_mode == 'vad_segment':
			if self.spk_mode == 'vad_segment': # recover sentence_list
			sentence_list = []
			for res, vadsegment in zip(restored_data, vadsegments):
			sentence_list.append({"start": vadsegment[0],\
			"end": vadsegment[1],
			"sentence": res['raw_text'],
			"timestamp": res['timestamp']})
			else: # punc_segment
			sentence_list = timestamp_sentence(punc_res[0]['punc_array'], \
			result['timestamp'], \
			result['text'])
			if 'timestamp' not in res:
			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'\
			can predict timestamp, and speaker diarization relies on timestamps.")
			sentence_list.append({"start": vadsegment[0],
			"end": vadsegment[1],
			"sentence": res['text'],
			"timestamp": res['timestamp']})
			elif self.spk_mode == 'punc_segment':
			if 'timestamp' not in result:
			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'\
			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)
			distribute_spk(sentence_list, sv_output)
			result['sentence_info'] = sentence_list
			elif kwargs.get("sentence_timestamp", False):
			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']

			result["key"] = key
			results_ret_list.append(result)
			end_asr_total = time.time()
			time_escape_total_per_sample = end_asr_total - beg_asr_total
			pbar_total.update(1)

			pbar_total.update(1)
			end_total = time.time()
			time_escape_total_all_samples = end_total - beg_total
			pbar_total.set_description(f"rtf_avg_all_samples: {time_escape_total_all_samples / time_speech_total_all_samples:0.3f}, "
			f"time_speech_total_all_samples: {time_speech_total_all_samples: 0.3f}, "
			f"time_escape_total_all_samples: {time_escape_total_all_samples:0.3f}")
			pbar_total.set_description(f"rtf_avg: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, "
			f"time_speech: {time_speech_total_per_sample: 0.3f}, "
			f"time_escape: {time_escape_total_per_sample:0.3f}")


			# end_total = time.time()
			# time_escape_total_all_samples = end_total - beg_total
			# print(f"rtf_avg_all: {time_escape_total_all_samples / time_speech_total_all_samples:0.3f}, "
			# f"time_speech_all: {time_speech_total_all_samples: 0.3f}, "
			# f"time_escape_all: {time_escape_total_all_samples:0.3f}")
			return results_ret_list