Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import whisper
 from pytube import YouTube
-import yake
 from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
 
 class GradioInference():
@@ -11,13 +10,12 @@ class GradioInference():
         self.current_size = "base"
         self.loaded_model = whisper.load_model(self.current_size)
         self.yt = None
-
-        # Initialize Facebook/BART-Large-CNN summarizer
         self.summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-
+        
+        # Initialize VoiceLabT5 model and tokenizer
         self.keyword_model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords")
-        self.tokenizer = T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords")
-    
+        self.keyword_tokenizer = T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords")
+
     def __call__(self, link, lang, size):
         if self.yt is None:
             self.yt = YouTube(link)
@@ -34,57 +32,78 @@ class GradioInference():
         # Perform summarization on the transcription
         transcription_summary = self.summarizer(results["text"], max_length=130, min_length=30, do_sample=False)
 
+        # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
-
-        input_sequence = task_prefix + transcription
-        input_ids = tokenizer(
-            input_sequence, return_tensors="pt", truncation=False,
-        ).input_ids
-        output = keyword_model.generate(input_ids, no_repeat_ngram_size=3, num_beams=4)
-        predicted = tokenizer.decode(output[0], skip_special_tokens=True)
+        input_sequence = task_prefix + results["text"]
+        input_ids = self.keyword_tokenizer(input_sequence, return_tensors="pt", truncation=False).input_ids
+        output = self.keyword_model.generate(input_ids, no_repeat_ngram_size=3, num_beams=4)
+        predicted = self.keyword_tokenizer.decode(output[0], skip_special_tokens=True)
         keywords = [x.strip() for x in predicted.split(',') if x.strip()]
-
+        
         return results["text"], transcription_summary[0]["summary_text"], keywords
 
     def populate_metadata(self, link):
         self.yt = YouTube(link)
         return self.yt.thumbnail_url, self.yt.title
 
+def transcribe_audio(audio_file):
+    model = whisper.load_model("base")
+    result = model.transcribe(audio_file)
+    return result["text"]
+
 
 gio = GradioInference()
 title = "Youtube Insights"
-description = "Your AI-powered Video Analytics"
+description = "Your AI-powered video analytics tool"
 
 block = gr.Blocks()
-with block:
+with block as demo:
     gr.HTML(
         """
         <div style="text-align: center; max-width: 500px; margin: 0 auto;">
           <div>
-            <h1>Youtube Insights</h1>
+            <h1>Youtube <span style="color: red;">Insights</span> 📹</h1>
           </div>
           <p style="margin-bottom: 10px; font-size: 94%">
-            Your AI-powered Video Analytics
+            Your AI-powered video analytics tool
           </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 URL")
-            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("Get Video Insights")
-            btn.click(gio, inputs=[link, lang, sz], outputs=[text, summary, keywords])
-            link.change(gio.populate_metadata, inputs=[link], outputs=[img, title])
+        with gr.Tab("From YouTube"):
+            with gr.Box():
+                with gr.Row().style(equal_height=True):
+                    size = 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("Get video insights")  # Updated button label
+                btn.click(gio, inputs=[link, lang, size], outputs=[text, summary, keywords])
+                link.change(gio.populate_metadata, inputs=[link], outputs=[img, title])
+
+        with gr.Tab("From Audio file"):
+            with gr.Box():
+                with gr.Row().style(equal_height=True):
+                    size = gr.Dropdown(label="Model Size", choices=gio.sizes, value='base')
+                    lang = gr.Dropdown(label="Language (Optional)", choices=gio.langs, value="none")
+                audio_file = gr.Audio(type="filepath")
+                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("Get video insights")  # Updated button label
+                btn.click(transcribe_audio, inputs=[audio_file], outputs=[text])
+                # link.change(gio.populate_metadata, inputs=[link], outputs=[img, title])
 
-block.launch()
+demo.launch()