import gradio as gr
from transformers import pipeline

# Load Whisper pipeline (speech-to-text)
asr_pipe = pipeline(
    task="automatic-speech-recognition",
    model="openai/whisper-small",
    return_timestamps=True
)

# Load summarization pipeline
summarizer = pipeline("summarization", model="google/pegasus-cnn_dailymail")

# Function to transcribe and summarize
def transcribe_and_summarize(audio_file):
    if audio_file is None:
        return "Please upload an audio file.", ""
    
    # Transcription
    transcript = asr_pipe(audio_file)["text"]

    # Summarization (limit to 1024 tokens for safety)
    if len(transcript) > 1000:
        transcript = transcript[:1000]  # truncate long transcripts
    summary = summarizer(transcript, max_length=80, min_length=20, do_sample=False)[0]["summary_text"]
    
    return transcript, summary

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## 🎙️ Whisper ASR + Summary\nUpload or record audio to get a transcription and summary.")

    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="Upload or Record Audio")
    
    with gr.Row():
        transcribe_btn = gr.Button("Transcribe & Summarize")
    
    with gr.Row():
        transcript_box = gr.Textbox(label="Transcription", lines=6)
        summary_box = gr.Textbox(label="Summary", lines=4)

    transcribe_btn.click(fn=transcribe_and_summarize, inputs=audio_input, outputs=[transcript_box, summary_box])

app.launch()