app.py

import streamlit as st
import pandas as pd
import numpy as np
import os
import json
import faiss
import torch

from sentence_transformers import SentenceTransformer
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM
)


# ---------------------------------------------------------
# PAGE CONFIG
# ---------------------------------------------------------

st.set_page_config(
    page_title="RxReview",
    page_icon="💊",
    layout="wide"
)


# ---------------------------------------------------------
# CONSTANTS
# ---------------------------------------------------------

DATA_PATH = "rxreview_priority_data.csv"

RAG_DIR = "rag"

RAG_CHUNKS_PATH = os.path.join(
    RAG_DIR,
    "rxreview_rag_chunks.csv"
)

RAG_INDEX_PATH = os.path.join(
    RAG_DIR,
    "rxreview_faiss.index"
)

RAG_CONFIG_PATH = os.path.join(
    RAG_DIR,
    "rxreview_rag_config.json"
)


# ---------------------------------------------------------
# LOAD DATA
# ---------------------------------------------------------

@st.cache_data
def load_priority_data():

    return pd.read_csv(
        DATA_PATH
    )


priority_df = load_priority_data()


# ---------------------------------------------------------
# VALUE SIMULATOR
# ---------------------------------------------------------

def select_by_capacity(
    data,
    score_col,
    capacity_fraction
):

    temp = data.sort_values(
        score_col,
        ascending=False
    ).reset_index(drop=True)

    n_selected = int(
        np.ceil(
            len(temp)
            * capacity_fraction
        )
    )

    return temp.iloc[
        :n_selected
    ].copy()


def rxreview_value_simulator(
    data,
    capacity_fraction,
    pharmacist_hourly_cost,
    review_minutes_per_patient,
    cost_per_readmission,
    preventable_fraction,
    intervention_effectiveness
):

    selected = select_by_capacity(
        data,
        "rxreview_priority_score",
        capacity_fraction
    )

    patients_selected = len(
        selected
    )

    total_readmissions = int(
        data[
            "readmit_30"
        ].sum()
    )

    captured_readmissions = int(
        selected[
            "readmit_30"
        ].sum()
    )

    pharmacist_hours = (
        patients_selected
        *
        review_minutes_per_patient
        /
        60
    )

    intervention_cost = (
        pharmacist_hours
        *
        pharmacist_hourly_cost
    )

    capture_rate = (
        captured_readmissions
        /
        total_readmissions
        if total_readmissions > 0
        else 0
    )

    lift = (
        capture_rate
        /
        capacity_fraction
        if capacity_fraction > 0
        else 0
    )

    potentially_preventable = (
        captured_readmissions
        *
        preventable_fraction
    )

    avoided_readmissions = (
        potentially_preventable
        *
        intervention_effectiveness
    )

    gross_savings = (
        avoided_readmissions
        *
        cost_per_readmission
    )

    net_value = (
        gross_savings
        -
        intervention_cost
    )

    denominator = (
        captured_readmissions
        *
        preventable_fraction
        *
        cost_per_readmission
    )

    break_even_effectiveness = (
        intervention_cost
        /
        denominator
        if denominator > 0
        else np.nan
    )

    return {

        "patients_selected":
            patients_selected,

        "pharmacist_hours":
            pharmacist_hours,

        "captured_readmissions":
            captured_readmissions,

        "capture_rate":
            capture_rate,

        "lift":
            lift,

        "intervention_cost":
            intervention_cost,

        "avoided_readmissions":
            avoided_readmissions,

        "gross_savings":
            gross_savings,

        "net_value":
            net_value,

        "break_even_effectiveness":
            break_even_effectiveness,

        "selected":
            selected
    }


# ---------------------------------------------------------
# RAG LOADERS
# ---------------------------------------------------------

@st.cache_resource
def load_embedding_model():

    return SentenceTransformer(
        "sentence-transformers/all-MiniLM-L6-v2"
    )


@st.cache_resource
def load_faiss_index():

    return faiss.read_index(
        RAG_INDEX_PATH
    )


@st.cache_data
def load_rag_chunks():

    return pd.read_csv(
        RAG_CHUNKS_PATH
    )


@st.cache_resource
def load_generation_model():

    model_name = (
        "Qwen/Qwen2.5-1.5B-Instruct"
    )

    tokenizer = (
        AutoTokenizer.from_pretrained(
            model_name
        )
    )

    model = (
        AutoModelForCausalLM
        .from_pretrained(
            model_name,
            torch_dtype="auto",
            device_map="auto"
        )
    )

    model.eval()

    return tokenizer, model


# ---------------------------------------------------------
# RAG FUNCTIONS
# ---------------------------------------------------------

def retrieve_rag_chunks(
    question,
    top_k=4
):

    embedding_model = (
        load_embedding_model()
    )

    rag_index = (
        load_faiss_index()
    )

    chunks = (
        load_rag_chunks()
    )

    query_embedding = (
        embedding_model.encode(
            [question],
            convert_to_numpy=True,
            normalize_embeddings=True
        )
    )

    scores, indices = (
        rag_index.search(
            query_embedding.astype(
                "float32"
            ),
            top_k
        )
    )

    results = (
        chunks.iloc[
            indices[0]
        ].copy()
    )

    results[
        "similarity_score"
    ] = scores[0]

    return results


def build_rag_messages(
    question,
    retrieved_chunks
):

    sections = []

    for i, (_, row) in enumerate(
        retrieved_chunks.iterrows(),
        start=1
    ):

        sections.append(
            f"""
SOURCE {i}
Source file: {row['source_file']}
Chunk ID: {row['chunk_id']}

{row['text']}
"""
        )

    context = "\n".join(
        sections
    )

    system_message = """
You are the RxReview discharge-support assistant.

Answer using only the retrieved source context.

Do not provide:
- patient-specific prescribing advice,
- medication dosing recommendations,
- recommendations to start, stop, or change medication.

If the source context does not explicitly support the requested information,
state that the retrieved RxReview sources do not provide enough information.

Keep answers concise and practical.
"""

    user_message = f"""
QUESTION:
{question}

SOURCE CONTEXT:
{context}

Answer using only the context above.
"""

    return [
        {
            "role": "system",
            "content": system_message
        },
        {
            "role": "user",
            "content": user_message
        }
    ]


def answer_rag_question(
    question,
    top_k=4,
    min_retrieval_score=0.40
):

    retrieved = (
        retrieve_rag_chunks(
            question,
            top_k
        )
    )

    top_score = (
        retrieved[
            "similarity_score"
        ].max()
        if len(retrieved) > 0
        else 0
    )

    if top_score < min_retrieval_score:

        answer = (
            "The retrieved RxReview sources "
            "do not provide enough information "
            "to answer this question."
        )

        return answer, retrieved

    tokenizer, model = (
        load_generation_model()
    )

    messages = (
        build_rag_messages(
            question,
            retrieved
        )
    )

    model_inputs = (
        tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True
        )
    )

    model_inputs = {
        key:
            value.to(
                model.device
            )
        for key, value
        in model_inputs.items()
    }

    input_length = (
        model_inputs[
            "input_ids"
        ].shape[-1]
    )

    with torch.no_grad():

        generated_ids = (
            model.generate(
                **model_inputs,
                max_new_tokens=250,
                do_sample=False,
                repetition_penalty=1.05
            )
        )

    new_tokens = (
        generated_ids[
            0,
            input_length:
        ]
    )

    answer = (
        tokenizer.decode(
            new_tokens,
            skip_special_tokens=True
        )
    )

    return (
        answer.strip(),
        retrieved
    )


# ---------------------------------------------------------
# HEADER
# ---------------------------------------------------------

st.title(
    "💊 RxReview"
)

st.subheader(
    "Pharmacist-Led Diabetes "
    "Discharge Review Decision Support"
)

st.caption(
    "Prioritize limited pharmacist capacity "
    "using readmission risk, medication complexity, "
    "cost/value analysis, and grounded discharge guidance."
)


# ---------------------------------------------------------
# TABS
# ---------------------------------------------------------

tab1, tab2, tab3, tab4, tab5 = (
    st.tabs([
        "Dashboard",
        "Review Queue",
        "Cost & Capacity",
        "Discharge Guidance",
        "About"
    ])
)


# ---------------------------------------------------------
# TAB 1 — DASHBOARD
# ---------------------------------------------------------

with tab1:

    st.header(
        "RxReview Dashboard"
    )

    capacity = st.slider(
        "Pharmacist Review Capacity (%)",
        min_value=5,
        max_value=30,
        value=10,
        step=5
    )

    selected = select_by_capacity(
        priority_df,
        "rxreview_priority_score",
        capacity / 100
    )

    total_readmissions = int(
        priority_df[
            "readmit_30"
        ].sum()
    )

    captured = int(
        selected[
            "readmit_30"
        ].sum()
    )

    capture_rate = (
        captured
        /
        total_readmissions
    )

    lift = (
        capture_rate
        /
        (capacity / 100)
    )

    col1, col2, col3, col4 = (
        st.columns(4)
    )

    col1.metric(
        "Patients Reviewed",
        f"{len(selected):,}"
    )

    col2.metric(
        "Readmissions Captured",
        f"{captured:,}"
    )

    col3.metric(
        "Capture Rate",
        f"{capture_rate:.1%}"
    )

    col4.metric(
        "Lift vs Random",
        f"{lift:.2f}×"
    )

    st.markdown(
        "### Capacity vs Capture"
    )

    chart_rows = []

    for c in [
        0.05,
        0.10,
        0.15,
        0.20,
        0.25,
        0.30
    ]:

        temp = select_by_capacity(
            priority_df,
            "rxreview_priority_score",
            c
        )

        captured_c = (
            temp[
                "readmit_30"
            ].sum()
        )

        chart_rows.append({
            "Capacity %":
                c * 100,

            "Readmissions Captured %":
                (
                    captured_c
                    /
                    total_readmissions
                    * 100
                )
        })

    chart_df = pd.DataFrame(
        chart_rows
    )

    st.line_chart(
        chart_df,
        x="Capacity %",
        y="Readmissions Captured %"
    )


# ---------------------------------------------------------
# TAB 2 — REVIEW QUEUE
# ---------------------------------------------------------

with tab2:

    st.header(
        "Patient Review Queue"
    )

    queue_capacity = (
        st.slider(
            "Queue Capacity (%)",
            min_value=5,
            max_value=30,
            value=10,
            step=5,
            key="queue_capacity"
        )
    )

    queue_df = select_by_capacity(
        priority_df,
        "rxreview_priority_score",
        queue_capacity / 100
    )

    display_cols = [
        col
        for col in [
            "patient_nbr",
            "xgb_calibrated_readmission_risk",
            "medication_complexity_score",
            "medication_complexity_weight",
            "rxreview_priority_score",
            "rxreview_priority_score_100",
            "medication_complexity_level",
            "rxreview_priority_level"
        ]
        if col in queue_df.columns
    ]

    st.dataframe(
        queue_df[
            display_cols
        ].head(100),
        use_container_width=True
    )

    st.caption(
        "The queue is ranked by the RxReview "
        "Priority Score. Scores support prioritization "
        "and do not replace clinical judgment."
    )


# ---------------------------------------------------------
# TAB 3 — COST & CAPACITY
# ---------------------------------------------------------

with tab3:

    st.header(
        "Cost & Capacity Simulator"
    )

    col1, col2 = (
        st.columns(2)
    )

    with col1:

        sim_capacity = st.slider(
            "Review Capacity (%)",
            5,
            30,
            10,
            5
        )

        hourly_cost = st.number_input(
            "Pharmacist Hourly Cost ($)",
            min_value=0.0,
            value=75.0,
            step=5.0
        )

        review_minutes = st.slider(
            "Minutes per Review",
            10,
            60,
            30,
            5
        )

    with col2:

        readmission_cost = (
            st.number_input(
                "Cost per Readmission ($)",
                min_value=0.0,
                value=15000.0,
                step=1000.0
            )
        )

        preventable_pct = (
            st.slider(
                "Potentially Preventable (%)",
                0,
                100,
                25,
                5
            )
        )

        effectiveness_pct = (
            st.slider(
                "Intervention Effectiveness (%)",
                0,
                100,
                20,
                5
            )
        )

    sim = rxreview_value_simulator(

        priority_df,

        sim_capacity / 100,

        hourly_cost,

        review_minutes,

        readmission_cost,

        preventable_pct / 100,

        effectiveness_pct / 100
    )

    st.markdown(
        "### Scenario Results"
    )

    c1, c2, c3, c4 = (
        st.columns(4)
    )

    c1.metric(
        "Patients Reviewed",
        f"{sim['patients_selected']:,}"
    )

    c2.metric(
        "Pharmacist Hours",
        f"{sim['pharmacist_hours']:,.0f}"
    )

    c3.metric(
        "Readmissions Captured",
        f"{sim['captured_readmissions']:,}"
    )

    c4.metric(
        "Lift",
        f"{sim['lift']:.2f}×"
    )

    c5, c6, c7 = (
        st.columns(3)
    )

    c5.metric(
        "Intervention Cost",
        f"${sim['intervention_cost']:,.0f}"
    )

    c6.metric(
        "Gross Savings",
        f"${sim['gross_savings']:,.0f}"
    )

    c7.metric(
        "Net Value",
        f"${sim['net_value']:,.0f}"
    )

    st.metric(
        "Break-Even Effectiveness",
        (
            f"{sim['break_even_effectiveness']:.1%}"
        )
    )

    st.caption(
        "Financial outputs are scenario estimates "
        "based on adjustable assumptions, not observed "
        "savings from the UCI dataset."
    )


# ---------------------------------------------------------
# TAB 4 — RAG
# ---------------------------------------------------------

with tab4:

    st.header(
        "Discharge Guidance Assistant"
    )

    st.write(
        "Ask questions about medication reconciliation, "
        "discharge planning, patient education, "
        "follow-up, care transitions, or "
        "readmission-reduction practices."
    )

    question = st.text_area(
        "Question",
        placeholder=(
            "Example: What should hospitals do "
            "to reconcile medications at discharge?"
        )
    )

    if st.button(
        "Get Grounded Guidance"
    ):

        if question.strip():

            with st.spinner(
                "Retrieving evidence and generating response..."
            ):

                answer, sources = (
                    answer_rag_question(
                        question
                    )
                )

            st.markdown(
                "### Answer"
            )

            st.write(
                answer
            )

            st.markdown(
                "### Retrieved Sources"
            )

            st.dataframe(
                sources[
                    [
                        "source_file",
                        "chunk_id",
                        "similarity_score"
                    ]
                ],
                use_container_width=True
            )

        else:

            st.warning(
                "Enter a question first."
            )


# ---------------------------------------------------------
# TAB 5 — ABOUT
# ---------------------------------------------------------

with tab5:

    st.header(
        "About RxReview"
    )

    st.markdown(
        """
**RxReview** is a decision-support prototype designed to
help hospitals prioritize patients with diabetes for
pharmacist-led discharge review when pharmacist capacity
is limited.

The application combines:

- calibrated 30-day readmission risk,
- medication complexity,
- operational capacity constraints,
- scenario-based cost/value analysis,
- and authoritative AHRQ/CMS discharge guidance.

### Important limitation

RxReview does not determine whether a medication regimen
is clinically correct and does not recommend medication
changes or doses.

It supports prioritization for human pharmacist review.
"""
    )