Create app.py

app.py

@@ -0,0 +1,181 @@
+
+import subprocess
+import torch
+import gradio as gr
+import yt_dlp
+import pandas as pd
+from nemo.collections.asr.models import ASRModel
+from nemo_align import align_tdt_to_ctc_timestamps
+import os
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def process_audio(input_file, output_file):
+    command = [
+        'sox', input_file,
+        output_file,
+        'channels', '1',
+        'rate', '16000'
+    ]
+    try:
+        subprocess.run(command, check=True)
+        return output_file
+    except:
+        raise gr.Error("Failed to convert audio to single channel and sampling rate to 16000")
+
+def get_dataframe_segments(segments):
+    df = pd.DataFrame(columns=['start_time', 'end_time', 'text'])
+    if len(segments) == 0:
+        df.loc[0] = 0, 0, ''
+        return df
+    
+    for segment in segments:
+        text, start_time, end_time = segment
+        if len(text)>0:
+            df.loc[len(df)] = round(start_time, 2), round(end_time, 2), text
+    
+    return df
+
+
+def get_video_info(url):
+    ydl_opts = {
+        'quiet': True,
+        'skip-download': True,
+    }
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        try:
+            info = ydl.extract_info(url, download=False)
+        except:
+            raise gr.Error("Failed to extract video info from Youtube")
+        return info
+
+def download_audio(url):
+    ydl_opts = {
+        'format': 'bestaudio/best,channels:1',
+        'quiet': True,
+        'outtmpl': 'audio_file',
+        'postprocessors': [{
+            'key': 'FFmpegExtractAudio',
+            'preferredcodec': 'flac',
+            'preferredquality': '192',
+        }],
+    }
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        try:
+            ydl.download([url])
+        except yt_dlp.utils.DownloadError as err:
+            raise gr.Error(str(err))
+
+    return process_audio('audio_file.flac', 'processed_file.flac')
+
+
+def get_audio_from_youtube(url):
+    info = get_video_info(url)
+    duration = info.get('duration', 0)  # Duration in seconds
+    video_id = info.get('id',None)
+
+    html = f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
+
+    if duration > 2*60*60: # 2 hrs change later based on GPU
+        return gr.Error(str("For GPU {}, single pass maximum audio can be 2hrs"))
+    else:
+        return download_audio(url), html
+
+
+def get_transcripts(audio_path, model):
+    with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
+        with torch.inference_mode():
+            text = model.transcribe(audio=[audio_path], )
+    return text
+
+def pick_asr_model():
+    model = 'nvidia/parakeet-tdt_ctc-1.1b'
+    asr_model = ASRModel.from_pretrained(model).to(device)
+    asr_model.eval()
+    return asr_model
+
+asr_model = pick_asr_model()
+
+def run_nemo_models(url, microphone, audio_path):
+    html = None
+    if url is None or len(url)<2:
+        path1 = microphone if microphone else audio_path
+    else:
+        gr.Info("Downloading and processing audio from Youtube")
+        path1, html = get_audio_from_youtube(url)
+       
+    gr.Info("Running NeMo Model")
+    text = get_transcripts(path1, asr_model)
+
+    segments = align_tdt_to_ctc_timestamps(text, asr_model, path1)
+    
+    df = get_dataframe_segments(segments)
+    
+    return df, html
+
+def clear_youtube_link():
+    # Remove .flac files in current directory
+    file_list = os.listdir()
+    for file in file_list:
+        if file.endswith(".flac"):
+            os.remove(file)
+    
+    return None
+
+
+# def run_speaker_diarization()
+
+with gr.Blocks(
+	title="NeMo Parakeet 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'>Transcription with timestamps using Parakeet TDT-CTC</h1>")
+    gr.Markdown('''
+    Choose between different sources of audio (Microphone, Audio File, Youtube Video) to transcribe along with timestamps. 
+    
+    Parakeet models with limited attention are quite fast due to their limited attention mechanism. The current model with 1.1B parameters can transcribe very long audios upto 11 hrs on A6000 GPU in a single pass.
+                
+    Model used: [nvidia/parakeet-tdt_ctc-1.1b](https://huggingface.co/nvidia/parakeet-tdt_ctc-1.1b).
+    ''')
+    # This block is for reading audio from MIC
+    with gr.Tab('Audio from Youtube'):
+        with gr.Row():
+            yt_link = gr.Textbox(value=None, label='Enter Youtube Link', type='text')
+            yt_render = gr.HTML()
+    
+    with gr.Tab('Audio From File'):
+        file_input = gr.Audio(sources='upload', label='Upload Audio', type='filepath')
+
+    with gr.Tab('Audio From Microphone'):
+        mic_input = gr.Audio(sources='microphone', label='Record Audio', type='filepath')
+        
+ 
+    # b1 = gr.Button("Get Transcription with Punctuation and Capitalization")
+
+    gr.Markdown('''Speech Recognition''')
+
+    # text_output = gr.Textbox(label='Transcription', type='text')
+
+    b2 = gr.Button("Get timestamps with text")
+
+    time_stamp = gr.DataFrame(wrap=True, label='Speech Recognition with TimeStamps',
+        row_count=(1, "dynamic"), headers=['start_time', 'end_time', 'text'])
+
+    # b1.click(run_nemo_models, inputs=[file_input, mic_input, yt_link], outputs=[text_output, yt_render])
+
+    b2.click(run_nemo_models, inputs=[yt_link, file_input, mic_input], outputs=[time_stamp, yt_render]).then(
+        clear_youtube_link, None, yt_link, queue=False) #here clean up passing None to audio.
+
+    demo.queue(True)
+    demo.launch(share=True, debug=True)
+