import os

import gradio as gr
from transformers import pipeline

RACISM_MODEL = "davidmasip/racism"
racism_analysis_pipe = pipeline(
    "text-classification",
    model=RACISM_MODEL,
    tokenizer=RACISM_MODEL,
    use_auth_token=os.getenv("access_token")
)


def racism_analysis(text):
    results = racism_analysis_pipe(text)[0]
    label = "Non-racist" if results["label"] == "LABEL_0" else "Racist"
    score = (
        1 - round(results["score"], 5)
        if results["label"] == "LABEL_0"
        else round(results["score"], 5)
    )
    return label, score


gradio_ui = gr.Interface(
    fn=racism_analysis,
    title="Racism Detector (Spanish)",
    description="Enter some text and check if model detects racism.",
    inputs=[
        gr.inputs.Textbox(lines=5, label="Paste some text here"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Label"),
        gr.outputs.Textbox(label="Racism score"),
    ],
    examples=[
        ["Unos menas roban a una mujer"],
        ["Unos chinos roban a una mujer"],
        ["Unos árabes roban a una mujer"],
        ["Unos moros roban a una mujer"],
        ["Unos panchitos roban a una mujer"],
        ["El gobierno levanta el confinamiento"]
    ],
)

gradio_ui.launch()