app.py

import os
import re
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ── Config ───────────────────────────────────────────────────────────────────
BARRIER_REPO = "marvin-cusm-chatbot/flan-t5-xl-barrier"
RISKLVL_REPO = "marvin-cusm-chatbot/flan-t5-large-risklvl"

# Set HF_TOKEN in Space Settings → Secrets if repos are private
HF_TOKEN = os.environ.get("HF_TOKEN", None)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE  = torch.float16 if torch.cuda.is_available() else torch.float32

# ── Label mapping (exactly as in your notebook) ──────────────────────────────
label_mapping = [
    [
        '1 - Your thoughts and feelings',
        '2 - Your habits and activities',
        '3 - Your social situation',
        '4 - Your economic situation',
        '5 - Your medication',
        '6 - Your care',
        '7 - Your health',
        '8 - None',
    ],
    [
        '1 - High',
        '2 - Medium',
        '3 - Low',
        '4 - None',
    ]
]

# ── Load merged models (no PEFT needed after merge_and_upload) ───────────────
print(f"Loading barrier model on {DEVICE}…")
tokenizer_barrier = AutoTokenizer.from_pretrained(BARRIER_REPO, token=HF_TOKEN)
model_barrier = AutoModelForSequenceClassification.from_pretrained(
    BARRIER_REPO, torch_dtype=DTYPE, token=HF_TOKEN,
)
model_barrier.eval().to(DEVICE)

print("Loading risk level model…")
tokenizer_risklvl = AutoTokenizer.from_pretrained(RISKLVL_REPO, token=HF_TOKEN)
model_risklvl = AutoModelForSequenceClassification.from_pretrained(
    RISKLVL_REPO, torch_dtype=DTYPE, token=HF_TOKEN,
)
model_risklvl.eval().to(DEVICE)

print("Both models ready.")

# ── Preprocessing (exactly as in your notebook) ──────────────────────────────
def preprocess_text(text):
    text = re.sub(r'<[^>]+>', '', str(text))
    text = re.sub(r'([^\w\s])\1+', r'\1', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

def split_into_sentences(text):
    sentences = re.split(r'[.?!]', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    return sentences

# ── Predict functions (exactly as in your notebook) ──────────────────────────
def predict_barrier(text, lm):
    preprocessed_text = preprocess_text(text)
    tokenized = tokenizer_barrier(preprocessed_text, return_tensors='pt', padding=True, truncation=True).to(DEVICE)
    with torch.no_grad():
        logits = model_barrier(**tokenized).logits
    predicted_class = torch.argmax(logits, dim=1).cpu()
    return lm[predicted_class]

def predict_risklvl(text, lm):
    preprocessed_text = preprocess_text(text)
    tokenized = tokenizer_risklvl(preprocessed_text, return_tensors='pt', padding=True, truncation=True).to(DEVICE)
    with torch.no_grad():
        logits = model_risklvl(**tokenized).logits
    predicted_class = torch.argmax(logits, dim=1).cpu()
    return lm[predicted_class]

# ── Risk ordering ─────────────────────────────────────────────────────────────
risk_order = ['High', 'Medium', 'Low', 'None']

# ── Full pipeline (mirrors your main() and gradio_predict exactly) ────────────
def run_triage(text: str):
    if not text.strip():
        return "", "", ""

    sentences = split_into_sentences(text)
    barrier_risk_levels = {}
    sentence_predictions = []

    for i, sentence in enumerate(sentences):
        prediction_barrier = predict_barrier(sentence, label_mapping[0])
        prediction_risklvl = predict_risklvl(sentence, label_mapping[1])

        sentence_predictions.append(
            f"S{i+1} — barrier = {prediction_barrier};   risk level = {prediction_risklvl}"
        )

        current_risk_level = prediction_risklvl.split(' - ')[1]
        if prediction_barrier not in barrier_risk_levels:
            barrier_risk_levels[prediction_barrier] = current_risk_level
        else:
            current_risk_index = risk_order.index(current_risk_level)
            stored_risk_index  = risk_order.index(barrier_risk_levels[prediction_barrier])
            if current_risk_index < stored_risk_index:
                barrier_risk_levels[prediction_barrier] = current_risk_level

    # ── Summary (mirrors your barriers_1_to_7 logic exactly) ──────────────
    summary_lines = []
    barriers_1_to_7_present = any(
        barrier.startswith(tuple(str(i) for i in range(1, 8)))
        for barrier in barrier_risk_levels
    )

    if barriers_1_to_7_present:
        for i in range(1, 8):
            matched = [b for b in barrier_risk_levels if b.startswith(f"{i} -")]
            for barrier in matched:
                risk = barrier_risk_levels[barrier]
                flag = "🔴" if risk == "High" else "🟡" if risk == "Medium" else "🟢" if risk == "Low" else "⚪"
                summary_lines.append(f"{flag}  {barrier}  →  {risk}")
    else:
        if '8 - None' in barrier_risk_levels:
            summary_lines.append("⚪  8 - None  →  No barrier identified")

    summary_text = "\n".join(summary_lines) if summary_lines else "No barriers detected."

    # ── Triage action ──────────────────────────────────────────────────────
    all_risks = list(barrier_risk_levels.values())
    if "High" in all_risks:
        action = "🚨  HIGH — Immediate follow-up recommended"
    elif "Medium" in all_risks:
        action = "⚠️   MEDIUM — Schedule check-in within 48 hours"
    elif "Low" in all_risks:
        action = "ℹ️   LOW — Note for next routine visit"
    else:
        action = "✅  NONE — No immediate action required"

    sentence_text = "\n".join(sentence_predictions)
    return sentence_text, summary_text, action

# ── Examples (from your notebook) ────────────────────────────────────────────
EXAMPLES = [
    ["Anyone on Biktarvy have any issues with insomnia? I was on Triumeq but had issues with insomnia. Switched back to my old meds. Now am considering Biktarvy."],
    ["I am considering a switch from Biktarvy to either Dovato or Cabenuva. Is there a difference in the possible negative effects to the liver when a medication is taken via injection rather than orally?"],
    ["I recently arrived in Canada as a visitor visa to visit my family, but my medication supply will only last for another 15-17 days. When I inquired at a drug store, the cost for one month's supply is $1,300, which I am unable to afford."],
    ["I'm having a very hard time lately and feel so alone. More than a year ago, I stopped taking my meds because I feel like no one cares and the meds are just delaying the inevitable. My family wants nothing to do with me. I don't have any friends I can count on."],
]

# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600&family=IBM+Plex+Sans:wght@400;600;700&display=swap');
* { font-family: 'IBM Plex Sans', sans-serif; }
#header { text-align:center; padding:1rem 0 0.25rem; }
#header h1 { font-size:1.5rem; font-weight:700; color:#0f172a; margin:0; letter-spacing:-0.02em; }
#header p  { color:#64748b; font-size:0.85rem; margin:0.25rem 0 0; }
.mono textarea { font-family:'IBM Plex Mono',monospace !important; font-size:0.82rem !important; }
.action-box textarea {
    font-size:1.05rem !important; font-weight:600 !important;
    text-align:center !important; border-radius:8px !important; min-height:48px !important;
}
.disclaimer { font-size:0.72rem; color:#cbd5e1; text-align:center; margin-top:0.5rem; }
"""

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="ISCORE Triage") as demo:

    gr.HTML("""
    <div id="header">
        <h1>ISCORE ART Adherence Triage</h1>
        <p>Barrier classification (Flan-T5-XL) + Risk level (Flan-T5-Large) · AIDS 2026 · THPEB043</p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=2):
            msg_in = gr.Textbox(
                label="Client message",
                placeholder="Paste a client message here — can be multiple sentences…",
                lines=5, max_lines=12,
            )
            with gr.Row():
                clear_btn  = gr.Button("Clear", variant="secondary", size="sm")
                submit_btn = gr.Button("Run triage →", variant="primary")

        with gr.Column(scale=1):
            action_out = gr.Textbox(
                label="Triage recommendation",
                interactive=False,
                elem_classes="action-box",
            )
            summary_out = gr.Textbox(
                label="Barrier summary (highest risk per barrier)",
                interactive=False,
                lines=6,
            )

    sentence_out = gr.Textbox(
        label="Per-sentence predictions",
        interactive=False,
        lines=6,
        elem_classes="mono",
    )

    gr.Examples(examples=EXAMPLES, inputs=msg_in, label="Example messages", examples_per_page=4)
    gr.HTML('<p class="disclaimer">Research demo · AIDS 2026 · THPEB043 · Not for clinical use</p>')

    submit_btn.click(fn=run_triage, inputs=msg_in, outputs=[sentence_out, summary_out, action_out])
    msg_in.submit(   fn=run_triage, inputs=msg_in, outputs=[sentence_out, summary_out, action_out])
    clear_btn.click( fn=lambda: ("", "", "", ""), outputs=[msg_in, sentence_out, summary_out, action_out])

demo.launch()