import gradio as gr
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")

# Move to CUDA if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 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)
    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()