Update app.py

app.py

@@ -1,24 +1,83 @@
 import gradio as gr
 import whisper
+from pytube import YouTube
+import yake
+from transformers import pipeline
 
+class GradioInference():
+    def __init__(self):
+        self.sizes = list(whisper._MODELS.keys())
+        self.langs = ["none"] + sorted(list(whisper.tokenizer.LANGUAGES.values()))
+        self.current_size = "base"
+        self.loaded_model = whisper.load_model(self.current_size)
+        self.yt = None
 
-def transcribe_audio(audio_file):
-    model = whisper.load_model("base")
-    result = model.transcribe(audio_file)
-    return result["text"]
+        # Initialize YAKE keyword extractor
+        self.keyword_extractor = yake.KeywordExtractor(lan="en", n=3, dedupLim=0.9, dedupFunc="seqm", windowsSize=1, top=5, features=None)
 
+        # Initialize Facebook/BART-Large-CNN summarizer
+        self.summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-def main():
-    audio_input = gr.Audio(source="upload", type="filepath")
-    output_text = gr.Textbox()
-    
-    iface = gr.Interface(fn=transcribe_audio, inputs=audio_input, 
-                         outputs=output_text, title="Audio Transcription App",
-                         description="Upload an audio file and hit the 'Submit'\
-                             button")
-    
-    iface.launch()
+    def __call__(self, link, lang, size):
+        if self.yt is None:
+            self.yt = YouTube(link)
+        path = self.yt.streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
 
+        if lang == "none":
+            lang = None
 
-if __name__ == '__main__':
-    main()
+        if size != self.current_size:
+            self.loaded_model = whisper.load_model(size)
+            self.current_size = size
+        results = self.loaded_model.transcribe(path, language=lang)
+
+        # Perform summarization on the transcription
+        transcription_summary = self.summarizer(results["text"], max_length=130, min_length=30, do_sample=False)
+
+        # Extract keywords from the transcription
+        keywords = self.keyword_extractor.extract_keywords(results["text"])
+
+        return results["text"], transcription_summary[0]["summary_text"], [kw[0] for kw in keywords]
+
+    def populate_metadata(self, link):
+        self.yt = YouTube(link)
+        return self.yt.thumbnail_url, self.yt.title
+
+
+gio = GradioInference()
+title = "Youtube Whisperer"
+description = "Speech to text transcription, summary, and keyword extraction of Youtube videos using OpenAI's Whisper, Facebook/BART-Large-CNN, and YAKE"
+
+block = gr.Blocks()
+with block:
+    gr.HTML(
+        """
+        <div style="text-align: center; max-width: 500px; margin: 0 auto;">
+          <div>
+            <h1>Youtube Whisperer</h1>
+          </div>
+          <p style="margin-bottom: 10px; font-size: 94%">
+            Speech to text transcription, summary, and keyword extraction of Youtube videos using OpenAI's Whisper, Facebook/BART-Large-CNN, and YAKE
+          </p>
+        </div>
+        """
+    )
+    with gr.Group():
+        with gr.Box():
+            with gr.Row().style(equal_height=True):
+                sz = gr.Dropdown(label="Model Size", choices=gio.sizes, value='base')
+                lang = gr.Dropdown(label="Language (Optional)", choices=gio.langs, value="none")
+            link = gr.Textbox(label="YouTube Link")
+            title = gr.Label(label="Video Title")
+            with gr.Row().style(equal_height=True):
+                img = gr.Image(label="Thumbnail")
+                text = gr.Textbox(label="Transcription", placeholder="Transcription Output", lines=10)
+            with gr.Row().style(equal_height=True):
+                summary = gr.Textbox(label="Summary", placeholder="Summary Output", lines=5)
+                keywords = gr.Textbox(label="Keywords", placeholder="Keywords Output", lines=5)
+            with gr.Row().style(equal_height=True):
+                btn = gr.Button("Transcribe, Summarize & Extract Keywords")
+            btn.click(gio, inputs=[link, lang, sz], outputs=[text, summary, keywords])
+            link.change(gio.populate_metadata, inputs=[link], outputs=[img, title])
+
+block.launch()