Update app.py

app.py

@@ -1,51 +1,41 @@
-import re
+import pytube
 from youtube_transcript_api import YouTubeTranscriptApi
-from youtube_transcript_api.formatters import TextFormatter
-import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import gradio as gr
-from transformers import pipeline
 
-text_summary = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)
+# Load the Hugging Face model and tokenizer
+model_name = "sshleifer/distilbart-cnn-12-6"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-def summary (input):
-    output = text_summary(input)
-    return output[0]['summary_text']
+def get_transcript(youtube_url):
+    # Extract the video ID from the YouTube URL
+    video_id = pytube.extract.video_id(youtube_url)
 
-def extract_video_id(url):
-    # Regex to extract the video ID from various YouTube URL formats
-    regex = r"(?:youtube\.com\/(?:[^\/\n\s]+\/\S+\/|(?:v|e(?:mbed)?)\/|\S*?[?&]v=)|youtu\.be\/)([a-zA-Z0-9_-]{11})"
-    match = re.search(regex, url)
-    if match:
-        return match.group(1)
-    return None
+    # Get the transcript using the YouTube Transcript API
+    try:
+        transcript_list = YouTubeTranscriptApi.get_transcript(video_id)
+    except Exception as e:
+        return f"Error retrieving transcript: {str(e)}"
 
+    # Join the transcript segments into a single string
+    transcript_text = " ".join([segment["text"] for segment in transcript_list])
 
-def get_youtube_transcript(video_url):
-    video_id = extract_video_id(video_url)
-    if not video_id:
-        return "Video ID could not be extracted."
+    # Summarize the transcript text using the Hugging Face model
+    inputs = tokenizer(transcript_text, return_tensors="pt", truncation=True, padding="longest")
+    summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=100, early_stopping=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-    try:
-        # Fetch the transcript
-        transcript = YouTubeTranscriptApi.get_transcript(video_id)
+    return summary
 
-        # Format the transcript into plain text
-        formatter = TextFormatter()
-        text_transcript = formatter.format_transcript(transcript)
-        summary_text = summary(text_transcript)
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=get_transcript,
+    inputs="text",
+    outputs="text",
+    title="@IT AI Enthusiast (Mayank Chugh) (https://www.youtube.com/@itaienthusiast/) - Project 2: YouTube Video Transcript Generator",
+    description="Enter a YouTube URL to generate and summarize the video transcript."
+)
 
-        return summary_text
-    except Exception as e:
-        return f"An error occurred: {e}"
-
-gr.close_all()
-
-# demo = gr.Interface(fn=summary, inputs="text",outputs="text")
-demo = gr.Interface(fn=get_youtube_transcript,
-                    inputs=[gr.Textbox(label="Input YouTube Url to summarize",lines=1)],
-                    outputs=[gr.Textbox(label="Summarized text",lines=4)],
-                    title="@IT AI Enthusiast (https://www.youtube.com/@itaienthusiast/) - Project 2: YouTube Script Summarizer",
-                    description="THIS APPLICATION WILL BE USED TO SUMMARIZE THE YOUTUBE VIDEO SCRIPT.",
-                    examples=['https://www.youtube.com/watch?v=tQb7bumjkIM'],
-                    concurrency_limit=8)
-demo.launch()
+# Launch the Gradio interface
+iface.launch()