Adding YT transcription code

app.py

@@ -1,66 +1,83 @@
 import gradio as gr
-import torch
-from transformers import WhisperForConditionalGeneration, WhisperProcessor
-import librosa
-import soundfile
+from transformers import pipeline
+import pytube as pt
 
 MODEL_NAME = "openai/whisper-small"
-lang = "ja"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME).to(device)
-processor = WhisperProcessor.from_pretrained(MODEL_NAME)
-
-
-def load_and_fix_data(input_file):  
-    speech, sample_rate = librosa.load(input_file)
-    if len(speech.shape) > 1: 
-        speech = speech[:,0] + speech[:,1]
-    if sample_rate !=16000:
-        speech = librosa.resample(speech, sample_rate,16000)
-    speech = librosa.to_mono(speech)
-    return speech
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=MODEL_NAME,
+    chunk_length_s=30,
+    device=device,
+)
 
 
-def transcribe(Microphone, File_Upload):
+def transcribe(microphone, file_upload):
     warn_output = ""
-    if (Microphone is not None) and (File_Upload is not None):
-        warn_output = "WARNING: You've uploaded an audio file and used the microphone. " \
-                      "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
-        file = Microphone
+    if (microphone is not None) and (file_upload is not None):
+        warn_output = (
+            "WARNING: You've uploaded an audio file and used the microphone. "
+            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
+        )
+        file = microphone
 
-    elif (Microphone is None) and (File_Upload is None):
+    elif (microphone is None) and (file_upload is None):
         return "ERROR: You have to either use the microphone or upload an audio file"
 
-    elif Microphone is not None:
-        file = Microphone
-    else:
-        file = File_Upload
-    
-    speech_data = load_and_fix_data(file)
+    file = microphone if microphone is not None else file_upload
 
-    inputs = processor(speech_data, return_tensors="pt", sampling_rate=16_000).input_features.to(device)
-    forced_decoder_ids = processor.get_decoder_prompt_ids(language=lang, task="transcribe")
-
-    predicted_ids = model.generate(inputs, max_length=480_000, forced_decoder_ids=forced_decoder_ids)
-    text = processor.batch_decode(predicted_ids, skip_special_tokens=True, normalize=True)[0]    
+    text = pipe(file)["text"]
 
     return warn_output + text
 
 
-iface = gr.Interface(
+def yt_transcribe(yt_url):
+
+    yt = pt.YouTube(yt_url)
+    stream = yt.streams.filter(only_audio=True)[0]
+    stream.download(filename="audio.mp3")
+
+    text = pipe("audio.mp3")["text"]
+
+    return text
+
+
+demo = gr.Blocks()
+
+mf_transcribe = gr.Interface(
     fn=transcribe,
     inputs=[
-        gr.inputs.Audio(source="microphone", type='filepath', optional=True),
-        gr.inputs.Audio(source="upload", type='filepath', optional=True),
+        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
+        gr.inputs.Audio(source="upload", type="filepath", optional=True),
     ],
     outputs="text",
     layout="horizontal",
     theme="huggingface",
-    title="[WFTE] Whisper model showcase",
-    description="Demo for showcasing fine-tuned OpenAI Whisper models from WFTE.",
-    allow_flagging='never',
+    title="Whisper Audio Transcribe",
+    description="Transcribe long audio/ microphone input (powered by 🤗transformers) with a click of a button!",
+    allow_flagging="never",
 )
 
-iface.launch(enable_queue=True)
+yt_transcribe = gr.Interface(
+    fn=yt_transcribe,
+    inputs=[
+        gr.inputs.Textbox(
+            lines=1, placeholder="Paste a URL to YT video here", label="yt_url"
+        )
+    ],
+    outputs="text",
+    layout="horizontal",
+    theme="huggingface",
+    title="Whisper YT Transcribe",
+    description="Transcribe long YouTube videos (powered by 🤗transformers) with a click of a button!",
+    allow_flagging="never",
+)
+
+with demo:
+    gr.TabbedInterface(
+        [mf_transcribe, yt_transcribe], ["Audio Transcribe", "YouTube Transcribe"]
+    )
+
+demo.launch(enable_queue=True)