Update app.py

app.py

@@ -1,50 +1,88 @@
-import time
-from transformers import AutoTokenizer, AutoModelForMaskedLM, pipeline
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import time
 
-# Model
-model_name = "emilyalsentzer/Bio_ClinicalBERT"
+# --- MODEL SETUP ---
+# CPU-friendly LLM
+model_name = "google/flan-t5-small"  # small enough for CPU
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForMaskedLM.from_pretrained(model_name)
-
-# Pipeline
-nlp_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)
-
-def analyze_malaria(text):
-    yield "Running analysis..."  # Status message
-    time.sleep(0.5)
-
-    # Diagnosis agent
-    try:
-        diagnosis_results = nlp_pipeline(f"{text} The patient may have [MASK].")
-        diagnosis = ", ".join([res['sequence'] for res in diagnosis_results[:3]])
-    except:
-        diagnosis = "Diagnosis analysis failed."
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-    # Treatment agent
-    try:
-        treatment_results = nlp_pipeline(f"For malaria, recommended treatment is [MASK].")
-        treatment = ", ".join([res['sequence'] for res in treatment_results[:3]])
-    except:
-        treatment = "Treatment analysis failed."
+generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-    # Prognosis agent
-    try:
-        prognosis_results = nlp_pipeline(f"Prognosis for malaria patient is [MASK].")
-        prognosis = ", ".join([res['sequence'] for res in prognosis_results[:3]])
-    except:
-        prognosis = "Prognosis analysis failed."
+# --- MALARIA 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"},
+}
 
-    yield f"Diagnosis:\n{diagnosis}\n\nTreatment:\n{treatment}\n\nPrognosis:\n{prognosis}"
+# --- 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(",")]
+    
+    diagnosis_list = []
+    treatment_list = []
+    prognosis_list = []
+    
+    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}'")
+    
+    # Convert lists to readable text
+    diagnosis_text = "\n".join(diagnosis_list)
+    treatment_text = "\n".join(treatment_list)
+    prognosis_text = "\n".join(prognosis_list)
+    
+    # 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"]
+    
+    return enhanced_diagnosis, enhanced_treatment, enhanced_prognosis
 
-# Gradio interface
-iface = gr.Interface(
-    fn=analyze_malaria,
-    inputs=gr.Textbox(lines=4, placeholder="Enter symptoms here..."),
-    outputs=gr.Textbox(label="Malaria Analysis"),
-    title="Malaria Multi-Agent AI",
-    description="Enter clinical symptoms. The AI provides diagnosis, treatment, and prognosis using Bio_ClinicalBERT.",
-)
+# --- 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")
+    
+    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)
+    
+    analyze_button.click(on_click, inputs=[symptoms_input], outputs=[diagnosis_output, treatment_output, prognosis_output])
 
-if __name__ == "__main__":
-    iface.launch()
+demo.launch()