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(""" """) 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('

Research demo · AIDS 2026 · THPEB043 · Not for clinical use

') 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()