Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, BartForConditionalGeneration, BartTokenizer
 from huggingface_hub import login
 import os
 
@@ -18,8 +18,9 @@ whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
 whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
 
 # Initialize the summarization model and tokenizer
-summarization_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
-summarization_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+# Use BART model for summarization
+summarization_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
+summarization_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
 
 # Function to transcribe audio
 def transcribe_audio(audio_file):
@@ -32,7 +33,7 @@ def transcribe_audio(audio_file):
 
 # Function to summarize text
 def summarize_text(text):
-    inputs = summarization_tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+    inputs = summarization_tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
     summary_ids = summarization_model.generate(inputs.input_ids, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
     summary = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
@@ -46,7 +47,7 @@ def process_audio(audio_file):
 # Gradio UI
 iface = gr.Interface(
     fn=process_audio,
-    inputs=gr.Audio(source="upload", type="file"),
+    inputs=gr.Audio(type="file"),
     outputs=[
         gr.Textbox(label="Transcription"),
         gr.Textbox(label="Summary")