Update app.py

app.py

@@ -2,8 +2,9 @@
 import gradio as gr
 import whisper
 from pytube import YouTube
-from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
+from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
 from wordcloud import WordCloud
+import re
 
 class GradioInference:
     def __init__(self):
@@ -37,6 +38,12 @@ class GradioInference:
         # Sentiment Classifier
         self.classifier = pipeline("text-classification")
 
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
+
+
+    
     def __call__(self, link, lang, size):
         """
         Call the Gradio Inference python class.
@@ -69,6 +76,32 @@ class GradioInference:
             results["text"], max_length=512, min_length=30, do_sample=False
         )
 
+        #### Prueba
+        WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
+        
+        input_ids = self.tokenizer(
+            [WHITESPACE_HANDLER(results["text"])],
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=512
+        )["input_ids"]
+        
+        output_ids = self.model.generate(
+            input_ids=input_ids,
+            max_length=130,
+            no_repeat_ngram_size=2,
+            num_beams=4
+        )[0]
+        
+        summary = tokenizer.decode(
+            output_ids,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False
+        )
+        #### Fin prueba
+
+        
         # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
         input_sequence = task_prefix + results["text"]
@@ -92,7 +125,7 @@ class GradioInference:
 
         return (
             results["text"],
-            transcription_summary[0]["summary_text"],
+            summary,
             keywords,
             label,
             wordcloud_image,
@@ -131,6 +164,31 @@ class GradioInference:
             results["text"], max_length=512, min_length=30, do_sample=False
         )
 
+        #### Prueba
+        WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
+        
+        input_ids = self.tokenizer(
+            [WHITESPACE_HANDLER(results["text"])],
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=512
+        )["input_ids"]
+        
+        output_ids = self.model.generate(
+            input_ids=input_ids,
+            max_length=130,
+            no_repeat_ngram_size=2,
+            num_beams=4
+        )[0]
+        
+        summary = tokenizer.decode(
+            output_ids,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False
+        )
+        #### Fin prueba
+        
         # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
         input_sequence = task_prefix + results["text"]
@@ -154,7 +212,7 @@ class GradioInference:
 
         return (
             results["text"],
-            transcription_summary[0]["summary_text"],
+            summary,
             keywords,
             label,
             wordcloud_image,