import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

model_dir = "crrodrvi/beto_lectura_facil_clasif"

# Configuración del modelo y carga del tokenizer
config = AutoConfig.from_pretrained(model_dir, num_labels=2, finetuning_task="text-classification")
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=config)

# Etiquetas personalizadas
id2label = {0: "Difícil", 1: "Fácil"}

def inference(input_text):
    inputs = tokenizer(
        input_text,
        max_length=512,
        truncation=True,
        padding="max_length",
        return_tensors="pt",
    )
    
    with torch.no_grad():
        logits = model(**inputs).logits
    
    predicted_class_id = logits.argmax().item()
    output = id2label[predicted_class_id]
    return output

# Crear la interfaz de Gradio
demo = gr.Interface(
    fn=inference,
    inputs=gr.Textbox(label="Input Text", placeholder="Introduce el texto aquí...", lines=10),
    outputs=gr.Textbox(label="Output Label"),
    examples=[
        ["En París se hace un desfile militar con aviones celebrando ese día. Además en las embajadas francesas de todo el mundo se realizan actos.", 1],
        ["Según Clean Cities, las ZBE son ampliamente utilizadas en otros países de Europa con el fin de mejorar la calidad del aire, salvaguardar la salud pública y abordar la crisis climática. Además, estas medidas fomentan la adopción de transportes limpios, activos y sostenibles, apoyan la transición hacia energías más limpias, reducen la contaminación sonora y contribuyen a la reorganización y naturalización de los entornos urbanos.", 0],
    ],
    title="BETO Clasificación lectura fácil-difícil",
)

# Lanzar la aplicación
demo.launch(debug=True)