Update app.py

app.py

@@ -1,57 +1,195 @@
+'''
+import gradio as gr
+from transformers import pipeline
+# Load pipelines for Canary ASR, LLama3 QA, and VITS TTS
+asr_pipeline = pipeline("automatic-speech-recognition", model="nvidia/canary-1b", device=0)
+qa_pipeline = pipeline("question-answering", model="LLAMA/llama3-base-qa", tokenizer="LLAMA/llama3-base-qa")
+tts_pipeline = pipeline("text-to-speech", model="patrickvonplaten/vits-large", device=0)
+'''
+
 import gradio as gr
 import json
+import librosa
 import os
+import soundfile as sf
 import tempfile
 import uuid
 from transformers import pipeline
-import librosa
-import soundfile as sf
 
-SAMPLE_RATE = 16000  # Hz
-MAX_AUDIO_SECS = 30  # Maximum duration of audio in seconds
+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_SECS = 30 # wont try to transcribe if longer than this
 src_lang = "en"
 tgt_lang = "en"
-pnc = "no"
+pnc="no"
 
-# Load the ASR pipeline
-asr_pipeline = pipeline("automatic-speech-recognition", model="nvidia/canary-1b")
+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
+
+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):
-    """
-    Convert audio file to 16 kHz mono WAV format.
-    Returns output filename and duration.
-    """
-    data, sr = librosa.load(audio_filepath, sr=None, mono=True)
-    duration = librosa.get_duration(y=data, sr=sr)
+	"""
+	Convert all files to monochannel 16 kHz wav files.
+	Do not convert and raise error if audio too long.
+	Returns output filename and duration.
+	"""
+	data, sr = librosa.load(audio_filepath, sr=None, mono=True)
 
-    if duration > MAX_AUDIO_SECS:
-        raise gr.Error(f"Maximum audio duration exceeded. Please provide an audio file of up to {MAX_AUDIO_SECS} seconds.")
+	duration = librosa.get_duration(y=data, sr=sr)
 
-    if sr != SAMPLE_RATE:
-        data = librosa.resample(data, orig_sr=sr, target_sr=SAMPLE_RATE)
+	if duration > MAX_AUDIO_SECS:
+		raise gr.Error(
+			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)."
+		)
 
-    out_filename = os.path.join(tmpdir, f"{utt_id}.wav")
-    sf.write(out_filename, data, SAMPLE_RATE)
+	if sr != SAMPLE_RATE:
+		data = librosa.resample(data, orig_sr=sr, target_sr=SAMPLE_RATE)
 
-    return out_filename, duration
+	out_filename = os.path.join(tmpdir, utt_id + '.wav')
 
+	# save output audio
+	sf.write(out_filename, data, SAMPLE_RATE)
 
-def transcribe(audio_filepath):
-    if audio_filepath is None:
-        raise gr.Error("Please provide some input audio: either upload an audio file or use the microphone")
+	return out_filename, duration
 
-    utt_id = uuid.uuid4()
-    with tempfile.TemporaryDirectory() as tmpdir:
-        converted_audio_filepath, duration = convert_audio(audio_filepath, tmpdir, str(utt_id))
-        transcribed_text = asr_pipeline(converted_audio_filepath, sampling_rate=SAMPLE_RATE)[0]["transcription"]
 
-    return transcribed_text
+def transcribe(audio_filepath, src_lang, tgt_lang, pnc):
+
+	if audio_filepath is None:
+		raise gr.Error("Please provide some input audio: either upload an audio file or use the microphone")
+	
+	utt_id = uuid.uuid4()
+	with tempfile.TemporaryDirectory() as tmpdir:
+		converted_audio_filepath, duration = convert_audio(audio_filepath, tmpdir, str(utt_id))
+
+		# make manifest file and save
+		manifest_data = {
+			"audio_filepath": converted_audio_filepath,
+			"source_lang": src_lang,
+			"target_lang": tgt_lang,
+			"taskname": taskname,
+			"pnc": pnc,
+			"answer": "predict",
+			"duration": str(duration),
+		}
+
+		manifest_filepath = os.path.join(tmpdir, f'{utt_id}.json')
+
+		with open(manifest_filepath, 'w') as fout:
+			line = json.dumps(manifest_data)
+			fout.write(line + '\n')
+
+		# call transcribe, passing in manifest filepath
+		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
+
+
+
+with gr.Blocks(
+	title="NeMo Canary Model",
+	css="""
+		textarea { font-size: 18px;}
+		#model_output_text_box span {
+			font-size: 18px;
+			font-weight: bold;
+		}
+	""",
+	theme=gr.themes.Default(text_size=gr.themes.sizes.text_lg) # make text slightly bigger (default is text_md )
+) as demo:
+
+	gr.HTML("<h1 style='text-align: center'>NeMo Canary model: Transcribe & Translate audio</h1>")
+
+	with gr.Row():
+		with gr.Column():
+			gr.HTML(
+				"<p><b>Step 1:</b> Upload an audio file or record with your microphone.</p>"
+
+				"<p style='color: #A0A0A0;'>This demo supports audio files up to 10 mins long. "
+				"You can transcribe longer files locally with this NeMo "
+				"<a href='https://github.com/NVIDIA/NeMo/blob/main/examples/asr/speech_multitask/speech_to_text_aed_chunked_infer.py'>script</a>.</p>"
+			)
+
+			audio_file = gr.Audio(sources=["microphone", "upload"], type="filepath")
+
+			gr.HTML("<p><b>Step 2:</b> Choose the input and output language.</p>")
+
+			
+		with gr.Column():
+
+			gr.HTML("<p><b>Step 3:</b> Run the model.</p>")
+
+			go_button = gr.Button(
+				value="Run model",
+				variant="primary", # make "primary" so it stands out (default is "secondary")
+			)
+
+			model_output_text_box = gr.Textbox(
+				label="Model Output",
+				elem_id="model_output_text_box",
+			)
+
+	with gr.Row():
+
+		gr.HTML(
+			"<p style='text-align: center'>"
+				"🐤 <a href='https://huggingface.co/nvidia/canary-1b' target='_blank'>Canary model</a> | "
+				"🧑‍💻 <a href='https://github.com/NVIDIA/NeMo' target='_blank'>NeMo Repository</a>"
+			"</p>"
+		)
 
+	go_button.click(
+		fn=transcribe, 
+		inputs = [audio_file],
+		outputs = [model_output_text_box]
+	)
 
-with gr.Interface(transcribe, gr.inputs.Audio(), "text", title="ASR with NeMo Canary Model") as iface:
-    iface.launch()
 
+demo.queue()
+demo.launch()
 
 '''