from transformers import pipeline
import gradio as gr

classifier = pipeline("text-classification", model="joortif/clasificador-climate-claim")

label_map = {
    "LABEL_0": "Afirma el cambio climático (SUPPORTS)",
    "LABEL_1": "Niega el cambio climático (REFUTES)",
    "LABEL_2": "No hay suficiente información (NOT_ENOUGH_INFO)",
    "LABEL_3": "Afirmación disputada (DISPUTED)"
}

def classify_text(text):
    results = classifier(text)
    return {label_map[res["label"]]: res["score"] for res in results}

demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=4, label="Escribe un texto sobre el clima"),
    outputs=gr.Label(num_top_classes=4),
    title="Clasificador de Afirmaciones Climáticas",
)

demo.launch()