from transformers import XLMRobertaForSequenceClassification, XLMRobertaTokenizer
import torch
import gradio as gr

# Nombre del modelo en Hugging Face
model_path = "natmarinn/sentimientos-bullrich"  

# Cargar el modelo y el tokenizador desde Hugging Face
model = XLMRobertaForSequenceClassification.from_pretrained(model_path)
tokenizer = XLMRobertaTokenizer.from_pretrained(model_path)

# Definir las etiquetas de las clases
clases = ["Negativo", "Neutro", "Positivo"]

def clasificar_texto(texto):
    # Tokenizar el texto de entrada
    inputs = tokenizer(texto, return_tensors="pt", truncation=True)
    
    # Realizar la predicción y obtener las probabilidades
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1).squeeze().tolist()  # Convertir a probabilidades
    
    # Crear una lista de categorías y probabilidades
    categorias_prob = list(zip(clases, probs))
    
    # Ordenar por probabilidad de mayor a menor
    categorias_prob = sorted(categorias_prob, key=lambda x: x[1], reverse=True)
    
    # Formatear el resultado en un string para mostrarlo en Gradio
    resultado = "\n".join([f"{categoria}: {probabilidad:.2%}" for categoria, probabilidad in categorias_prob])
    return resultado

# Crear la interfaz en Gradio
iface = gr.Interface(
    fn=clasificar_texto,                # Función que ejecuta la clasificación
    inputs="text",                      # Tipo de entrada: texto
    outputs="text",                     # Tipo de salida: texto
    title="Clasificación de Sentimientos - Patricia Bullrich - Debate Presidencial",
    description="Escribe un comentario sobre el debate presidencial y el modelo clasificará el sentimiento con sus probabilidades.",
    examples=[["Vamos pato"], ["Montonera"], ["Jardin de infantes"]]
)

# Ejecutar la aplicación
iface.launch(share=True)