allow buffered inference up to 10 mins

app.py

@@ -6,16 +6,39 @@ import soundfile as sf
 import tempfile
 import uuid
 
+import torch
+
 from nemo.collections.asr.models import ASRModel
+from nemo.collections.asr.parts.utils.streaming_utils import FrameBatchMultiTaskAED
+from nemo.collections.asr.parts.utils.transcribe_utils import get_buffered_pred_feat_multitaskAED
 
 SAMPLE_RATE = 16000 # Hz
+MAX_AUDIO_MINUTES = 10 # wont try to transcribe if longer than this
 
 model = ASRModel.from_pretrained("nvidia/canary-1b")
 model.eval()
 
+# make sure beam size always 1 for consistency
+model.change_decoding_strategy(None)
+decoding_cfg = model.cfg.decoding
+decoding_cfg.beam.beam_size = 1
+model.change_decoding_strategy(decoding_cfg)
+
+# setup for buffered inference
+model.cfg.preprocessor.dither = 0.0
+model.cfg.preprocessor.pad_to = 0
+
+feature_stride = model.cfg.preprocessor['window_stride']
+model_stride_in_secs = feature_stride * 8 # 8 = model stride, which is 8 for FastConformer
 
-MAX_AUDIO_SECONDS = 40
+frame_asr = FrameBatchMultiTaskAED(
+	asr_model=model,
+	frame_len=40.0,
+	total_buffer=40.0,
+	batch_size=16,
+)
 
+amp_dtype = torch.float16
 
 def convert_audio(audio_filepath, tmpdir, utt_id):
 	"""
@@ -24,21 +47,20 @@ def convert_audio(audio_filepath, tmpdir, utt_id):
 	Returns output filename and duration.
 	"""
 
-	data, sr = librosa.load(audio_filepath)
+	data, sr = librosa.load(audio_filepath, sr=None, mono=False)
 
 	duration = librosa.get_duration(y=data, sr=sr)
 
-	if duration > MAX_AUDIO_SECONDS:
+	if duration / 60.0 > MAX_AUDIO_MINUTES:
 		raise gr.Error(
-			f"This demo can transcribe up to {MAX_AUDIO_SECONDS} seconds of audio."
+			f"This demo can transcribe up to {MAX_AUDIO_MINUTES} minutes of audio. "
+			"If you wish, you may trim the audio using the Audio viewer in Step 1 "
+			"(click on the scissors icon to start trimming audio)."
 		)
 
 	if sr != SAMPLE_RATE:
 		data = librosa.resample(data, orig_sr=sr, target_sr=SAMPLE_RATE)
 
-	# monochannel
-	data = librosa.to_mono(data)
-
 	out_filename = os.path.join(tmpdir, utt_id + '.wav')
 
 	# save output audio
@@ -54,7 +76,6 @@ def transcribe(audio_filepath, src_lang, tgt_lang, pnc):
 	
 	utt_id = uuid.uuid4()
 	with tempfile.TemporaryDirectory() as tmpdir:
-
 		converted_audio_filepath, duration = convert_audio(audio_filepath, tmpdir, str(utt_id))
 
 		# map src_lang and tgt_lang from long versions to short
@@ -102,9 +123,23 @@ def transcribe(audio_filepath, src_lang, tgt_lang, pnc):
 			fout.write(line + '\n')
 
 		# call transcribe, passing in manifest filepath
-		model_output = model.transcribe(manifest_filepath)
-
-	return model_output[0]
+		if duration < 40:
+			output_text = model.transcribe(manifest_filepath)[0]
+		else: # do buffered inference
+			with torch.cuda.amp.autocast(dtype=amp_dtype): # TODO: make it work if no cuda
+				with torch.no_grad():
+					hyps = get_buffered_pred_feat_multitaskAED(
+						frame_asr,
+						model.cfg.preprocessor,
+						model_stride_in_secs,
+						model.device,
+						manifest=manifest_filepath,
+						filepaths=None,
+					)
+
+					output_text = hyps[0].text
+
+	return output_text
 
 # add logic to make sure dropdown menus only suggest valid combos
 def on_src_or_tgt_lang_change(src_lang_value, tgt_lang_value, pnc_value):
@@ -124,15 +159,15 @@ def on_src_or_tgt_lang_change(src_lang_value, tgt_lang_value, pnc_value):
 	tgt_lang_value, and then which states you can go to from there.
 
 			tgt lang
-		- |EN |ES |FR |DE 
-		------------------
-		EN| Y | Y | Y | Y
-		------------------
+			- |EN |ES |FR |DE
+			------------------
+			EN| Y | Y | Y | Y
+			------------------
 		src 	ES| Y | Y |   |
 		lang	------------------
-		FR| Y |   | Y |
-		------------------
-		DE| Y |   |   | Y
+			FR| Y |   | Y |
+			------------------
+			DE| Y |   |   | Y
 	"""
 
 	if src_lang_value == "English" and tgt_lang_value == "English":

requirements.txt

@@ -1 +1 @@
-git+https://github.com/NVIDIA/NeMo.git@r1.23.0#egg=nemo_toolkit[all]
+git+https://github.com/NVIDIA/NeMo.git@61325fe0c70ef4294d8562991f6841d26b238e85#egg=nemo_toolkit[all] # commit from canary_buffer_infer branch