Update app.py

app.py

@@ -1,105 +1,113 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# --- MODEL SETUP ---
-MODEL_ID = "TheBloke/guanaco-3B-GPTQ"  # public, CPU-friendly
-print("Loading model (this may take a minute on CPU)...")
+# -----------------------------
+# Load stronger CPU-friendly AI model
+# -----------------------------
+model_name = "google/flan-t5-base"  # upgraded from small
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto",        # CPU only
-    torch_dtype="auto",
-    low_cpu_mem_usage=True
-)
+# -----------------------------
+# AI response function
+# -----------------------------
+def malaria_ai(age, gender, location, travel_endemic, travel_details,
+               symptoms, temperature, blood_pressure, heart_rate,
+               previous_malaria, medications, additional_notes, agent):
+    
+    # Basic input check
+    if not age and not symptoms:
+        return "<p style='color:red;'>Please provide at least age or symptoms for analysis.</p>"
+    
+    symptoms_list = ", ".join(symptoms) if symptoms else "No symptoms reported"
+    
+    # Structured prompt with clear instructions
+    prompt = f"""
+Patient Information:
 
-generator = pipeline(
-    task="text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_length=512
-)
+- Age: {age}
+- Gender: {gender or 'Not specified'}
+- Location: {location or 'Not specified'}
+- Recent travel to malaria-endemic areas: {"Yes" if travel_endemic else "No"}
+- Travel details: {travel_details or 'None'}
+- Symptoms: {symptoms_list}
+- Temperature: {temperature}°C
+- Blood Pressure: {blood_pressure or 'Not recorded'}
+- Heart Rate: {heart_rate or 'Not recorded'}
+- Previous malaria episodes: {"Yes" if previous_malaria else "No"}
+- Medications/Allergies: {medications or 'None'}
+- Additional Notes: {additional_notes or 'None'}
 
-# --- HELPER FUNCTIONS ---
-def format_patient_data(data):
-    symptoms = ", ".join(data["symptoms"]) if data["symptoms"] else "No symptoms reported"
-    text = f"""
-Patient Data:
-
-- Age: {data.get('age', 'Not specified')}
-- Gender: {data.get('gender', 'Not specified')}
-- Location: {data.get('location', 'Not specified')}
-- Travel to endemic areas: {"Yes" if data.get("travel") else "No"}
-- Travel details: {data.get('travel_details', '')}
-
-Symptoms: {symptoms}
-
-Vital Signs:
-- Temperature: {data.get('temperature', 'Not recorded')}
-- Blood Pressure: {data.get('blood_pressure', 'Not recorded')}
-- Heart Rate: {data.get('heart_rate', 'Not recorded')}
-
-Medical History:
-- Previous malaria: {"Yes" if data.get("previous_malaria") else "No"}
-- Medications/allergies: {data.get('medications', '')}
-
-Additional notes: {data.get('additional_notes', '')}
+Task:
+You are a medical assistant AI. Based on the patient information, provide a detailed and realistic {agent.lower()} analysis of malaria. Use evidence-based reasoning. Include any necessary precautions or warnings. Format your response clearly with bullet points and paragraphs. Keep it informative but concise.
 """
-    return text.strip()
 
-def run_analysis(agent, patient_data):
-    prompt = format_patient_data(patient_data)
-    prompt += f"\n\n{agent.upper()} ANALYSIS REQUEST:\nProvide evidence-based medical information for educational purposes."
+    # Generate AI response
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=600,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.7
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    response = generator(prompt, max_length=512, do_sample=True, temperature=0.7)[0]['generated_text']
-    return response
+    # Replace line breaks safely outside f-string
+    response_html = response.replace("\n", "<br>")
 
-# --- GRADIO UI ---
-def build_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("# 🦟 Malaria AI Assistant (CPU-friendly, Open-Source)")
-        with gr.Row():
-            with gr.Column():
-                age = gr.Number(label="Age", value=None)
-                gender = gr.Dropdown(["Male", "Female", "Other"], label="Gender")
-                location = gr.Textbox(label="Location")
-                travel = gr.Checkbox(label="Traveled to endemic area?")
-                travel_details = gr.Textbox(label="Travel Details")
-                symptoms = gr.CheckboxGroup(
-                    ["Fever", "Chills", "Headache", "Nausea/Vomiting", "Muscle aches", "Fatigue"],
-                    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="Medications / Allergies")
-                additional_notes = gr.Textbox(label="Additional notes")
-            with gr.Column():
-                agent = gr.Radio(["diagnostic", "treatment", "prognosis"], label="Analysis Type", value="diagnostic")
-                output = gr.Textbox(label="Analysis Result", interactive=False)
+    # Scrollable styled HTML card
+    formatted_response = f"""
+<div style="
+    background-color:#dbeafe;
+    padding:20px;
+    border-radius:10px;
+    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
+    max-height:400px;
+    overflow-y:auto;
+    font-family: sans-serif;
+">
+<p style="color:#111;">{response_html}</p>
+</div>
+"""
+    return formatted_response
 
-                btn = gr.Button("Run Analysis")
-                btn.click(
-                    run_analysis,
-                    inputs=[agent, {
-                        "age": age,
-                        "gender": gender,
-                        "location": location,
-                        "travel": travel,
-                        "travel_details": travel_details,
-                        "symptoms": symptoms,
-                        "temperature": temperature,
-                        "blood_pressure": blood_pressure,
-                        "heart_rate": heart_rate,
-                        "previous_malaria": previous_malaria,
-                        "medications": medications,
-                        "additional_notes": additional_notes
-                    }],
-                    outputs=output
-                )
-    return demo
+# -----------------------------
+# Gradio interface with default demo data
+# -----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## 🦟 Malaria AI Assistant\nDiagnostic, treatment, and prognostic analysis")
+    
+    with gr.Row():
+        with gr.Column():
+            age = gr.Number(label="Age", value=25)
+            gender = gr.Dropdown(["", "Male", "Female", "Other"], label="Gender", value="Male")
+            location = gr.Textbox(label="Location (City, Country)", value="Lagos, Nigeria")
+            travel_endemic = gr.Checkbox(label="Recent travel to malaria-endemic areas", value=True)
+            travel_details = gr.Textbox(label="Travel Details", value="Visited rural Northern Nigeria for 2 weeks")
+            symptoms = gr.CheckboxGroup(
+                ["Fever","Chills","Headache","Nausea/Vomiting","Muscle aches","Fatigue"],
+                label="Symptoms",
+                value=["Fever","Chills","Headache"]
+            )
+            temperature = gr.Number(label="Temperature (°C)", value=38.5)
+            blood_pressure = gr.Textbox(label="Blood Pressure", value="120/80")
+            heart_rate = gr.Number(label="Heart Rate (bpm)", value=88)
+            previous_malaria = gr.Checkbox(label="Previous malaria episodes", value=True)
+            medications = gr.Textbox(label="Current medications/allergies", value="None")
+            additional_notes = gr.Textbox(label="Additional information", value="Patient shows early signs of fatigue.")
+            agent = gr.Radio(["Diagnostic", "Treatment", "Prognostic"], label="AI Analysis Type", value="Diagnostic")
+            submit_btn = gr.Button("Run Analysis")
+        
+        with gr.Column():
+            output = gr.HTML(label="AI Analysis Result")
+    
+    submit_btn.click(
+        fn=malaria_ai,
+        inputs=[age, gender, location, travel_endemic, travel_details,
+                symptoms, temperature, blood_pressure, heart_rate,
+                previous_malaria, medications, additional_notes, agent],
+        outputs=output
+    )
 
-if __name__ == "__main__":
-    demo = build_ui()
-    demo.launch()
+demo.launch()