import gradio as gr
from faster_whisper import WhisperModel
from transformers import pipeline

# Load Faster-Whisper and BART model
print("Loading models...")
whisper_model = WhisperModel("base")  # You can use "tiny" if it's still too heavy
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

def transcribe(audio_path):
    if audio_path is None:
        return "Please record some audio."
    
    segments, _ = whisper_model.transcribe(audio_path)
    transcription = " ".join([segment.text for segment in segments])
    return transcription

def summarize(text):
    if not text.strip():
        return "No transcription available to summarize."
    summary = summarizer(text, max_length=150, min_length=40, do_sample=False)
    return summary[0]['summary_text']

# Build Gradio UI
with gr.Blocks() as app:
    gr.Markdown("## 🎙️ Real-Time Transcription & Summarization\nSpeak into your mic and generate a summary.")

    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="🎧 Record Audio (Max ~120 sec)")
        transcription_output = gr.Textbox(label="📝 Transcription", lines=6, interactive=False)

    transcribe_button = gr.Button("Transcribe")
    transcribe_button.click(fn=transcribe, inputs=audio_input, outputs=transcription_output)

    with gr.Row():
        summarize_button = gr.Button("Generate Summary")
        summary_output = gr.Textbox(label="📄 Summary", lines=6, interactive=False)

    summarize_button.click(fn=summarize, inputs=transcription_output, outputs=summary_output)

app.launch()