updated app.py

app.py

@@ -1,40 +1,58 @@
+import torch
 import gradio as gr
+import whisper
 from transformers import BertTokenizer, EncoderDecoderModel
-import torch
 
-# Load model and tokenizer
-model = EncoderDecoderModel.from_pretrained("Imsachinsingh00/bert2bert-mts-summary")
-tokenizer = BertTokenizer.from_pretrained("Imsachinsingh00/bert2bert-mts-summary")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+# Load Whisper model for transcription
+model_whisper = whisper.load_model("base")
 
-def summarize_text(dialogue):
-    inputs = tokenizer(dialogue, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    outputs = model.generate(**inputs, max_length=64)
-    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return summary
+# Load Summarization Model & Tokenizer
+model_name = "Imsachinsingh00/bert2bert-mts-summary"
+tokenizer = BertTokenizer.from_pretrained(model_name)
+model = EncoderDecoderModel.from_pretrained(model_name)
 
-# Gradio app
-with gr.Blocks() as demo:
-    gr.Markdown("## 🎤 Medical Dialogue Summarization App")
+# ✅ REQUIRED for encoder-decoder models
+model.config.decoder_start_token_id = tokenizer.cls_token_id
+model.config.pad_token_id = tokenizer.pad_token_id
 
-    with gr.Row():
-        mic_input = gr.Microphone(label="🎙️ Record Dialogue")
-        text_input = gr.Textbox(label="📝 Or Paste Dialogue", lines=10, placeholder="Paste or speak a conversation here...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
 
-    summarize_btn = gr.Button("Summarize")
-    summary_output = gr.Textbox(label="📄 Summary", interactive=False)
+# 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
 
-    def transcribe_and_summarize(audio, text):
-        import whisper
-        model_whisper = whisper.load_model("base")
-        result = model_whisper.transcribe(audio)
-        transcribed_text = result["text"]
-        text = text or transcribed_text
-        summary = summarize_text(text)
-        return transcribed_text, summary
+# Gradio Functions
+def transcribe_and_summarize(audio_file):
+    result = model_whisper.transcribe(audio_file)
+    transcription = result["text"]
+    summary = generate_summary(transcription)
+    return transcription, summary
 
-    summarize_btn.click(fn=transcribe_and_summarize, inputs=[mic_input, text_input], outputs=[text_input, summary_output])
+def summarize_text_input(text_input):
+    summary = generate_summary(text_input)
+    return text_input, summary
 
-demo.launch()
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🏥 Medical Dialogue Summarizer")
+
+    with gr.Tab("🎙️ Record & Summarize"):
+        audio_input = gr.Audio(type="filepath", label="Record Doctor-Patient Conversation")
+        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_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)

requirements.txt

@@ -6,4 +6,5 @@ torch
 gradio
 scikit-learn
 huggingface_hub
-whisper
+whisper
+git+https://github.com/openai/whisper.git