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):
			"""
			@@ -139,7 +141,7 @@
			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", 1) == 0:
			device = "cpu"
			@@ -159,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
			@@ -206,13 +207,14 @@
			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)
			@@ -225,6 +227,7 @@
			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 kwargs.get("data_type", None) == "fbank": # fbank
			batch["data_in"] = data_batch[0]
			batch["data_lengths"] = input_len
			@@ -376,16 +379,22 @@
			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, disable_pbar=True, **cfg)
			import copy; raw_text = copy.copy(result["text"])
			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 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=kwargs.get('preset_spk_num', None))
			@@ -394,22 +403,32 @@
			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']})
			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':
			sentence_list = timestamp_sentence(punc_res[0]['punc_array'], \
			result['timestamp'], \
			result['raw_text'])
			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'], \
			result['raw_text'])
			sentence_list = timestamp_sentence(punc_res[0]['punc_array'],
			result['timestamp'],
			raw_text,
			return_raw_text=return_raw_text)
			result['sentence_info'] = sentence_list
			del result['spk_embedding']
			if "spk_embedding" in result: del result['spk_embedding']

			result["key"] = key
			results_ret_list.append(result)