Update app.py

app.py

@@ -10,10 +10,7 @@ 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.tokenizer_model = AutoTokenizer.from_pretrained("google/pegasus-large")
-        self.summarizer_model = AutoModelForSeq2SeqLM.from_pretrained("google/pegasus-large")
+        self.summarizer = pipeline("summarization", model="google/pegasus-large")
                 
         # Initialize VoiceLabT5 model and tokenizer
         self.keyword_model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords")
@@ -35,16 +32,9 @@ class GradioInference():
             self.current_size = size
         
         results = self.loaded_model.transcribe(path, language=lang)
-
-        inputs = self.tokenizer_model(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: "
@@ -56,7 +46,7 @@ class GradioInference():
 
         label = self.classifier(results["text"])[0]["label"]
         
-        return results["text"], summary[0], keywords, label
+        return results["text"], transcription_summary[0]["summary_text"], keywords, label
 
     def populate_metadata(self, link):
         self.yt = YouTube(link)