Update app.py

app.py

@@ -4,9 +4,20 @@ from youtube_transcript_api import YouTubeTranscriptApi
 from youtube_transcript_api.formatters import TextFormatter
 import re
 
-# Load the summarization model
-text_summary = pipeline("summarization", model="Falconsai/text_summarization")
+# Define the models
+models = {
+    "Falconsai/text_summarization": "Falconsai/text_summarization",
+    "suriya7/bart-finetuned-text-summarization": "suriya7/bart-finetuned-text-summarization"
+}
 
+# Default model
+default_model = "Falconsai/text_summarization"
+
+# Function to create a summarization pipeline
+def create_summarization_pipeline(model_name):
+    return pipeline("summarization", model=model_name)
+
+# Function to extract video ID from URL
 def extract_video_id(url):
     regex = r"(?:youtube\.com\/(?:[^\/\n\s]+\/\S+\/|(?:v|e(?:mbed)?)\/|\S*?[?&]v=)|youtu\.be\/)([a-zA-Z0-9_-]{11})"
     match = re.search(regex, url)
@@ -14,6 +25,7 @@ def extract_video_id(url):
         return match.group(1)
     return None
 
+# Function to get YouTube transcript
 def get_youtube_transcript(video_url):
     video_id = extract_video_id(video_url)
     if not video_id:
@@ -26,20 +38,25 @@ def get_youtube_transcript(video_url):
     except Exception as e:
         return f"An error occurred: {e}"
 
-def summarize_youtube_video(url):
+# Function to summarize YouTube video with selected model
+def summarize_youtube_video(url, model_name):
     transcript = get_youtube_transcript(url)
     if "An error occurred" in transcript:
         return transcript
-    summary = text_summary(transcript, min_length=10, max_length=1000, do_sample=False)
+    summarization_pipeline = create_summarization_pipeline(model_name)
+    summary = summarization_pipeline(transcript, min_length=10, max_length=1000, do_sample=False)
     return summary[0]['summary_text']
 
 # Define the Gradio interface
 iface = gr.Interface(
     fn=summarize_youtube_video,
-    inputs=gr.Textbox(label="Enter YouTube Video URL", placeholder="e.g. https://www.youtube.com/watch?v=abcdef12345"),
+    inputs=[
+        gr.Textbox(label="Enter YouTube Video URL", placeholder="e.g. https://www.youtube.com/watch?v=abcdef12345"),
+        gr.Dropdown(choices=list(models.keys()), value=default_model, label="Select Summarization Model")
+    ],
     outputs=gr.Textbox(label="Video Summary"),
     title="YouTube Video Summarizer",
-    description="Enter the URL of a YouTube video to get a summary of its transcript."
+    description="Enter the URL of a YouTube video and select a summarization model to get a summary of its transcript."
 )
 
 if __name__ == "__main__":