Update app.py

app.py

@@ -1,21 +1,16 @@
 import gradio as gr
 from transformers import pipeline
 
-# 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")
+# Load a pre-trained question-answering model
+model = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-# Function to diagnose based on symptoms
+# Function to answer based on symptoms
 def diagnose(symptoms):
-    # 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}"
+    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']
 
-# Triage function to assess symptom severity
+# Triage function (same as before)
 def triage(symptoms):
     if "shortness of breath" in symptoms or "chest pain" in symptoms:
         return "Urgent: Seek immediate medical attention."
@@ -26,11 +21,11 @@ def triage(symptoms):
 
 # Combine diagnosis and triage into one function
 def full_check(symptoms):
-    diagnosis = diagnose(symptoms)  # Get diagnosis
-    severity = triage(symptoms)     # Get severity
+    diagnosis = diagnose(symptoms)  # Get diagnosis from QA model
+    severity = triage(symptoms)     # Get severity level
     return diagnosis, severity
 
-# Create Gradio interface
+# Create the Gradio interface
 iface = gr.Interface(
     fn=full_check, 
     inputs="text",