CognitivePulse / app.py
kshamaasuresh's picture
Update app.py
99fe06b verified
Raw
History Blame Contribute Delete
28 kB
"""
app.py — CognitivePulse: Biomarker Intelligence & Coaching Assistant
A five-stage Streamlit application demonstrating applied ML and RAG system design
in the preventive brain health domain. Built as a proof-of-concept for BetterBrain's
existing product pipeline: biomarker intake → risk stratification → personalized
priority ranking → grounded coaching brief → patient Q&A chatbot.
Disclaimer: This is a research and engineering demonstration prototype.
It is not a validated clinical or diagnostic tool.
"""
import os
import json
import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
st.set_page_config(
page_title="CognitivePulse",
page_icon="🔬",
layout="wide",
initial_sidebar_state="expanded",
)
# ---------------------------------------------------------------------------
# Sidebar
# ---------------------------------------------------------------------------
with st.sidebar:
st.markdown("## 🔬 CognitivePulse")
st.markdown("**Biomarker Intelligence & Coaching Assistant**")
st.markdown("---")
st.warning(
"This is a research and engineering demonstration prototype. "
"It is not a validated clinical or diagnostic tool. "
"Any real deployment would require clinical validation, regulatory review, "
"and licensed oversight.",
icon="⚠️",
)
st.markdown("**What it demonstrates:**")
st.markdown(
"- XGBoost risk stratification on a real Alzheimer's research dataset\n"
"- SHAP-based explainability for per-patient biomarker contribution\n"
"- Modifiability-weighted intervention priority ranking\n"
"- RAG coaching brief grounded in the prevention literature\n"
"- RAGAS-style faithfulness evaluation\n"
"- Patient Q&A chatbot grounded in individual biomarker profile"
)
st.markdown("---")
st.markdown("**Data source:**")
st.markdown(
"El Kharoua, R. (2024). [Alzheimer's Disease Dataset](https://doi.org/10.34740/KAGGLE/DSV/8668279). "
"Kaggle. 2,149 patients, 33 features."
)
if not os.environ.get("GROQ_API_KEY"):
st.info("Set `GROQ_API_KEY` to enable AI features (Tabs 4 & 5).", icon="🔑")
# ---------------------------------------------------------------------------
# Session state + initialisation
# ---------------------------------------------------------------------------
@st.cache_resource(show_spinner="Loading dataset and training risk model…")
def init_pipeline():
from data_loader import load_dataset, get_population_stats
from risk_model import load_or_train
from rag_engine import LiteratureRetriever
df, source = load_dataset()
model, explainer, cv_results = load_or_train(df)
pop_stats = get_population_stats(df)
retriever = LiteratureRetriever()
return df, source, model, explainer, cv_results, pop_stats, retriever
df, data_source, model, explainer, cv_results, pop_stats, retriever = init_pipeline()
for key, default in [
("current_patient", None),
("risk_result", None),
("interventions", None),
("coaching", None),
("chat_messages", []),
]:
if key not in st.session_state:
st.session_state[key] = default
from data_loader import FEATURE_COLS, FEATURE_META, REFERENCE_RANGES
# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
st.title("🔬 CognitivePulse")
st.caption("Biomarker Intelligence & Coaching Assistant — Preventive Brain Health Platform")
if data_source == "synthetic":
st.info(
"Running on **synthetic data** (statistically matched to the Kaggle dataset). "
"To use real data: set `KAGGLE_USERNAME` and `KAGGLE_KEY` environment variables.",
icon="ℹ️",
)
elif data_source == "kaggle":
st.success(f"Loaded real dataset from Kaggle: {len(df):,} patients, {df.shape[1]} features.", icon="✅")
else:
st.success(f"Loaded dataset from local file: {len(df):,} patients.", icon="✅")
if cv_results:
c1, c2, c3, c4 = st.columns(4)
c1.metric("Model", "XGBoost")
c2.metric("CV AUC", f"{cv_results['auc_mean']:.3f} ± {cv_results['auc_std']:.3f}")
c3.metric("CV F1", f"{cv_results['f1_mean']:.3f} ± {cv_results['f1_std']:.3f}")
c4.metric("Training samples", f"{len(df):,}")
st.markdown("---")
# ---------------------------------------------------------------------------
# Navigation
# ---------------------------------------------------------------------------
tab1, tab2, tab3, tab4, tab5 = st.tabs([
"📊 Population Dashboard",
"👤 Patient Risk Assessment",
"🎯 Intervention Priorities",
"🤖 AI Coaching Brief",
"💬 Patient Q&A",
])
# ============================================================
# TAB 1 — Population Dashboard
# ============================================================
with tab1:
st.header("Population Dashboard")
st.caption("Global feature importance and population-level distributions across the dataset.")
col1, col2 = st.columns([1, 1])
with col1:
st.subheader("Global Feature Importance (XGBoost Gain)")
from risk_model import global_feature_importance
fi_df = global_feature_importance(model)
fig_fi = px.bar(
fi_df.head(15), x="Importance", y="Feature", orientation="h",
color="Importance", color_continuous_scale="Blues",
height=480,
)
fig_fi.update_layout(
margin=dict(l=10, r=10, t=10, b=10),
coloraxis_showscale=False,
yaxis={"autorange": "reversed"},
plot_bgcolor="rgba(0,0,0,0)",
paper_bgcolor="rgba(0,0,0,0)",
)
st.plotly_chart(fig_fi, use_container_width=True)
with col2:
st.subheader("Distribution by Diagnosis")
continuous_feats = [f for f in FEATURE_COLS if FEATURE_META[f]["type"] == "continuous"]
feat = st.selectbox(
"Select feature to visualise",
continuous_feats,
index=continuous_feats.index("MMSE") if "MMSE" in continuous_feats else 0,
)
fig_hist = px.histogram(
df, x=feat, color="Diagnosis", barmode="overlay",
color_discrete_map={0: "#4A90D9", 1: "#E05252"},
labels={"Diagnosis": "Diagnosis", feat: FEATURE_META[feat]["label"]},
opacity=0.7, height=380,
)
fig_hist.update_layout(
margin=dict(l=10, r=10, t=10, b=10),
legend=dict(title="", orientation="h", y=1.05),
plot_bgcolor="rgba(0,0,0,0)",
paper_bgcolor="rgba(0,0,0,0)",
)
st.plotly_chart(fig_hist, use_container_width=True)
st.subheader("Correlation Matrix — Modifiable Risk Factors")
mod_cols = [f for f in FEATURE_COLS if FEATURE_META[f]["modifiable"]]
corr = df[mod_cols + ["Diagnosis"]].corr()
fig_corr = px.imshow(
corr, text_auto=".2f", color_continuous_scale="RdBu_r",
zmin=-1, zmax=1, aspect="auto", height=420,
)
fig_corr.update_layout(margin=dict(l=10, r=10, t=10, b=10))
st.plotly_chart(fig_corr, use_container_width=True)
# ============================================================
# TAB 2 — Patient Risk Assessment
# ============================================================
with tab2:
st.header("Patient Risk Assessment")
st.caption(
"Enter a patient's biomarker profile below to generate an individual risk score "
"with SHAP-based feature attribution."
)
st.subheader("Enter Patient Profile")
# Human-readable option maps for categorical/binary fields.
# Gender includes "Prefer to self-identify" — this maps to 0 internally
# since the underlying dataset is binary; noted in the UI below.
OPTION_MAPS = {
"Gender": {0: "Male", 1: "Female", 2: "Prefer to self-identify"},
"Ethnicity": {0: "Caucasian", 1: "African American", 2: "Asian", 3: "Other"},
"EducationLevel":{0: "No formal education", 1: "High school",
2: "Bachelor's degree", 3: "Higher degree"},
"Smoking": {0: "No", 1: "Yes"},
"FamilyHistoryAlzheimers": {0: "No", 1: "Yes"},
"CardiovascularDisease": {0: "No", 1: "Yes"},
"Diabetes": {0: "No", 1: "Yes"},
"Depression": {0: "No", 1: "Yes"},
"HeadInjury": {0: "No", 1: "Yes"},
"Hypertension": {0: "No", 1: "Yes"},
"MemoryComplaints": {0: "No", 1: "Yes"},
"BehavioralProblems": {0: "No", 1: "Yes"},
"Confusion": {0: "No", 1: "Yes"},
"Disorientation": {0: "No", 1: "Yes"},
"PersonalityChanges": {0: "No", 1: "Yes"},
"DifficultyCompletingTasks": {0: "No", 1: "Yes"},
"Forgetfulness": {0: "No", 1: "Yes"},
}
# Pre-fill with a realistic example
defaults = {
"Age": 68, "Gender": 0, "Ethnicity": 0, "EducationLevel": 2,
"BMI": 29.5, "Smoking": 0, "AlcoholConsumption": 5.0,
"PhysicalActivity": 2.5, "DietQuality": 5.0, "SleepQuality": 6.0,
"FamilyHistoryAlzheimers": 1, "CardiovascularDisease": 1,
"Diabetes": 0, "Depression": 0, "HeadInjury": 0, "Hypertension": 1,
"SystolicBP": 148, "DiastolicBP": 88, "CholesterolTotal": 240,
"CholesterolLDL": 158, "CholesterolHDL": 45, "CholesterolTriglycerides": 185,
"MMSE": 25, "FunctionalAssessment": 7.0, "MemoryComplaints": 1,
"BehavioralProblems": 0, "ADL": 7.5, "Confusion": 0,
"Disorientation": 0, "PersonalityChanges": 0,
"DifficultyCompletingTasks": 0, "Forgetfulness": 1,
}
patient = {}
sections = {
"Demographics": ["Age", "Gender", "Ethnicity", "EducationLevel"],
"Lifestyle": ["BMI", "Smoking", "AlcoholConsumption", "PhysicalActivity",
"DietQuality", "SleepQuality"],
"Medical History": ["FamilyHistoryAlzheimers", "CardiovascularDisease", "Diabetes",
"Depression", "HeadInjury", "Hypertension"],
"Clinical Measurements": ["SystolicBP", "DiastolicBP", "CholesterolTotal",
"CholesterolLDL", "CholesterolHDL", "CholesterolTriglycerides"],
"Cognitive & Functional": ["MMSE", "FunctionalAssessment", "MemoryComplaints",
"BehavioralProblems", "ADL"],
"Symptoms": ["Confusion", "Disorientation", "PersonalityChanges",
"DifficultyCompletingTasks", "Forgetfulness"],
}
gender_val_for_model = None # tracks whether gender needs remapping
for section, features in sections.items():
with st.expander(f"**{section}**", expanded=(section in ("Lifestyle", "Clinical Measurements"))):
cols = st.columns(3)
for i, feat in enumerate(features):
meta = FEATURE_META[feat]
default = defaults.get(feat, 0)
with cols[i % 3]:
if feat in OPTION_MAPS:
opt_map = OPTION_MAPS[feat]
options = list(opt_map.keys())
selected = st.selectbox(
meta["label"],
options,
index=options.index(int(default)) if int(default) in options else 0,
format_func=lambda x, m=opt_map: m.get(x, str(x)),
key=f"pt_{feat}",
)
# "Prefer to self-identify" (value 2) maps to 0 for the binary model
if feat == "Gender" and selected == 2:
gender_val_for_model = 0
patient[feat] = 0
else:
patient[feat] = selected
elif meta["type"] in ("binary", "categorical", "ordinal"):
options = list(range(4)) if meta["type"] != "binary" else [0, 1]
patient[feat] = st.selectbox(
meta["label"], options,
index=min(int(default), len(options) - 1),
key=f"pt_{feat}",
)
else:
rng = REFERENCE_RANGES.get(feat)
if rng:
mn = float(rng["flag"][0] if rng.get("flag") else 0)
mx = float(rng["flag"][1] if rng.get("flag") else 100)
else:
mn, mx = 0.0, 100.0
val = max(mn, min(mx, float(default)))
patient[feat] = st.number_input(
meta["label"], min_value=mn, max_value=mx,
value=val, step=1.0, key=f"pt_{feat}",
)
if gender_val_for_model is not None:
st.caption(
"ℹ️ Gender is used as a binary feature in the underlying dataset (Male=0, Female=1). "
"'Prefer to self-identify' is recorded as Male for model computation only. "
"This is a known limitation of the current dataset."
)
if st.button("Generate Risk Assessment →", type="primary"):
from risk_model import predict_patient
from intervention_engine import rank_interventions
result = predict_patient(model, explainer, patient)
interventions = rank_interventions(result["shap_contributions"], patient)
st.session_state.current_patient = patient
st.session_state.risk_result = result
st.session_state.interventions = interventions
st.session_state.coaching = None
st.session_state.chat_messages = [] # reset chat when new assessment runs
if st.session_state.risk_result:
result = st.session_state.risk_result
st.markdown("---")
st.subheader("Risk Assessment Results")
band_colors = {"low": "#4CAF50", "moderate": "#FF9800", "elevated": "#FF5722", "high": "#D32F2F"}
band = result["risk_band"]
color = band_colors[band]
rc1, rc2, rc3 = st.columns(3)
rc1.metric("Risk Score", f"{result['risk_score']}/100")
rc2.metric("Risk Band", band.upper())
rc3.metric("Probability", f"{result['risk_probability']:.1%}")
fig_gauge = go.Figure(go.Indicator(
mode="gauge+number",
value=result["risk_score"],
domain={"x": [0, 1], "y": [0, 1]},
title={"text": "Brain Health Risk Score"},
gauge={
"axis": {"range": [0, 100]},
"bar": {"color": color},
"steps": [
{"range": [0, 25], "color": "#C8E6C9"},
{"range": [25, 50], "color": "#FFF9C4"},
{"range": [50, 75], "color": "#FFE0B2"},
{"range": [75, 100],"color": "#FFCDD2"},
],
"threshold": {"line": {"color": color, "width": 4}, "value": result["risk_score"]},
},
))
fig_gauge.update_layout(height=280, margin=dict(l=20, r=20, t=40, b=20),
paper_bgcolor="rgba(0,0,0,0)")
st.plotly_chart(fig_gauge, use_container_width=True)
st.subheader("Feature Attribution (SHAP Waterfall)")
shap_df = pd.DataFrame([
{"Feature": FEATURE_META.get(k, {}).get("label", k), "SHAP": v}
for k, v in result["shap_contributions"].items()
]).sort_values("SHAP", key=abs, ascending=False).head(12)
fig_shap = px.bar(
shap_df, x="SHAP", y="Feature", orientation="h",
color="SHAP", color_continuous_scale="RdBu_r",
color_continuous_midpoint=0, height=420,
)
fig_shap.update_layout(
margin=dict(l=10, r=10, t=10, b=10),
coloraxis_showscale=False,
yaxis={"autorange": "reversed"},
plot_bgcolor="rgba(0,0,0,0)",
paper_bgcolor="rgba(0,0,0,0)",
)
st.plotly_chart(fig_shap, use_container_width=True)
st.caption(
"Red bars increase predicted risk; blue bars decrease it. "
"Bar length reflects the magnitude of each feature's contribution for this specific patient."
)
# ============================================================
# TAB 3 — Intervention Priorities
# ============================================================
with tab3:
st.header("Intervention Priority Engine")
if not st.session_state.interventions:
st.info("Complete the Patient Risk Assessment (Tab 2) first.", icon="ℹ️")
else:
ivs = st.session_state.interventions
st.caption(
"Modifiable risk factors ranked by: |SHAP contribution| × actionability weight. "
"Only features at adverse levels are included."
)
if not ivs:
st.success(
"No modifiable risk factors were identified as adverse for this patient profile. "
"General brain-health maintenance is appropriate.",
icon="✅",
)
else:
for i, iv in enumerate(ivs, 1):
with st.container():
cols = st.columns([0.05, 0.95])
with cols[1]:
st.markdown(f"**{i}. {iv['intervention_summary']}**")
mc1, mc2, mc3 = st.columns(3)
mc1.metric("Priority Score", f"{iv['priority_score']:.3f}")
mc2.metric("SHAP Contribution", f"{iv['shap_value']:+.3f}")
if iv["patient_value"] is not None:
mc3.metric(iv["label"], iv["patient_value"])
st.markdown("---")
if ivs:
fig_iv = px.bar(
pd.DataFrame(ivs).rename(columns={"intervention_summary": "Intervention",
"priority_score": "Priority Score"}),
x="Priority Score", y="Intervention", orientation="h",
color="Priority Score", color_continuous_scale="Oranges", height=300,
)
fig_iv.update_layout(
margin=dict(l=10, r=10, t=10, b=10),
coloraxis_showscale=False,
yaxis={"autorange": "reversed"},
plot_bgcolor="rgba(0,0,0,0)",
paper_bgcolor="rgba(0,0,0,0)",
)
st.plotly_chart(fig_iv, use_container_width=True)
# ============================================================
# TAB 4 — AI Coaching Brief
# ============================================================
with tab4:
st.header("AI Coaching Brief")
st.caption(
"Retrieves relevant prevention literature for each priority intervention, "
"generates a grounded coaching brief, "
"then runs a RAGAS-style faithfulness check against the retrieved sources."
)
if not st.session_state.interventions:
st.info("Complete the Patient Risk Assessment (Tab 2) first.", icon="ℹ️")
elif not os.environ.get("GROQ_API_KEY"):
st.error("Set `GROQ_API_KEY` to enable the coaching assistant.", icon="🔑")
else:
if st.session_state.coaching is None:
if st.button("Generate Coaching Brief →", type="primary"):
from intervention_engine import build_coach_brief
from rag_engine import generate_coaching, check_faithfulness
from groq import Groq
client = Groq()
with st.spinner("Retrieving literature and generating coaching brief…"):
coach_brief = build_coach_brief(
st.session_state.current_patient,
st.session_state.risk_result,
st.session_state.interventions,
)
retrieved = retriever.retrieve_for_interventions(st.session_state.interventions)
coaching = generate_coaching(coach_brief, retrieved, client=client)
faithfulness = check_faithfulness(
coaching["text"], coaching["sources_used"], client=client
)
st.session_state.coaching = {
"coaching": coaching,
"faithfulness": faithfulness,
"backend": retriever.backend,
}
st.rerun()
if st.session_state.coaching:
result = st.session_state.coaching
st.markdown("### Coaching Brief")
st.markdown(result["coaching"]["text"])
if result["coaching"]["sources_used"]:
with st.expander(f"📚 Sources retrieved ({len(result['coaching']['sources_used'])})"):
for s in result["coaching"]["sources_used"]:
st.markdown(f"**{s['title']}** \n*{s['source']}* \n{s['summary']} \n[→ Source]({s['url']})")
st.markdown("---")
with st.expander("✅ Faithfulness Evaluation (RAGAS-style)"):
f = result["faithfulness"]
score = f.get("faithfulness_score")
if score is not None:
col_s, col_b = st.columns(2)
col_s.metric("Faithfulness Score", f"{score:.0%}")
col_b.caption(f"Retrieval backend: `{result['backend']}`")
for c in f.get("claims", []):
icon = {"SUPPORTED": "✅", "PARTIAL": "🟡", "UNSUPPORTED": "❌"}.get(c["verdict"], "❓")
st.write(f"{icon} **{c['verdict']}** — {c['claim']}")
st.caption(c.get("reason", ""))
if st.button("← Regenerate"):
st.session_state.coaching = None
st.rerun()
# ============================================================
# TAB 5 — Patient Q&A Chatbot
# ============================================================
with tab5:
st.header("Patient Q&A")
st.caption(
"Ask questions about your biomarker results, risk score, or brain health in general. "
"Responses are grounded in your specific profile and the prevention research literature."
)
if not st.session_state.risk_result:
st.info("Complete the Patient Risk Assessment (Tab 2) first to activate the Q&A assistant.", icon="ℹ️")
elif not os.environ.get("GROQ_API_KEY"):
st.error("Set `GROQ_API_KEY` to enable the Q&A assistant.", icon="🔑")
else:
st.warning(
"This assistant is an AI prototype, not a medical professional. "
"It cannot diagnose conditions or replace clinical advice. "
"Always consult a qualified healthcare provider for medical decisions.",
icon="⚠️",
)
# Build the system prompt from the patient's actual profile
result = st.session_state.risk_result
interventions = st.session_state.interventions or []
patient = st.session_state.current_patient or {}
# Summarise the patient's key biomarkers for the system prompt
key_biomarkers = []
for feat in ["SystolicBP", "CholesterolLDL", "CholesterolHDL", "BMI",
"PhysicalActivity", "DietQuality", "SleepQuality", "MMSE"]:
if feat in patient:
label = FEATURE_META.get(feat, {}).get("label", feat)
key_biomarkers.append(f" - {label}: {patient[feat]}")
iv_summary = "\n".join(
f" {i+1}. {iv['intervention_summary']} (priority score: {iv['priority_score']:.3f})"
for i, iv in enumerate(interventions)
) or " No modifiable adverse factors identified."
top_drivers_summary = "\n".join(
f" - {d['label']}: SHAP={d['shap_value']:+.3f} ({d['direction']})"
for d in result.get("top_drivers", [])
)
CHAT_SYSTEM_PROMPT = f"""You are a brain health assistant for a preventive neurology platform. \
You are speaking directly with a patient who has just completed a biomarker risk assessment.
PATIENT'S PROFILE:
- Risk score: {result['risk_score']}/100 ({result['risk_band'].upper()} risk band)
- Risk probability: {result['risk_probability']:.1%}
Key biomarker values:
{chr(10).join(key_biomarkers)}
Top risk drivers (SHAP-based):
{top_drivers_summary}
Prioritised intervention areas:
{iv_summary}
YOUR ROLE:
- Answer the patient's questions about their results, their biomarkers, and brain health in general.
- Refer to their specific numbers when relevant (e.g. "Your LDL of X is...").
- Be warm, clear, and non-alarmist. Use plain language, not medical jargon.
- Never make a diagnosis. Never tell someone they have or will get Alzheimer's.
- Frame everything as risk factors and evidence-based lifestyle guidance.
- If asked about something outside brain health or their results, politely redirect.
- Always end responses that touch on medical decisions with a reminder to consult their healthcare provider.
- Keep responses concise — 3 to 5 sentences unless a longer explanation is genuinely needed."""
# Display existing chat history
for msg in st.session_state.chat_messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
# Suggested starter questions (only show if no messages yet)
if not st.session_state.chat_messages:
st.markdown("**Not sure where to start? Try one of these:**")
starter_questions = [
"What does my risk score mean?",
"Which of my biomarkers are most concerning?",
"What can I do to improve my brain health?",
"How does sleep affect dementia risk?",
"What is the MIND diet?",
]
cols = st.columns(len(starter_questions))
for col, question in zip(cols, starter_questions):
if col.button(question, key=f"starter_{question}"):
st.session_state.chat_messages.append({"role": "user", "content": question})
st.rerun()
# Chat input
if prompt := st.chat_input("Ask a question about your results or brain health…"):
st.session_state.chat_messages.append({"role": "user", "content": prompt})
st.rerun()
# Generate a response if the last message is from the user
if st.session_state.chat_messages and st.session_state.chat_messages[-1]["role"] == "user":
from groq import Groq
client = Groq()
messages_for_api = [{"role": "system", "content": CHAT_SYSTEM_PROMPT}] + [
{"role": m["role"], "content": m["content"]}
for m in st.session_state.chat_messages
]
with st.chat_message("assistant"):
with st.spinner(""):
response = client.chat.completions.create(
model="openai/gpt-oss-120b",
max_tokens=600,
reasoning_effort="low",
messages=messages_for_api,
)
reply = (response.choices[0].message.content or "").strip()
st.markdown(reply)
st.session_state.chat_messages.append({"role": "assistant", "content": reply})
# Clear chat button
if st.session_state.chat_messages:
if st.button("Clear conversation", key="clear_chat"):
st.session_state.chat_messages = []
st.rerun()