Update app.py

app.py

@@ -1,29 +1,16 @@
 import gradio as gr
-from transformers import AutoModelForMaskedLM, AutoTokenizer, pipeline
+from transformers import pipeline
 
-# Load ClinicalBERT model
 model_name = "emilyalsentzer/Bio_ClinicalBERT"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForMaskedLM.from_pretrained(model_name)
+nlp_pipeline = pipeline("fill-mask", model=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):
     if "[MASK]" not in user_input:
-        return "Please include '[MASK]' in your input. Example: 'The patient is showing signs of [MASK].'"
-
+        user_input += " [MASK]"
+    
     response = nlp_pipeline(user_input)
-    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]'."
-)
+    return {"prediction": response[0]["sequence"], "confidence": response[0]["score"]}
 
-interface.launch()
+# Create an API interface
+interface = gr.Interface(fn=medical_chatbot, inputs="text", outputs="json")
+interface.launch(share=True)