import gradio as gr
from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline

# Load ClinicalBERT model
model_name = "emilyalsentzer/Bio_ClinicalBERT"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

# Create a text generation pipeline
nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)

# Function to interact with ClinicalBERT
def medical_chatbot(user_input):
    response = nlp_pipeline(user_input.replace("[MASK]", ""))
    return response[0]["sequence"]  # Returns the most likely sentence

# Gradio UI
interface = gr.Interface(
    fn=medical_chatbot,
    inputs=gr.Textbox(lines=2, placeholder="Enter medical query with [MASK]..."),
    outputs="text",
    title="Medical Chatbot",
    description="Ask medical questions. Example: 'Patient shows symptoms of [MASK]'."
)

interface.launch()