Update app.py

app.py

@@ -1,39 +1,43 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the ClinicalBERT model
-model = pipeline("text-classification", model="emilyalsentzer/Bio_ClinicalBERT")
+# Load a general-purpose zero-shot classification model (can be BioBERT or another fine-tuned model)
+model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 # Function to diagnose based on symptoms
 def diagnose(symptoms):
-    prediction = model(symptoms)
-    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}"
+    # Using zero-shot classification, no predefined labels are needed; model dynamically predicts the best category
+    result = model(symptoms, candidate_labels=["disease", "illness", "symptom", "medical condition"])
+    
+    # Dynamic prediction
+    diagnosis = result['labels'][0]  # Top predicted label
+    confidence = result['scores'][0]  # Confidence score
+    
+    return f"Predicted diagnosis: {diagnosis} with confidence: {confidence:.2f}"
 
-# Triage function based on symptoms
+# Triage function to assess symptom severity
 def triage(symptoms):
-    if "chest pain" in symptoms or "shortness of breath" in symptoms:
+    if "shortness of breath" in symptoms or "chest pain" in symptoms:
         return "Urgent: Seek immediate medical attention."
     elif "fever" in symptoms and "cough" in symptoms:
         return "Mild: Likely a viral infection, monitor symptoms."
     else:
         return "Non-urgent: No immediate concern."
 
-# Combine Diagnosis and Triage
+# Combine diagnosis and triage into one function
 def full_check(symptoms):
-    diagnosis = diagnose(symptoms)
-    severity = triage(symptoms)
+    diagnosis = diagnose(symptoms)  # Get diagnosis
+    severity = triage(symptoms)     # Get severity
     return diagnosis, severity
 
-# Create the Gradio interface
-iface = gr.Interface(fn=full_check, inputs="text", outputs=["text", "text"], title="Sehat Guard - Symptom Checker")
+# Create Gradio interface
+iface = gr.Interface(
+    fn=full_check, 
+    inputs="text", 
+    outputs=["text", "text"], 
+    title="Sehat Guard - Symptom Checker",
+    description="Enter your symptoms to get a possible diagnosis and severity of the condition."
+)
 
 # Launch the app
 iface.launch()