import gradio as gr


from transformers import AutoModelForSequenceClassification
import torch

from transformers import AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("somosnlp/bertin_base_climate_detection_spa_v2")

model = AutoModelForSequenceClassification.from_pretrained("somosnlp/bertin_base_climate_detection_spa_v2")

id2label = {0: "NEGATIVE", 1: "POSITIVE"}
label2id = {"NEGATIVE": 0, "POSITIVE": 1}

def inference_fun(Texto):
    inputs = tokenizer(Texto, return_tensors="pt", truncation=True)
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class_id = logits.argmax().item()
    output_tag = model.config.id2label[predicted_class_id]
    output_logit = logits
    return output_tag, output_logit

examples = [
    [
        "Los gases refrigerantes fluorados son potentes gases de efecto invernadero. Las alternativas, como el amoniaco o el dióxido de carbono capturado, pueden sustituirlos con el tiempo."
    ],
    [
        "En la química de coordinación, un ligando es un ión o molécula (o grupo funcional) que se une a un átomo metálico central para formar un complejo de coordinación. El enlace entre el metal y el ligando generalmente involucra la donación de uno o más pares de electrones del ligando. La naturaleza del enlace metal-ligando oscila entre el enlace covalente y el enlace iónico.Además, el orden de enlace metal-ligando puede variar de uno a tres. Los ligandos se consideran bases de Lewis, aunque se conocen casos raros de ""ligandos"" ácidos de Lewis.."
    ],
    [
        "El uso excesivo de fertilizantes nitrogenados -un fenómeno frecuente en la agricultura- da lugar a la producción de óxido nitroso, un potente gas de efecto invernadero. Un uso más juicioso de los fertilizantes puede frenar estas emisiones y reducir la producción de fertilizantes, que consume mucha energía."
    ]
]

description = "En la siguiente interfaz puedes introducir un texto cualquiera, luego puedes dar click en submit y devolverá una etiqueta de POSITIVO: para un texto que aborde temas de cambio climatico y NEGATIVO: para un texto en general que no hable de estos temas"

iface = gr.Interface(fn=inference_fun, inputs="text", outputs=["label", "text"], examples=examples, description=description)
iface.launch()