Update app.py

app.py

@@ -1,54 +1,106 @@
-import gradio as gr
-from transformers import pipeline
-
-# Load the model with accelerate support
-generator = pipeline(
-    "text-generation", 
-    model="ibm-granite/granite-3.3-2b-instruct", 
-    device_map="auto"
-)
-
-# Predict disease from symptoms
-def predict_disease(symptoms):
-    prompt = f"User has the following symptoms: {symptoms}. What is the most probable disease?"
-    try:
-        output = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
-        if prompt in output:
-            response = output.split(prompt)[-1].strip()
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import pandas as pd
+
+# ✅ Load IBM Granite model with cache to speed up
+@st.cache_resource
+def load_model():
+    model_id = "ibm-granite/granite-3.3-2b-instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+generator = load_model()
+
+# ✅ Set Streamlit page configuration
+st.title("🩺 HealthAI – Intelligent Healthcare Assistant")
+
+# ✅ Define tabs
+tab1, tab2, tab3, tab4 = st.tabs([
+    "🧠 Patient Chat", "🧾 Disease Prediction", 
+    "💊 Treatment Plans", "📊 Health Analytics"
+])
+
+# ------------------------------
+# 🧠 TAB 1: Patient Chat
+# ------------------------------
+with tab1:
+    st.subheader("Ask any health-related question")
+    query = st.text_area("Enter your question here")
+    
+    if st.button("Get Advice", key="chat"):
+        if query.strip() == "":
+            st.warning("Please enter a question.")
         else:
-            response = output.strip()
-        return response
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Suggest home remedy from disease
-def home_remedy(disease):
-    prompt = f"Suggest a natural home remedy for the disease: {disease}."
-    try:
-        output = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
-        if prompt in output:
-            response = output.split(prompt)[-1].strip()
+            with st.spinner("Thinking..."):
+                response = generator(query, max_new_tokens=200)[0]["generated_text"]
+                st.success("AI Response:")
+                st.markdown(f"markdown\n{response}\n")
+
+# ------------------------------
+# 🧾 TAB 2: Disease Prediction
+# ------------------------------
+with tab2:
+    st.subheader("Enter your symptoms (comma-separated)")
+    symptoms = st.text_input("E.g. persistent fever, fatigue, dry cough")
+    
+    if st.button("AI Diagnose", key="predict"):
+        if symptoms.strip() == "":
+            st.warning("Please enter your symptoms.")
         else:
-            response = output.strip()
-        return response
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# UI with Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 HealthAI")
-    gr.Markdown("Identify diseases from symptoms and get natural home remedies.")
-
-    with gr.Tab("🔍 Symptoms Identifier"):
-        symptoms_input = gr.Textbox(label="Enter your symptoms", placeholder="e.g. fever, cough, sore throat")
-        symptoms_output = gr.Textbox(label="Predicted Disease")
-        symptoms_btn = gr.Button("Predict Disease")
-        symptoms_btn.click(fn=predict_disease, inputs=symptoms_input, outputs=symptoms_output)
-
-    with gr.Tab("🌿 Home Remedies"):
-        disease_input = gr.Textbox(label="Enter your disease", placeholder="e.g. Cold, Headache")
-        remedy_output = gr.Textbox(label="Suggested Remedy")
-        remedy_btn = gr.Button("Get Remedy")
-        remedy_btn.click(fn=home_remedy, inputs=disease_input, outputs=remedy_output)
-
-demo.launch()
+            prompt = (
+                f"I am feeling unwell. My symptoms are: {symptoms}.\n"
+                "Can you please suggest what possible conditions I might have based on this?\n"
+                "List top 3 possible diseases with a short reason for each, and give a seriousness score out of 10."
+            )
+            with st.spinner("Analyzing symptoms..."):
+                result = generator(prompt, max_new_tokens=300, do_sample=True)[0]['generated_text']
+                st.success("AI Prediction:")
+                st.markdown(f"markdown\n{result}\n")
+
+# ------------------------------
+# 💊 TAB 3: Treatment Plans
+# ------------------------------
+with tab3:
+    st.header("💊 Treatment Plan Generator")
+condition = st.text_input("Enter the known condition (e.g., Asthma, Diabetes)")
+
+if st.button("Get Full Treatment Plan"):
+    if not condition.strip():
+        st.warning("Please enter a condition.")
+    else:
+        with st.spinner("Generating treatment plan..."):
+
+            def get_response(prompt):
+                return generator(prompt, max_new_tokens=1000, temperature=0.7, do_sample=True)[0]['generated_text'].strip()
+
+            prompts = {
+                "1️⃣ Medications": f"What medications are usually prescribed for {condition}?",
+                "2️⃣ Diet": f"What diet is recommended for someone with {condition}?",
+                "3️⃣ Exercise": f"What type of physical activities should a person with {condition} follow?",
+                "4️⃣ Follow-Up & Monitoring": f"What follow-up steps and monitoring should be done for {condition}?",
+                "5️⃣ Precautions": f"What precautions should be taken by someone with {condition}?",
+                "6️⃣ Mental Health & Stress": f"How can someone with {condition} manage stress and mental health?"
+            }
+
+            for section, prompt in prompts.items():
+                st.subheader(section)
+                st.markdown(f"markdown\n{get_response(prompt)}\n")
+
+# ------------------------------
+# 📊 TAB 4: Health Analytics
+# ------------------------------
+with tab4:
+    st.subheader("Track your health data over time")
+    uploaded = st.file_uploader("Upload your CSV file (with columns like 'blood_pressure', 'heart_rate')", type=["csv"])
+    
+    if uploaded:
+        df = pd.read_csv(uploaded)
+        st.dataframe(df)
+
+        for col in df.select_dtypes(include=['float', 'int']).columns:
+            st.line_chart(df[col])
+            if df[col].mean() > df[col].iloc[-1]:
+                st.info(f"📉 {col} is improving.")
+            else:
+                st.warning(f"📈 {col} is rising — consider medical advice.")