Update app.py

app.py

@@ -1,26 +1,65 @@
 import gradio as gr
 import whisper
+from transformers import pipeline
 
-# Function to process the audio file
+# Load Whisper model
+whisper_model = whisper.load_model("base")
+
+# Load traditional summarization models
+def get_summarizer(model_name):
+    if model_name == "BART (facebook/bart-large-cnn)":
+        return pipeline("summarization", model="facebook/bart-large-cnn")
+    elif model_name == "T5 (t5-small)":
+        return pipeline("summarization", model="t5-small")
+    elif model_name == "Pegasus (google/pegasus-xsum)":
+        return pipeline("summarization", model="google/pegasus-xsum")
+    else:
+        return None
+
+# Function to transcribe audio file using Whisper
 def transcribe_audio(model_size, audio):
-    # Load the Whisper model based on the user's choice
     model = whisper.load_model(model_size)
-    
-    # Transcribe the audio file
     result = model.transcribe(audio)
+    transcription = result['text']
+    return transcription
+
+# Function to summarize the transcribed text
+def summarize_text(transcription, model_name):
+    if len(transcription.strip()) == 0:
+        return "No text to summarize."
+    
+    summarizer = get_summarizer(model_name)
     
-    return result['text']
+    if summarizer:
+        summary = summarizer(transcription, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
+        return summary
+    else:
+        return "Invalid summarization model selected."
 
-# Gradio interface with model selection
+# Create a Gradio interface that combines transcription and summarization
+def combined_transcription_and_summarization(model_size, summarizer_model, audio):
+    # Step 1: Transcribe the audio using Whisper
+    transcription = transcribe_audio(model_size, audio)
+    
+    # Step 2: Summarize the transcribed text using the chosen summarizer model
+    summary = summarize_text(transcription, summarizer_model)
+    
+    return transcription, summary
+
+# Gradio interface for transcription and summarization
 iface = gr.Interface(
-    fn=transcribe_audio,                           # The function that will process the audio and model choice
+    fn=combined_transcription_and_summarization,   # The combined function
     inputs=[
-        gr.Dropdown(label="Choose Whisper Model", choices=["tiny", "base", "small", "medium", "large"], value="base"),  # Model selection
-        gr.Audio(type="filepath")  # Audio upload input (removed 'source')
+        gr.Dropdown(label="Choose Whisper Model", choices=["tiny", "base", "small", "medium", "large"], value="base"),  # Whisper model selection
+        gr.Dropdown(label="Choose Summarizer Model", choices=["BART (facebook/bart-large-cnn)", "T5 (t5-small)", "Pegasus (google/pegasus-xsum)"], value="BART (facebook/bart-large-cnn)"),  # Summarizer model selection
+        gr.Audio(type="filepath")  # Audio upload
+    ],
+    outputs=[
+        gr.Textbox(label="Transcription"),         # Output for the transcribed text
+        gr.Textbox(label="Summary")                # Output for the summary
     ],
-    outputs="text",                                # Output transcription as text
-    title="Whisper Audio Transcription",
-    description="Upload an audio file and select a Whisper model to get the transcription."
+    title="Whisper Audio Transcription and Summarization",
+    description="Upload an audio file, choose a Whisper model for transcription, and a summarization model to summarize the transcription."
 )
 
 # Launch the interface