Update app.py

app.py

@@ -57,125 +57,131 @@ class GradioInference:
             - Sentiment Analysis: using Hugging Face's default sentiment classifier
             - WordCloud: using the wordcloud python library.
         """
-
-        gr.Info("Starting process")
-        progress(0, desc="Starting analysis")
-        
-        if self.yt is None:
-            self.yt = YouTube(link)
-        
-        # Pytube library to access to YouTube audio stream
-        path = self.yt.streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
-
-        if lang == "none":
-            lang = None
-
-        if size != self.current_size:
-            self.loaded_model = whisper.load_model(size)
-            self.current_size = size
-
-        progress(0.20, desc="Transcribing")
-        
-        # Transcribe the audio extracted from pytube
-        results = self.loaded_model.transcribe(path, language=lang)
-
-        progress(0.40, desc="Summarizing")
-        
-        # Perform summarization on the transcription
-        transcription_summary = self.bart_summarizer(
-            results["text"], 
-            max_length=256, 
-            min_length=30, 
-            do_sample=False, 
-            truncation=True
-        )
-
-        # Multilingual summary with mt5 
-        WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
-        
-        input_ids_sum = self.mt5_tokenizer(
-            [WHITESPACE_HANDLER(results["text"])],
-            return_tensors="pt",
-            padding="max_length",
-            truncation=True,
-            max_length=512
-        )["input_ids"]
-        
-        output_ids_sum = self.mt5_model.generate(
-            input_ids=input_ids_sum,
-            max_length=256,
-            no_repeat_ngram_size=2,
-            num_beams=4
-        )[0]
-        
-        summary = self.mt5_tokenizer.decode(
-            output_ids_sum,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False
-        )
-        # End multilingual summary
-        
-        progress(0.60, desc="Extracting Keywords")
-        
-        # 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_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()]
-        formatted_keywords = "\n".join([f"• {keyword}" for keyword in keywords])
-
-        progress(0.80, desc="Extracting Sentiment")
-       
-        # Define a dictionary to map labels to emojis
-        sentiment_emojis = {
-            "positive": "Positive 👍🏼",
-            "negative": "Negative 👎🏼",
-            "neutral": "Neutral 😶",
-        }
-        
-        # Sentiment label    
-        label = self.classifier(summary)[0]["label"]
-
-        # Format the label with emojis
-        formatted_sentiment = sentiment_emojis.get(label, label)
-
-        progress(0.90, desc="Generating Wordcloud")
-        
-        # Generate WordCloud object
-        wordcloud = WordCloud(colormap = "Oranges").generate(results["text"])
-
-        # WordCloud image to display
-        wordcloud_image = wordcloud.to_image()
-
-        if lang == "english" or lang == "none":
-            return (
-                results["text"],
-                transcription_summary[0]["summary_text"],
-                formatted_keywords,
-                formatted_sentiment,
-                wordcloud_image,
+        try:
+            
+            progress(0, desc="Starting analysis")
+            
+            if self.yt is None:
+                self.yt = YouTube(link)
+            
+            # Pytube library to access to YouTube audio stream
+            path = self.yt.streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
+    
+            if lang == "none":
+                lang = None
+    
+            if size != self.current_size:
+                self.loaded_model = whisper.load_model(size)
+                self.current_size = size
+    
+            progress(0.20, desc="Transcribing")
+            
+            # Transcribe the audio extracted from pytube
+            results = self.loaded_model.transcribe(path, language=lang)
+    
+            progress(0.40, desc="Summarizing")
+            
+            # Perform summarization on the transcription
+            transcription_summary = self.bart_summarizer(
+                results["text"], 
+                max_length=256, 
+                min_length=30, 
+                do_sample=False, 
+                truncation=True
             )
-        else:
-            return (
-                results["text"],
-                summary,
-                formatted_keywords,
-                formatted_sentiment,
-                wordcloud_image,
+    
+            # Multilingual summary with mt5 
+            WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
+            
+            input_ids_sum = self.mt5_tokenizer(
+                [WHITESPACE_HANDLER(results["text"])],
+                return_tensors="pt",
+                padding="max_length",
+                truncation=True,
+                max_length=512
+            )["input_ids"]
+            
+            output_ids_sum = self.mt5_model.generate(
+                input_ids=input_ids_sum,
+                max_length=256,
+                no_repeat_ngram_size=2,
+                num_beams=4
+            )[0]
+            
+            summary = self.mt5_tokenizer.decode(
+                output_ids_sum,
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False
+            )
+            # End multilingual summary
+            
+            progress(0.60, desc="Extracting Keywords")
+            
+            # 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_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()]
+            formatted_keywords = "\n".join([f"• {keyword}" for keyword in keywords])
+    
+            progress(0.80, desc="Extracting Sentiment")
+           
+            # Define a dictionary to map labels to emojis
+            sentiment_emojis = {
+                "positive": "Positive 👍🏼",
+                "negative": "Negative 👎🏼",
+                "neutral": "Neutral 😶",
+            }
+            
+            # Sentiment label    
+            label = self.classifier(summary)[0]["label"]
+    
+            # Format the label with emojis
+            formatted_sentiment = sentiment_emojis.get(label, label)
+    
+            progress(0.90, desc="Generating Wordcloud")
+            
+            # Generate WordCloud object
+            wordcloud = WordCloud(colormap = "Oranges").generate(results["text"])
+    
+            # WordCloud image to display
+            wordcloud_image = wordcloud.to_image()
+    
+            if lang == "english" or lang == "none":
+                return (
+                    results["text"],
+                    transcription_summary[0]["summary_text"],
+                    formatted_keywords,
+                    formatted_sentiment,
+                    wordcloud_image,
+                )
+            else:
+                return (
+                    results["text"],
+                    summary,
+                    formatted_keywords,
+                    formatted_sentiment,
+                    wordcloud_image,
+                )
+
+        except:
+            gr.Error("Restricted Content. Choose a different video")
+            return None, None, None, None, None
+
+        gr.Info("Success")
 
 
     def populate_metadata(self, link):