app.py

import gradio as gr
from transformers import pipeline, BertTokenizer, EncoderDecoderModel
import torch

# Load Whisper for speech-to-text
asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")

# Load your fine-tuned summarization model
model_name = "Imsachinsingh00/bert2bert-mts-summary"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = EncoderDecoderModel.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")

# Summarization function
def summarize_text(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
    outputs = model.generate(**inputs, max_length=64)
    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return summary

# Pipeline: audio → transcription
def transcribe(audio):
    return asr(audio)["text"]

# App UI
with gr.Blocks() as demo:
    gr.Markdown("## 🩺 Medical Dialogue Summarizer")

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(source="microphone", type="filepath", label="🎙️ Record Dialogue")
            transcribed_text = gr.Textbox(lines=10, label="📝 Transcribed Text (editable)")

            record_button = gr.Button("🎧 Transcribe")
            record_button.click(transcribe, inputs=audio_input, outputs=transcribed_text)

        with gr.Column():
            summary_output = gr.Textbox(lines=10, label="📋 Summary (output)", interactive=False)
            summarize_button = gr.Button("✂️ Summarize")
            summarize_button.click(summarize_text, inputs=transcribed_text, outputs=summary_output)

    gr.Markdown("Built for Voize Interview — Powered by Whisper + BERT")

demo.launch()