Update app.py

app.py

@@ -1,58 +1,48 @@
 import torch
 import gradio as gr
 import whisper
-from transformers import BertTokenizer, EncoderDecoderModel
+from transformers import pipeline
 
-# Load Whisper model for transcription
+# 🔹 Load Whisper model (base) for transcription
 model_whisper = whisper.load_model("base")
 
-# Load Summarization Model & Tokenizer
-model_name = "Imsachinsingh00/bert2bert-mts-summary"
-tokenizer = BertTokenizer.from_pretrained(model_name)
-model = EncoderDecoderModel.from_pretrained(model_name)
+# 🔹 Load summarization pipeline using BART
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-# ✅ REQUIRED for encoder-decoder models
-model.config.decoder_start_token_id = tokenizer.cls_token_id
-model.config.pad_token_id = tokenizer.pad_token_id
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
-
-# Summarization Function
-def generate_summary(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    summary_ids = model.generate(**inputs, max_length=64, num_beams=4, early_stopping=True)
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
-
-# Gradio Functions
+# 🔹 Transcription + Summarization
 def transcribe_and_summarize(audio_file):
-    result = model_whisper.transcribe(audio_file)
-    transcription = result["text"]
-    summary = generate_summary(transcription)
-    return transcription, summary
-
+    try:
+        result = model_whisper.transcribe(audio_file)
+        transcription = result.get("text", "")
+        if not transcription.strip():
+            return "No speech detected", "Summary not generated"
+        summary = summarizer(transcription, max_length=100, min_length=30, do_sample=False)
+        return transcription, summary[0]['summary_text']
+    except Exception as e:
+        return "Transcription failed", f"Error: {str(e)}"
+
+# 🔹 Text-only summarization
 def summarize_text_input(text_input):
-    summary = generate_summary(text_input)
-    return text_input, summary
+    if not text_input.strip():
+        return "", "Input is empty"
+    summary = summarizer(text_input, max_length=100, min_length=30, do_sample=False)
+    return text_input, summary[0]['summary_text']
 
-# Gradio UI
+# 🔹 Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏥 Medical Dialogue Summarizer")
+    gr.Markdown("## 🏥 Medical Summarization App (Whisper + BART)")
 
     with gr.Tab("🎙️ Record & Summarize"):
-        audio_input = gr.Audio(type="filepath", label="Record Doctor-Patient Conversation")
+        audio_input = gr.Audio(type="filepath", label="Upload or Record Audio")
         mic_transcript = gr.Textbox(label="Transcript")
         mic_summary = gr.Textbox(label="Summary", interactive=False)
         mic_button = gr.Button("Transcribe & Summarize")
         mic_button.click(transcribe_and_summarize, inputs=[audio_input], outputs=[mic_transcript, mic_summary])
 
     with gr.Tab("📋 Paste & Summarize"):
-        text_input = gr.Textbox(lines=8, label="Paste Dialogue Here")
+        text_input = gr.Textbox(lines=8, label="Paste Medical Report or Dialogue")
         text_output = gr.Textbox(label="Summary", interactive=False)
         text_button = gr.Button("Summarize")
         text_button.click(summarize_text_input, inputs=[text_input], outputs=[text_input, text_output])
 
-# Launch with sharing for local + link
 demo.launch(share=True)