import os
from transformers import pipeline
import gradio as gr


hf_writer = gr.HuggingFaceDatasetSaver(os.getenv('HF_API_TOKEN'), "suicide-comment-es-space-human-feedback")
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="manual",
    flagging_options=['Correct Classification', 'Incorrect Classification'],
    flagging_callback=hf_writer,
    examples=[
        "Quiero acabar con todo. No merece la pena vivir.",
        "El partido de fútbol fue igualado, todos disfrutamos mucho jugando juntos."
    ],
    article="""
## Motivación

Introduce cualquier frase en español, un modelo de inteligencia artificial tratará de predecir si la frase contiene algún patrón indicativo de ideación o comportamiento suicida.

## Feedback

¡Ayúdanos a mejorar la aplicación marcando la respuesta como Correcta o Incorrecta según corresponda!

Nosotros usaremos este feedback para ampliar el dataset que se usará para fine-tunear el modelo en versiones posteriores.

¡Gracias!

## Dataset

https://huggingface.co/datasets/hackathon-somos-nlp-2023/suicide-comments-es

## Model

https://huggingface.co/hackathon-somos-nlp-2023/roberta-base-bne-finetuned-suicide-es

## Team

* [dariolopez](https://huggingface.co/dariolopez)
* [diegogd](https://huggingface.co/diegogd)

""",

    
)
iface.launch()