app.py

"""
app.py — Undertone
The signal beneath the surface.
Mental Health Triage Assistant · Gradio interface for HuggingFace Spaces
"""

import torch
import gradio as gr
from pathlib import Path
from transformers import (
    DistilBertTokenizerFast,
    DistilBertForSequenceClassification,
)

MODEL_DIR = Path(__file__).parent / "model" / "saved"

LABEL2ID = {"informational": 0, "emotional_support": 1, "escalation": 2}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}

LABEL_CONFIG = {
    "informational": {
        "emoji": "🟢",
        "title": "Informational",
        "response": (
            "This message appears to be asking for information or resources. "
            "We can help point you to reliable mental health information and services."
        ),
        "action": "Provide relevant resources and psychoeducation.",
    },
    "emotional_support": {
        "emoji": "🟡",
        "title": "Emotional Support Needed",
        "response": (
            "It sounds like you're going through a difficult time. "
            "You're not alone — a counselor is here to listen and support you."
        ),
        "action": "Connect with a counselor for empathetic listening and coping support.",
    },
    "escalation": {
        "emoji": "🔴",
        "title": "Immediate Attention Required",
        "response": (
            "Your message suggests you may be in crisis. "
            "Please reach out immediately — help is available right now."
        ),
        "action": "Escalate to crisis counselor. Provide crisis hotline numbers immediately.",
        "crisis_resources": [
            "🆘 988 Suicide & Crisis Lifeline — Call or text **988**",
            "💬 Crisis Text Line — Text HOME to **741741**",
            "🌐 International resources — https://www.iasp.info/resources/Crisis_Centres/",
        ],
    },
}

# ── 高危 emoji → 文字映射 ─────────────────────────────────
# DistilBERT tokenizer 把 emoji 变成 [UNK]，预处理转成文字保留语义
EMOJI_RISK_MAP = {
    # 明确危机信号
    "🔪": " knife weapon harm ",
    "💊": " pills medication overdose ",
    "🔫": " gun weapon ",
    "🪓": " weapon harm ",
    "⚰️": " death coffin ",
    "🪦": " grave death ",
    "☠️": " death skull ",
    "💀": " death dying ",
    "🩸": " blood injury ",
    # 强烈情绪信号
    "😭": " crying devastated ",
    "😢": " crying sad ",
    "🆘": " emergency help urgent ",
    "😶": " numb empty ",
    "💔": " heartbroken pain ",
    "🥺": " desperate pleading ",
    "😫": " exhausted overwhelmed ",
    "😩": " hopeless drained ",
    "🌊": " drowning overwhelmed ",
    "🖤": " darkness despair ",
    "😔": " sad down ",
    "😞": " disappointed hopeless ",
    "🙃": " ironic resigned ",
}

def preprocess(text: str) -> str:
    """把高危 emoji 替换为对模型有意义的文字"""
    for emoji_char, replacement in EMOJI_RISK_MAP.items():
        text = text.replace(emoji_char, replacement)
    return " ".join(text.split())  # 清理多余空格


DEVICE = (
    "mps"  if torch.backends.mps.is_available() else
    "cuda" if torch.cuda.is_available() else
    "cpu"
)

# ── 模型加载 ──────────────────────────────────────────────
print(f"Loading model from {MODEL_DIR} on {DEVICE}...")
tokenizer = DistilBertTokenizerFast.from_pretrained(str(MODEL_DIR))
model     = DistilBertForSequenceClassification.from_pretrained(str(MODEL_DIR)).to(DEVICE)
model.eval()
print("Model loaded.")


@torch.no_grad()
def triage(message: str):
    if not message or not message.strip():
        return (
            "**Please enter a message to analyze.**",
            {},
            "—",
            "—",
        )

    # emoji 预处理
    processed = preprocess(message.strip())

    enc = tokenizer(
        [processed],
        truncation=True,
        padding=True,
        max_length=128,
        return_tensors="pt",
    )
    logits = model(
        input_ids=enc["input_ids"].to(DEVICE),
        attention_mask=enc["attention_mask"].to(DEVICE),
    ).logits
    probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]

    label_id   = int(probs.argmax())
    label_name = ID2LABEL[label_id]
    config     = LABEL_CONFIG[label_name]
    confidence = float(probs[label_id])

    prob_dict = {
        LABEL_CONFIG[ID2LABEL[i]]["title"]: float(probs[i])
        for i in range(3)
    }

    classification_md = (
        f"## {config['emoji']} {config['title']}\n\n"
        f"**Confidence:** {confidence:.1%}\n\n"
        f"---\n\n"
        f"{config['response']}"
    )

    crisis_md = ""
    if label_name == "escalation":
        crisis_md = "### 🆘 Crisis Resources\n\n" + "\n\n".join(config["crisis_resources"])
    else:
        crisis_md = "_No immediate crisis resources needed._"

    return (
        classification_md,
        prob_dict,
        config["action"],
        crisis_md,
    )


# ── 示例消息 ─────────────────────────────────────────────
EXAMPLES = [
    ["What are the signs of anxiety disorder?"],
    ["I've been feeling really low lately and don't know what to do 😔"],
    ["idk i just want everything to stop ngl"],
    ["i fel so aloen and its getting worse"],
    ["😭🔪 cant do this anymore"],
    ["lol whatever doesn't matter anyway 💀"],
    ["Where can I find free therapy in my area?"],
    ["I've been giving away my stuff to friends lately"],
]

CUSTOM_CSS = """
:root {
    --bg:      #0F1117;
    --surface: #1A1D27;
    --border:  #2A2D3A;
    --text:    #E8EAF0;
    --muted:   #6B7280;
}
body, .gradio-container {
    background: var(--bg) !important;
    color: var(--text) !important;
    font-family: 'Inter', system-ui, sans-serif !important;
}
.app-header {
    text-align: center;
    padding: 2rem 1rem 1rem;
    border-bottom: 1px solid var(--border);
    margin-bottom: 1.5rem;
}
.app-header h1 { font-size: 1.8rem; font-weight: 700; color: #fff; margin: 0; }
.app-header p  { color: var(--muted); font-size: 0.9rem; margin-top: 0.4rem; }
.badge {
    display: inline-block; padding: 2px 10px; border-radius: 99px;
    font-size: 0.7rem; font-weight: 600; text-transform: uppercase;
    background: #2A2D3A; color: #9CA3AF; margin-top: 0.5rem;
}
.disclaimer {
    background: #1E1B2E; border: 1px solid #3B2F6B; border-radius: 8px;
    padding: 0.75rem 1rem; font-size: 0.8rem; color: #9CA3AF; margin-bottom: 1.5rem;
}
.label-legend {
    display: flex; gap: 1.5rem; justify-content: center;
    padding: 0.75rem; border-top: 1px solid var(--border);
    margin-top: 1.5rem; font-size: 0.8rem; color: var(--muted);
}
.dot { display: inline-block; width: 8px; height: 8px; border-radius: 50%; margin-right: 5px; }
"""

# ── 界面 ─────────────────────────────────────────────────
with gr.Blocks(css=CUSTOM_CSS, title="Undertone") as demo:

    gr.HTML("""
    <div class="app-header">
        <h1>🧠 Undertone</h1>
        <p>The signal beneath the surface — mental health message triage</p>
        <span class="badge">Research Demo · DistilBERT Fine-Tuned</span>
    </div>
    """)

    gr.HTML("""
    <div class="disclaimer">
        ⚠️ <strong>Disclaimer:</strong> This is a research prototype for educational purposes only.
        It does not provide medical advice and should never replace professional mental health support.
        If you or someone you know is in crisis, please contact a crisis service immediately.
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            msg_input = gr.Textbox(
                label="Message",
                placeholder="Type a support message here...",
                lines=4,
                max_lines=8,
            )
            submit_btn = gr.Button("Analyze Message", variant="primary")
            gr.Examples(examples=EXAMPLES, inputs=msg_input, label="Try an example")

        with gr.Column(scale=1):
            classification_out = gr.Markdown(label="Classification")
            confidence_out     = gr.Label(label="Confidence Scores", num_top_classes=3)

    with gr.Row():
        with gr.Column():
            action_out = gr.Textbox(label="Recommended Action", interactive=False)
        with gr.Column():
            crisis_out = gr.Markdown(label="Crisis Resources")

    gr.HTML("""
    <div class="label-legend">
        <span><span class="dot" style="background:#2ECC71"></span>Informational</span>
        <span><span class="dot" style="background:#F39C12"></span>Emotional Support</span>
        <span><span class="dot" style="background:#E74C3C"></span>Escalation</span>
    </div>
    """)

    submit_btn.click(
        fn=triage,
        inputs=msg_input,
        outputs=[classification_out, confidence_out, action_out, crisis_out],
    )
    msg_input.submit(
        fn=triage,
        inputs=msg_input,
        outputs=[classification_out, confidence_out, action_out, crisis_out],
    )

if __name__ == "__main__":
    demo.launch()