Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load a pre-trained zero-shot classification model (e.g., BART)
-model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+# Load a more suitable model for medical diagnosis (BioBERT)
+model = pipeline("zero-shot-classification", model="dmis-lab/biobert-base-cased-v1.1")
 
-# Broad dynamic categories for diagnosis (you can modify these if needed)
+# Candidate labels for medical diagnosis
 candidate_labels = [
     "respiratory infection", "viral infection", "bacterial infection", "autoimmune disease", 
     "cardiovascular disease", "endocrine disorders", "gastrointestinal disorders", 
@@ -16,14 +16,28 @@ def diagnose(symptoms):
     result = model(symptoms, candidate_labels=candidate_labels)
     diagnosis_category = result['labels'][0]  # Top predicted category
     confidence = result['scores'][0]  # Confidence score
+
+    # Only accept predictions with confidence >= 0.60
+    if confidence < 0.60:
+        return "Unable to provide a confident diagnosis. Please consult a healthcare professional."
+
     return f"Predicted diagnosis category: {diagnosis_category} with confidence: {confidence:.2f}"
 
 # Triage function to assess symptom severity
 def triage(symptoms):
+    # Check for critical symptoms that need urgent attention
     if "shortness of breath" in symptoms or "chest pain" in symptoms:
         return "Urgent: Seek immediate medical attention."
+
+    # Check for milder symptoms
     elif "fever" in symptoms and "cough" in symptoms:
         return "Mild: Likely a viral infection, monitor symptoms."
+
+    # New logic for symptoms like "red nose and cheeks"
+    elif "red nose" in symptoms or "red cheeks" in symptoms:
+        return "Non-urgent: Likely a mild skin reaction. Monitor for any additional symptoms."
+
+    # Default case if no specific symptoms are detected
     else:
         return "Non-urgent: No immediate concern."