import spaces

import gradio as gr
import torch
from transformers import pipeline


pipe = pipeline(
    "text-classification",
    model="SHK4K/suicide-roberta",
    device=0 if torch.cuda.is_available() else -1,
)


@spaces.GPU
def is_safe(text):
    if not text or not text.strip():
        return {
            "safe": 0.0,
            "unsafe": 0.0,
        }

    results = pipe(
        text,
        truncation=True,
        max_length=512,
        top_k=2,
    )

    if isinstance(results[0], list):
        results = results[0]

    return {
        "Safe" if r["label"].lower() == 'label_0' else 'Unsafe': r["score"]
        for r in results
    }

demo = gr.Interface(
    fn=is_safe,

    inputs=gr.Textbox(
        lines=6,
        label="Text",
        placeholder="Enter text to classify...",
    ),

    outputs=gr.Label(
        num_top_classes=2,
        label="Prediction",
    ),

    title="Suicide Risk Detector",

    description=(
        "A research NLP model that detects potential "
        "suicide-risk signals in text. "
        "This is not a medical diagnosis."
    ),
)

demo.launch()