From 5a8f37908469d9550f905ba0876c7c4e6f9b8026 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 21 十二月 2023 21:08:46 +0800
Subject: [PATCH] vad + asr

---
 funasr/models/bici_paraformer/model.py |   88 +++++++++++++++++++++++++-------------------
 1 files changed, 50 insertions(+), 38 deletions(-)

diff --git a/funasr/models/bici_paraformer/model.py b/funasr/models/bici_paraformer/model.py
index 52eac87..03c8896 100644
--- a/funasr/models/bici_paraformer/model.py
+++ b/funasr/models/bici_paraformer/model.py
@@ -29,6 +29,7 @@
 from funasr.utils.timestamp_tools import ts_prediction_lfr6_standard
 from funasr.register import tables
 from funasr.models.ctc.ctc import CTC
+from funasr.utils.timestamp_tools import time_stamp_sentence
 
 from funasr.models.paraformer.model import Paraformer
 
@@ -211,10 +212,11 @@
 		
 		loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
 		return loss, stats, weight
-	
+
+
 	def generate(self,
-	             data_in: list,
-	             data_lengths: list = None,
+	             data_in,
+	             data_lengths=None,
 	             key: list = None,
 	             tokenizer=None,
 	             frontend=None,
@@ -230,17 +232,23 @@
 			self.nbest = kwargs.get("nbest", 1)
 		
 		meta_data = {}
-		# extract fbank feats
-		time1 = time.perf_counter()
-		audio_sample_list = load_audio(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000))
-		time2 = time.perf_counter()
-		meta_data["load_data"] = f"{time2 - time1:0.3f}"
-		speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"),
-		                                       frontend=self.frontend)
-		time3 = time.perf_counter()
-		meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
-		meta_data[
-			"batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
+		if isinstance(data_in, torch.Tensor):  # fbank
+			speech, speech_lengths = data_in, data_lengths
+			if len(speech.shape) < 3:
+				speech = speech[None, :, :]
+			if speech_lengths is None:
+				speech_lengths = speech.shape[1]
+		else:
+			# extract fbank feats
+			time1 = time.perf_counter()
+			audio_sample_list = load_audio(data_in, fs=frontend.fs, audio_fs=kwargs.get("fs", 16000))
+			time2 = time.perf_counter()
+			meta_data["load_data"] = f"{time2 - time1:0.3f}"
+			speech, speech_lengths = extract_fbank(audio_sample_list, data_type=kwargs.get("data_type", "sound"),
+			                                       frontend=frontend)
+			time3 = time.perf_counter()
+			meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
+			meta_data["batch_data_time"] = speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
 		
 		speech.to(device=kwargs["device"]), speech_lengths.to(device=kwargs["device"])
 		
@@ -261,9 +269,8 @@
 		decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1]
 		
 		# BiCifParaformer, test no bias cif2
-
 		_, _, us_alphas, us_peaks = self.calc_predictor_timestamp(encoder_out, encoder_out_lens,
-			                                                                    pre_token_length)
+		                                                          pre_token_length)
 		
 		results = []
 		b, n, d = decoder_out.size()
@@ -302,27 +309,32 @@
 				# remove blank symbol id, which is assumed to be 0
 				token_int = list(filter(lambda x: x != self.eos and x != self.sos and x != self.blank_id, token_int))
 				
-				# Change integer-ids to tokens
-				token = tokenizer.ids2tokens(token_int)
-				text = tokenizer.tokens2text(token)
-				
-				_, timestamp = ts_prediction_lfr6_standard(us_alphas[i][:encoder_out_lens[i] * 3],
-				                                           us_peaks[i][:encoder_out_lens[i] * 3],
-				                                           copy.copy(token),
-				                                           vad_offset=kwargs.get("begin_time", 0))
-				
-				text_postprocessed, time_stamp_postprocessed, word_lists = postprocess_utils.sentence_postprocess(token, timestamp)
-				
-				result_i = {"key": key[i], "token": token, "text": text, "text_postprocessed": text_postprocessed,
-				            "time_stamp_postprocessed": time_stamp_postprocessed,
-				            "word_lists": word_lists
-				            }
-				results.append(result_i)
-				
-				if ibest_writer is not None:
-					ibest_writer["token"][key[i]] = " ".join(token)
-					ibest_writer["text"][key[i]] = text
-					ibest_writer["text_postprocessed"][key[i]] = text_postprocessed
+				if tokenizer is not None:
+					# Change integer-ids to tokens
+					token = tokenizer.ids2tokens(token_int)
+					text = tokenizer.tokens2text(token)
 					
+					_, timestamp = ts_prediction_lfr6_standard(us_alphas[i][:encoder_out_lens[i] * 3],
+					                                           us_peaks[i][:encoder_out_lens[i] * 3],
+					                                           copy.copy(token),
+					                                           vad_offset=kwargs.get("begin_time", 0))
+					
+					text_postprocessed, time_stamp_postprocessed, word_lists = postprocess_utils.sentence_postprocess(
+						token, timestamp)
+					sentences = time_stamp_sentence(None, time_stamp_postprocessed, text_postprocessed)
+					result_i = {"key": key[i], "token": token, "text": text, "text_postprocessed": text_postprocessed,
+					            "timestamp": time_stamp_postprocessed,
+					            "word_lists": word_lists,
+					            "sentences": sentences
+					            }
+					
+					if ibest_writer is not None:
+						ibest_writer["token"][key[i]] = " ".join(token)
+						ibest_writer["text"][key[i]] = text
+						ibest_writer["timestamp"][key[i]] = time_stamp_postprocessed
+						ibest_writer["text_postprocessed"][key[i]] = text_postprocessed
+				else:
+					result_i = {"key": key[i], "token_int": token_int}
+				results.append(result_i)
 		
-		return results, meta_data
+		return results, meta_data
\ No newline at end of file

--
Gitblit v1.9.1