| | |
|---|
| | | 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): |
|---|
| | | """ |
|---|
| | |
|---|
| | | 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" |
|---|
| | |
|---|
| | | 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.to(device) |
|---|
| | | |
|---|
| | | # init_param |
|---|
| | |
|---|
| | | # 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) |
|---|
| | |
|---|
| | | 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 |
|---|
| | | |
|---|
| | | time1 = time.perf_counter() |
|---|
| | | with torch.no_grad(): |
|---|
| | | pdb.set_trace() |
|---|
| | | results, meta_data = model.inference(**batch, **kwargs) |
|---|
| | | time2 = time.perf_counter() |
|---|
| | | |
|---|
| | | pdb.set_trace() |
|---|
| | | asr_result_list.extend(results) |
|---|
| | | |
|---|
| | | # batch_data_time = time_per_frame_s * data_batch_i["speech_lengths"].sum().item() |
|---|
| | |
|---|
| | | 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) |
|---|
| | | 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 |
|---|
| | |
|---|
| | | for res, vadsegment in zip(restored_data, vadsegments): |
|---|
| | | 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']}) |
|---|
| | | 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) |
|---|
| | | 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) |
|---|
| | | 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'] |
|---|
| | | |
|---|
