Update app.py

app.py

@@ -1,16 +1,24 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load a pre-trained question-answering model
-model = pipeline("question-answering", model="deepset/roberta-base-squad2")
+# Load a pre-trained zero-shot classification model (e.g., BART)
+model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
-# Function to answer based on symptoms
+# Broad dynamic categories for diagnosis (you can modify these if needed)
+candidate_labels = [
+    "respiratory infection", "viral infection", "bacterial infection", "autoimmune disease", 
+    "cardiovascular disease", "endocrine disorders", "gastrointestinal disorders", 
+    "neurological disorders", "skin disorders", "genetic disorders", "cancer", "kidney disease"
+]
+
+# Function to diagnose based on symptoms
 def diagnose(symptoms):
-    question = f"Given these symptoms: {symptoms}, what is the possible diagnosis?"
-    answer = model(question=question, context="The model will use this context to infer the diagnosis.")  # Placeholder context
-    return answer['answer']
+    result = model(symptoms, candidate_labels=candidate_labels)
+    diagnosis_category = result['labels'][0]  # Top predicted category
+    confidence = result['scores'][0]  # Confidence score
+    return f"Predicted diagnosis category: {diagnosis_category} with confidence: {confidence:.2f}"
 
-# Triage function (same as before)
+# Triage function to assess symptom severity
 def triage(symptoms):
     if "shortness of breath" in symptoms or "chest pain" in symptoms:
         return "Urgent: Seek immediate medical attention."
@@ -21,7 +29,7 @@ def triage(symptoms):
 
 # Combine diagnosis and triage into one function
 def full_check(symptoms):
-    diagnosis = diagnose(symptoms)  # Get diagnosis from QA model
+    diagnosis = diagnose(symptoms)  # Get diagnosis category
     severity = triage(symptoms)     # Get severity level
     return diagnosis, severity