Update app.py

app.py

@@ -28,7 +28,11 @@ class GradioInference:
         self.yt = None
 
         # Initialize summary model for English
-        self.summarizer = pipeline("summarization", model="facebook/bart-large-cnn", truncation=True)
+        self.bart_summarizer = pipeline("summarization", model="facebook/bart-large-cnn", truncation=True)
+
+        # Initialize Multilingual summary model 
+        self.mt5_tokenizer = AutoTokenizer.from_pretrained("csebuetnlp/mT5_multilingual_XLSum", truncation=True)
+        self.mt5_model = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
 
         # Initialize VoiceLabT5 model and tokenizer
         self.keyword_model = T5ForConditionalGeneration.from_pretrained(
@@ -41,9 +45,6 @@ class GradioInference:
         # Sentiment Classifier
         self.classifier = pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student", return_all_scores=False)
 
-        # Initialize Multilingual summary model 
-        self.tokenizer = AutoTokenizer.from_pretrained("csebuetnlp/mT5_multilingual_XLSum", truncation=True)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
     
     def __call__(self, link, lang, size, progress=gr.Progress()):
         """
@@ -57,6 +58,7 @@ class GradioInference:
             - WordCloud: using the wordcloud python library.
         """
         progress(0, desc="Starting analysis")
+        
         if self.yt is None:
             self.yt = YouTube(link)
         
@@ -78,14 +80,18 @@ class GradioInference:
         progress(0.40, desc="Summarizing")
         
         # Perform summarization on the transcription
-        transcription_summary = self.summarizer(
-            results["text"], max_length=256, min_length=30, do_sample=False, truncation=True
+        transcription_summary = self.bart_summarizer(
+            results["text"], 
+            max_length=256, 
+            min_length=30, 
+            do_sample=False, 
+            truncation=True
         )
 
-        #### Resumen multilingue
+        #### Resumen multilingue con mt5 
         WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
         
-        input_ids_sum = self.tokenizer(
+        input_ids_sum = self.mt5_tokenizer(
             [WHITESPACE_HANDLER(results["text"])],
             return_tensors="pt",
             padding="max_length",
@@ -93,14 +99,14 @@ class GradioInference:
             max_length=512
         )["input_ids"]
         
-        output_ids_sum = self.model.generate(
+        output_ids_sum = self.mt5_model.generate(
             input_ids=input_ids_sum,
-            max_length=130,
+            max_length=256,
             no_repeat_ngram_size=2,
             num_beams=4
         )[0]
         
-        summary = self.tokenizer.decode(
+        summary = self.mt5_tokenizer.decode(
             output_ids_sum,
             skip_special_tokens=True,
             clean_up_tokenization_spaces=False
@@ -112,12 +118,19 @@ class GradioInference:
         # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
         input_sequence = task_prefix + results["text"]
+        
         input_ids = self.keyword_tokenizer(
-            input_sequence, return_tensors="pt", truncation=False
+            input_sequence, 
+            return_tensors="pt", 
+            truncation=False
         ).input_ids
+        
         output = self.keyword_model.generate(
-            input_ids, no_repeat_ngram_size=3, num_beams=4
+            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()]
         formatted_keywords = "\n".join([f"• {keyword}" for keyword in keywords])
@@ -201,14 +214,14 @@ class GradioInference:
         progress(0.40, desc="Summarizing")
         
         # Perform summarization on the transcription
-        transcription_summary = self.summarizer(
+        transcription_summary = self.bart_summarizer(
             results["text"], max_length=150, min_length=30, do_sample=False, truncation=True
         )
         
         #### Resumen multilingue
         WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
         
-        input_ids_sum = self.tokenizer(
+        input_ids_sum = self.mt5_tokenizer(
             [WHITESPACE_HANDLER(results["text"])],
             return_tensors="pt",
             padding="max_length",
@@ -216,14 +229,14 @@ class GradioInference:
             max_length=512
         )["input_ids"]
         
-        output_ids_sum = self.model.generate(
+        output_ids_sum = self.mt5_model.generate(
             input_ids=input_ids_sum,
             max_length=130,
             no_repeat_ngram_size=2,
             num_beams=4
         )[0]
         
-        summary = self.tokenizer.decode(
+        summary = self.mt5_tokenizer.decode(
             output_ids_sum,
             skip_special_tokens=True,
             clean_up_tokenization_spaces=False
@@ -235,11 +248,17 @@ class GradioInference:
         # Extract keywords using VoiceLabT5
         task_prefix = "Keywords: "
         input_sequence = task_prefix + results["text"]
+        
         input_ids = self.keyword_tokenizer(
-            input_sequence, return_tensors="pt", truncation=False
+            input_sequence, 
+            return_tensors="pt", 
+            truncation=False
         ).input_ids
+        
         output = self.keyword_model.generate(
-            input_ids, no_repeat_ngram_size=3, num_beams=4
+            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()]
@@ -267,10 +286,9 @@ class GradioInference:
         )
         wordcloud_image = wordcloud.to_image()
 
-        if lang == "english":
+        if lang == "english" or lang == "none":
             return (
                 results["text"],
-                # summ,
                 transcription_summary[0]["summary_text"],
                 formatted_keywords,
                 formatted_sentiment,
@@ -279,7 +297,6 @@ class GradioInference:
         else:
             return (
                 results["text"],
-                # summ,
                 summary,
                 formatted_keywords,
                 formatted_sentiment,
@@ -306,7 +323,7 @@ with block as demo:
         </div>
         """
     )
-    with gr.Group():
+    with gr.Group(spacing_size="md", radius_size="md"):
         with gr.Tab("From YouTube 📹"):
             with gr.Box():