Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import whisper
 from pytube import YouTube
-from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
+from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
 
 class GradioInference():
     def __init__(self):
@@ -10,8 +10,11 @@ class GradioInference():
         self.current_size = "base"
         self.loaded_model = whisper.load_model(self.current_size)
         self.yt = None
-        self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
-        
+        # self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+
+        self.tokenizer_model = AutoTokenizer.from_pretrained("google/pegasus-large")
+        self.summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("google/pegasus-large")
+                
         # Initialize VoiceLabT5 model and tokenizer
         self.keyword_model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords")
         self.keyword_tokenizer = T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords")
@@ -33,8 +36,15 @@ class GradioInference():
         
         results = self.loaded_model.transcribe(path, language=lang)
 
+        inputs = tokenizer(results["text"], max_length=1024, truncation=True, return_tensors="pt")
+    
+        summary_ids = self.keyword_model.generate(inputs["input_ids"])
+        summary = self.keyword_tokenizer.batch_decode(summary_ids, 
+                                  skip_special_tokens=True, 
+                                  clean_up_tokenization_spaces=False)
+        
         # Perform summarization on the transcription
-        transcription_summary = self.summarizer(results["text"], max_length=130, min_length=30, do_sample=False)
+        # transcription_summary = self.summarizer(results["text"], max_length=130, min_length=30, do_sample=False)
 
         # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
@@ -46,7 +56,7 @@ class GradioInference():
 
         label = self.classifier(results["text"])[0]["label"]
         
-        return results["text"], transcription_summary[0]["summary_text"], keywords, label
+        return results["text"], summary[0], keywords, label
 
     def populate_metadata(self, link):
         self.yt = YouTube(link)