Update app.py

app.py

@@ -1,101 +1,78 @@
-import gradio as gr
-import whisper
 import os
-import re
+import gradio as gr
 import yt_dlp
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import whisper
+from transformers import pipeline, MarianMTModel, MarianTokenizer
+from pytube import YouTube
+import torch
 
-# Load models
-whisper_model = whisper.load_model("base")
-multilingual_model = "csebuetnlp/mT5_multilingual_XLSum"
-tokenizer = AutoTokenizer.from_pretrained(multilingual_model)
-summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(multilingual_model)
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
+# Load Whisper (tiny for CPU)
+whisper_model = whisper.load_model("small")
 
-SUPPORTED_LANGUAGES = {
-    "bn": "Bengali", "en": "English", "gu": "Gujarati", "hi": "Hindi",
-    "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi", "ta": "Tamil",
-    "te": "Telugu", "ur": "Urdu"
-}
+# Load summarizer
+summarizer = pipeline("summarization", model="Falconsai/text_summarization")
+
+# Load translation model (multilingual to English)
+translation_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-mul-en")
+translation_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-mul-en")
 
 def download_audio(youtube_url):
+    output_file = "audio.webm"
     ydl_opts = {
         'format': 'bestaudio/best',
-        'outtmpl': 'audio.%(ext)s',
-        'postprocessors': [{
-            'key': 'FFmpegExtractAudio',
-            'preferredcodec': 'mp3',
-        }],
-        'cookiefile': 'cookies.txt'  # <— Uses your uploaded cookie file
+        'outtmpl': output_file,
+        'quiet': True,
     }
-
     with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        info_dict = ydl.extract_info(youtube_url, download=True)
-        filename = ydl.prepare_filename(info_dict).replace(".webm", ".mp3").replace(".m4a", ".mp3")
-    return filename
-
-def extract_thumbnail(youtube_url):
-    match = re.search(r"(?:v=|\/)([0-9A-Za-z_-]{11})", youtube_url)
-    if match:
-        video_id = match.group(1)
-        return f"https://img.youtube.com/vi/{video_id}/0.jpg"
-    return ""
+        ydl.download([youtube_url])
+    return output_file
 
-def summarize_text(text):
-    input_text = f"summarize: {text}"
-    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
-    summary_ids = summarizer_model.generate(inputs, max_length=150, min_length=30, num_beams=4)
-    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+def translate_to_english(text):
+    inputs = translation_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    translated = translation_model.generate(**inputs, max_length=512)
+    return translation_tokenizer.decode(translated[0], skip_special_tokens=True)
 
-def transcribe_and_summarize(youtube_url, translate_to_english):
-    try:
-        audio_file = download_audio(youtube_url)
-        result = whisper_model.transcribe(audio_file)
-        transcript = result["text"]
-        lang_code = result["language"]
+def process_video(url):
+    audio_path = download_audio(url)
+    result = whisper_model.transcribe(audio_path)
+    transcription = result["text"]
 
-        thumbnail_url = extract_thumbnail(youtube_url)
+    translated_text = translate_to_english(transcription)
 
-        if lang_code not in SUPPORTED_LANGUAGES:
-            return None, f"❌ Language '{lang_code}' not supported.", "", "", None
+    # Summarize
+    summary = summarizer(translated_text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]
 
-        summary = summarize_text(transcript)
+    # Get thumbnail
+    yt = YouTube(url)
+    thumbnail_url = yt.thumbnail_url
 
-        if translate_to_english and lang_code != "en":
-            translated_summary = translator(summary)[0]["translation_text"]
-        else:
-            translated_summary = summary
+    return transcription, translated_text, summary, thumbnail_url, summary
 
-        os.remove(audio_file)
-
-        summary_text = f"Transcript:\n{transcript}\n\nSummary:\n{translated_summary}"
-
-        with open("summary.txt", "w", encoding="utf-8") as f:
-            f.write(summary_text)
-
-        return thumbnail_url, f"🗣️ Language: {SUPPORTED_LANGUAGES[lang_code]}", transcript, translated_summary, "summary.txt"
-
-    except Exception as e:
-        return None, f"❌ Error: {str(e)}", "", "", None
+def download_summary(text):
+    filename = "summary.txt"
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(text)
+    return filename
 
-with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("<h1 style='text-align: center;'>🎬 Multilingual YouTube Summarizer</h1>")
-    gr.Markdown("Paste any YouTube video link, and get transcript + summary. Works for Hindi, Bengali, Tamil, Urdu, and more!")
+with gr.Blocks(theme=gr.themes.Soft(), title="🎥 YouTube Video Summarizer with LLaMA") as demo:
+    gr.Markdown("## 🧠 Multilingual YouTube Summarizer")
+    gr.Markdown("Upload a video link and get the transcript, English translation, and summary.")
 
     with gr.Row():
-        youtube_url = gr.Textbox(label="YouTube Video URL")
-        translate_check = gr.Checkbox(label="Translate Summary to English", value=True)
-
-    thumbnail = gr.Image(label="Video Thumbnail", type="filepath")
-    lang_out = gr.Text(label="Detected Language")
-    transcript_out = gr.Textbox(label="Transcript", lines=8)
-    summary_out = gr.Textbox(label="Summary", lines=6)
-    download_btn = gr.File(label="Download .txt")
+        youtube_input = gr.Text(label="YouTube Video URL", placeholder="https://www.youtube.com/watch?v=...")
+        submit_btn = gr.Button("Transcribe & Summarize")
 
-    btn = gr.Button("Generate Summary")
-
-    btn.click(fn=transcribe_and_summarize,
-              inputs=[youtube_url, translate_check],
-              outputs=[thumbnail, lang_out, transcript_out, summary_out, download_btn])
+    with gr.Row():
+        with gr.Column():
+            transcript_output = gr.Textbox(label="🔊 Original Transcript", lines=10)
+            translation_output = gr.Textbox(label="🌍 Translated to English", lines=10)
+            summary_output = gr.Textbox(label="🧾 Summary", lines=10)
+            download_btn = gr.Button("📥 Download Summary")
+            download_file = gr.File(label="Download Link")
+        video_thumb = gr.Image(label="🎞️ Video Thumbnail", width=256)
+
+    submit_btn.click(fn=process_video, inputs=[youtube_input],
+                     outputs=[transcript_output, translation_output, summary_output, video_thumb, download_file])
+    download_btn.click(fn=download_summary, inputs=[summary_output], outputs=[download_file])
 
 demo.launch()