updated app.py

app.py

@@ -1,33 +1,43 @@
 import gradio as gr
+from transformers import pipeline, BertTokenizer, EncoderDecoderModel
 import torch
-from transformers import BertTokenizer, EncoderDecoderModel, pipeline
 
-# Load model and tokenizer
-model = EncoderDecoderModel.from_pretrained("imsachinsingh00/bert2bert-mts-summary")
-tokenizer = BertTokenizer.from_pretrained("imsachinsingh00/bert2bert-mts-summary")
+# Load Whisper for speech-to-text
+asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")
 
-# Move to CUDA if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# 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_dialogue(dialogue):
-    inputs = tokenizer(dialogue, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
-    summary_ids = model.generate(inputs.input_ids, max_length=64, num_beams=4, early_stopping=True)
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+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
 
-# Gradio interface
-demo = gr.Interface(
-    fn=summarize_dialogue,
-    inputs=[
-        gr.Textbox(lines=10, label="Doctor-Patient Dialogue"),
-        gr.Audio(source="microphone", type="filepath", optional=True)
-    ],
-    outputs="text",
-    title="Medical Dialogue Summarizer",
-    description="Enter or speak a conversation. The model will summarize it."
-)
-
-if __name__ == "__main__":
-    demo.launch()
+# 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()

requirements.txt

@@ -2,4 +2,5 @@ transformers
 datasets 
 evaluate
 rouge_score
-torch
+torch
+gradio