Update app.py

app.py

@@ -1,88 +1,125 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-import time
 
-# --- MODEL SETUP ---
-# CPU-friendly LLM
-model_name = "google/flan-t5-small"  # small enough for CPU
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
-
-# --- MALARIA RULES ---
+# ===========================
+# Malaria Symptom Rules
+# ===========================
 malaria_rules = {
-    "fever": {"diagnosis": "Malaria likely", "treatment": "ACT (artemisinin-based combination therapy)", "prognosis": "Good if treated early"},
-    "chills": {"diagnosis": "Malaria likely", "treatment": "ACT", "prognosis": "Good if treated early"},
-    "headache": {"diagnosis": "Malaria possible", "treatment": "Supportive care and ACT if confirmed", "prognosis": "Good with treatment"},
-    "sweating": {"diagnosis": "Malaria possible", "treatment": "Monitor temperature, ACT if confirmed", "prognosis": "Good with treatment"},
-    "nausea": {"diagnosis": "Malaria possible", "treatment": "Hydration, ACT if confirmed", "prognosis": "Good with treatment"},
-    "vomiting": {"diagnosis": "Malaria possible", "treatment": "Hydration, ACT if confirmed", "prognosis": "Good with treatment"},
-    "fatigue": {"diagnosis": "Malaria possible", "treatment": "Rest and ACT if confirmed", "prognosis": "Good with treatment"},
-    "anemia": {"diagnosis": "Severe malaria possible", "treatment": "Hospitalization, blood transfusion, ACT", "prognosis": "Guarded, monitor closely"},
-    "jaundice": {"diagnosis": "Severe malaria possible", "treatment": "Hospitalization, supportive care, ACT", "prognosis": "Guarded"},
-    "convulsions": {"diagnosis": "Cerebral malaria possible", "treatment": "Emergency care, IV antimalarials", "prognosis": "Poor if untreated"},
-    "respiratory distress": {"diagnosis": "Severe malaria possible", "treatment": "Oxygen therapy, IV antimalarials", "prognosis": "Guarded"},
-    "abdominal pain": {"diagnosis": "Malaria possible", "treatment": "Supportive care, ACT if confirmed", "prognosis": "Good with treatment"},
-    "diarrhea": {"diagnosis": "Malaria possible", "treatment": "Hydration, ACT if confirmed", "prognosis": "Good with treatment"},
-    "muscle pain": {"diagnosis": "Malaria possible", "treatment": "Rest, analgesics, ACT if confirmed", "prognosis": "Good with treatment"},
+    "fever": {
+        "diagnosis": ["Patient may have malaria.", "Consider differential: viral infection, bacterial infection."],
+        "treatment": ["Recommend antimalarial therapy as appropriate."],
+        "prognosis": ["Mild if treated early."]
+    },
+    "chills": {
+        "diagnosis": ["Chills often accompany malaria fever.", "Monitor for shaking chills indicating severe malaria."],
+        "treatment": ["Supportive care, hydration, and antipyretics as needed."],
+        "prognosis": ["Typically resolves with treatment."]
+    },
+    "headache": {"diagnosis": ["Common in malaria patients."], "treatment": ["Analgesics as needed."], "prognosis": ["Improves with treatment."]},
+    "nausea": {"diagnosis": ["May indicate malaria or gastrointestinal infection."], "treatment": ["Anti-nausea medications if severe."], "prognosis": ["Resolves with malaria treatment."]},
+    "vomiting": {"diagnosis": ["May indicate severe malaria if combined with fever."], "treatment": ["Consider oral or IV hydration."], "prognosis": ["Monitor for dehydration."]},
+    "muscle_aches": {"diagnosis": ["Common symptom of malaria."], "treatment": ["Analgesics if necessary."], "prognosis": ["Improves with antimalarial therapy."]},
+    "fatigue": {"diagnosis": ["Common symptom, may persist during recovery."], "treatment": ["Rest and supportive care."], "prognosis": ["Gradual improvement expected."]},
+    "sweating": {"diagnosis": ["Sweating often follows fever spikes."], "treatment": ["Supportive care."], "prognosis": ["Normalizes with treatment."]},
+    "anemia": {"diagnosis": ["Severe malaria can cause hemolytic anemia."], "treatment": ["Monitor hemoglobin; blood transfusion if critical."], "prognosis": ["Serious if untreated."]},
+    "jaundice": {"diagnosis": ["May indicate liver involvement; risk for severe malaria."], "treatment": ["Monitor liver function, supportive care."], "prognosis": ["Can be severe if untreated."]},
+    "seizures": {"diagnosis": ["Cerebral malaria, medical emergency."], "treatment": ["Immediate hospitalization, IV antimalarials."], "prognosis": ["Life-threatening without prompt care."]},
+    # Combined symptom rules
+    ("fever", "chills"): {"diagnosis": ["Classic malaria presentation."], "treatment": ["Prompt antimalarial therapy."], "prognosis": ["Good if treated early."]},
+    ("fever", "vomiting"): {"diagnosis": ["Possible severe malaria."], "treatment": ["Monitor hydration, antimalarials."], "prognosis": ["Serious if untreated."]},
+    ("jaundice", "anemia"): {"diagnosis": ["Severe malaria with liver involvement."], "treatment": ["Hospitalization recommended."], "prognosis": ["Severe, close monitoring required."]},
+    ("seizures", "fever"): {"diagnosis": ["Cerebral malaria; emergency."], "treatment": ["Immediate IV therapy and ICU support."], "prognosis": ["Life-threatening."]},
 }
 
-# --- ANALYSIS FUNCTION ---
-def analyze_symptoms(symptoms_input):
-    # show temporary running feedback
-    output = "Running analysis...\n"
-    
-    symptoms = [s.strip().lower() for s in symptoms_input.split(",")]
+# ===========================
+# Analysis Functions
+# ===========================
+def analyze_malaria(symptoms):
+    diagnosis = set()
+    treatment = set()
+    prognosis = set()
+    # Check combined symptoms first
+    for combo in [k for k in malaria_rules if isinstance(k, tuple)]:
+        if all(sym in symptoms for sym in combo):
+            diagnosis.update(malaria_rules[combo]["diagnosis"])
+            treatment.update(malaria_rules[combo]["treatment"])
+            prognosis.update(malaria_rules[combo]["prognosis"])
+    # Check single symptoms
+    for sym in symptoms:
+        if sym in malaria_rules:
+            diagnosis.update(malaria_rules[sym]["diagnosis"])
+            treatment.update(malaria_rules[sym]["treatment"])
+            prognosis.update(malaria_rules[sym]["prognosis"])
+    return {"Diagnosis": "\n".join(diagnosis) if diagnosis else "No malaria-related diagnosis detected.",
+            "Treatment": "\n".join(treatment) if treatment else "No treatment recommended.",
+            "Prognosis": "\n".join(prognosis) if prognosis else "No prognosis available."}
+
+def run_analysis(age, gender, location, travel, travel_details, symptoms_list,
+                 temperature, blood_pressure, heart_rate, previous_malaria, medications, additional_notes):
     
-    diagnosis_list = []
-    treatment_list = []
-    prognosis_list = []
+    symptoms = symptoms_list
+    result = analyze_malaria(symptoms)
     
-    for symptom in symptoms:
-        if symptom in malaria_rules:
-            rule = malaria_rules[symptom]
-            diagnosis_list.append(rule["diagnosis"])
-            treatment_list.append(rule["treatment"])
-            prognosis_list.append(rule["prognosis"])
-        else:
-            diagnosis_list.append(f"No rule for '{symptom}'")
-            treatment_list.append(f"No rule for '{symptom}'")
-            prognosis_list.append(f"No rule for '{symptom}'")
+    severity = "Mild"
+    color = "#10B981"
+    if all(sym in symptoms for sym in ["seizures", "fever"]):
+        severity = "Life-threatening"
+        color = "#EF4444"
+    elif all(sym in symptoms for sym in ["jaundice", "anemia"]):
+        severity = "Severe"
+        color = "#F97316"
     
-    # Convert lists to readable text
-    diagnosis_text = "\n".join(diagnosis_list)
-    treatment_text = "\n".join(treatment_list)
-    prognosis_text = "\n".join(prognosis_list)
+    alert_banner = f"<div style='border-radius:0.5rem; padding:1rem; background-color:{color}; color:white; font-weight:bold; margin-bottom:1rem;'>⚠️ Severity Alert: {severity}</div>"
     
-    # Enhance outputs with LLM
-    enhanced_diagnosis = generator(f"Summarize and explain clearly: {diagnosis_text}", max_length=150)[0]["generated_text"]
-    enhanced_treatment = generator(f"Summarize and explain clearly: {treatment_text}", max_length=150)[0]["generated_text"]
-    enhanced_prognosis = generator(f"Summarize and explain clearly: {prognosis_text}", max_length=150)[0]["generated_text"]
+    diag_output = f"{alert_banner}<b>Diagnosis:</b><br>{result['Diagnosis'].replace(chr(10), '<br>')}"
+    treat_output = f"{alert_banner}<b>Treatment:</b><br>{result['Treatment'].replace(chr(10), '<br>')}"
+    prog_output = f"{alert_banner}<b>Prognosis:</b><br>{result['Prognosis'].replace(chr(10), '<br>')}"
     
-    return enhanced_diagnosis, enhanced_treatment, enhanced_prognosis
+    return diag_output, treat_output, prog_output
 
-# --- GRADIO UI ---
-with gr.Blocks() as demo:
-    gr.Markdown("## Malaria Multi-Agent AI")
-    
-    with gr.Row():
-        symptoms_input = gr.Textbox(label="Enter symptoms (comma-separated)", placeholder="fever, chills, headache")
-        analyze_button = gr.Button("Run Analysis")
+# ===========================
+# Gradio UI (Tailwind-inspired)
+# ===========================
+symptoms_options = ["fever", "chills", "headache", "nausea", "vomiting",
+                    "muscle_aches", "fatigue", "sweating", "anemia", "jaundice", "seizures"]
+
+with gr.Blocks(css="""
+    body {background-color:#F9FAFB;}
+    .card {border-radius:1rem; box-shadow:0 10px 25px rgba(0,0,0,0.1); background:white; padding:1.5rem;}
+    .header {background:linear-gradient(135deg,#667eea 0%,#764ba2 100%); color:white; padding:2rem; text-align:center; border-radius:1rem; margin-bottom:1rem;}
+    .btn {background-color:#5D5CDE; color:white; padding:0.75rem 1.5rem; border-radius:0.75rem; font-weight:bold;}
+""") as demo:
+    gr.HTML("<div class='header'><h1 class='text-4xl font-bold'>🦟 Malaria AI Assistant</h1><p>Advanced diagnostic, treatment, and prognostic analysis</p></div>")
     
     with gr.Row():
-        diagnosis_output = gr.Textbox(label="Diagnosis", interactive=False)
-        treatment_output = gr.Textbox(label="Treatment", interactive=False)
-        prognosis_output = gr.Textbox(label="Prognosis", interactive=False)
-    
-    def on_click(symptoms):
-        # Show temporary running message first
-        diagnosis_output.value = "Running analysis..."
-        treatment_output.value = "Running analysis..."
-        prognosis_output.value = "Running analysis..."
-        time.sleep(0.5)  # short pause to simulate processing
-        return analyze_symptoms(symptoms)
+        with gr.Column(scale=1):
+            with gr.Group(elem_classes="card"):
+                age = gr.Number(label="Age")
+                gender = gr.Dropdown(["male","female","other"], label="Gender")
+                location = gr.Textbox(label="Location")
+                travel = gr.Checkbox(label="Recent Travel to Endemic Areas")
+                travel_details = gr.Textbox(label="Travel Details")
+                symptoms_list = gr.CheckboxGroup(symptoms_options, label="Symptoms")
+                temperature = gr.Number(label="Temperature (°C)")
+                blood_pressure = gr.Textbox(label="Blood Pressure")
+                heart_rate = gr.Number(label="Heart Rate")
+                previous_malaria = gr.Checkbox(label="Previous Malaria Episodes")
+                medications = gr.Textbox(label="Current Medications / Allergies")
+                additional_notes = gr.Textbox(label="Additional Notes")
+        
+        with gr.Column(scale=1):
+            with gr.Group(elem_classes="card"):
+                diag_output = gr.HTML(label="Diagnosis")
+                treat_output = gr.HTML(label="Treatment")
+                prog_output = gr.HTML(label="Prognosis")
     
-    analyze_button.click(on_click, inputs=[symptoms_input], outputs=[diagnosis_output, treatment_output, prognosis_output])
+    run_btn = gr.Button("Run Analysis", elem_classes="btn")
+    run_btn.click(
+        run_analysis,
+        inputs=[age, gender, location, travel, travel_details, symptoms_list,
+                temperature, blood_pressure, heart_rate, previous_malaria,
+                medications, additional_notes],
+        outputs=[diag_output, treat_output, prog_output]
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()