Update app.py

app.py

@@ -1,27 +1,15 @@
 import gradio as gr
-import yt_dlp
 import os
 import torch
 import gc
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import tempfile
-import whisper
-
-# Load summarizer
-@torch.no_grad()
-def load_summarizer():
-    model_name = "facebook/bart-large-cnn"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    return pipeline("summarization", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
-
-summarizer = load_summarizer()
+import yt_dlp
 
-# Load Whisper model
-whisper_model = whisper.load_model("base")  # or "small" for better accuracy
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-base")
 
 def download_audio(url: str, temp_dir: str) -> str:
-    """Download audio using yt-dlp and return path"""
     output_path = os.path.join(temp_dir, "audio.%(ext)s")
     ydl_opts = {
         'format': 'bestaudio/best',
@@ -37,31 +25,21 @@ def download_audio(url: str, temp_dir: str) -> str:
         ydl.download([url])
     return output_path.replace('%(ext)s', 'mp3')
 
-def transcribe_audio(audio_path: str) -> str:
-    """Transcribe audio with Whisper"""
-    result = whisper_model.transcribe(audio_path)
-    return result['text']
-
-def summarize_text(text: str) -> str:
-    """Summarize text"""
-    if len(text.strip()) < 50:
-        return "❌ Transcription too short to summarize"
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    summary = summarizer(text, max_length=150, min_length=50, do_sample=False)
-    return summary[0]['summary_text']
-
 def process_video(url: str) -> str:
     with tempfile.TemporaryDirectory() as tmpdir:
         audio_path = download_audio(url, tmpdir)
-        transcription = transcribe_audio(audio_path)
-        summary = summarize_text(transcription)
-        return summary
-
-def main(youtube_url):
-    return process_video(youtube_url)
+        transcription_result = asr_pipeline(audio_path)
+        text = transcription_result['text']
+        if len(text.strip()) < 50:
+            return "Transcription too short or unclear"
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        summary_result = summarizer(text, max_length=150, min_length=50, do_sample=False)
+        return summary_result[0]['summary_text']
+
+def main(url):
+    return process_video(url)
 
 iface = gr.Interface(fn=main, inputs="text", outputs="text", title="YouTube Audio Summarizer")
-
 iface.launch()