File size: 1,119 Bytes
aab2a91 2ca7d5d aab2a91 2ca7d5d aab2a91 2ca7d5d aab2a91 2ca7d5d aab2a91 4b2a98e bbf40ae 4b2a98e bbf40ae 4b2a98e 2ca7d5d 4b2a98e aab2a91 4b2a98e 2ca7d5d 4b2a98e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | 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() |