Update app.py

app.py

@@ -1,13 +1,21 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the pre-trained model (you can replace this with any healthcare model from Hugging Face)
-model = pipeline("fill-mask", model="medicalai/ClinicalBERT")
+# Load the ClinicalBERT model
+model = pipeline("text-classification", model="emilyalsentzer/Bio_ClinicalBERT")
 
 # Function to diagnose based on symptoms
 def diagnose(symptoms):
     prediction = model(symptoms)
-    return f"Possible diagnosis: {prediction[0]['label']}"
+    LABEL_TO_DIAGNOSIS = {
+        "LABEL_0": "Common Cold",
+        "LABEL_1": "Flu",
+        "LABEL_2": "Migraine",
+        "LABEL_3": "Pneumonia",
+    }
+    label = prediction[0]['label']
+    diagnosis = LABEL_TO_DIAGNOSIS.get(label, "Diagnosis not found")
+    return f"Possible diagnosis: {diagnosis}"
 
 # Triage function based on symptoms
 def triage(symptoms):