from fastai.learner import load_learner
import gradio as gr
from huggingface_hub import from_pretrained_fastai

# Lista de etiquetas de emociones en español
emotion_labels = [
    'admiración', 'diversión', 'ira', 'molestia', 'aprobación', 'cuidado',
    'confusión', 'curiosidad', 'deseo', 'decepción', 'desaprobación',
    'disgusto', 'vergüenza', 'emoción', 'miedo', 'gratitud', 'duelo',
    'alegría', 'amor', 'nerviosismo', 'optimismo', 'orgullo', 'comprensión',
    'alivio', 'remordimiento', 'tristeza', 'sorpresa', 'neutral'
]

# Carga del modelo desde Hugging Face
learn = from_pretrained_fastai("diribes/go_emotions")

# Clasificación con umbral personalizado
def clasificar_texto(texto, umbral=0.2):
    _, _, probs = learn.predict(texto)
    emociones_detectadas = [
        f"{emotion_labels[i]} ({probs[i]:.2f})"
        for i in range(len(probs))
        if probs[i] > umbral
    ]
    return ", ".join(emociones_detectadas) if emociones_detectadas else "Sin emoción detectada"

# Interfaz Gradio
iface = gr.Interface(
    fn=clasificar_texto,
    inputs=gr.Textbox(lines=5, label="Introduce el texto"),
    outputs=gr.Label(label="Emoción(es) detectada(s)"),
    title="Clasificador de Emociones",
    description="Modelo de clasificación de emociones fine-tuned con FastAI y desplegado en Hugging Face Spaces."
)

if __name__ == "__main__":
    iface.launch()