from transformers import pipeline
import gradio as gr


pipe = pipeline("text-classification", model='hackathon-somos-nlp-2023/roberta-base-bne-finetuned-suicide-es')


def detect_suicide_comment(input):
    return pipe(input)


iface = gr.Interface(
    fn=detect_suicide_comment, inputs="text", outputs="text", allow_flagging="auto"
)
iface.launch()