import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Cargar el modelo y el tokenizador
model_name = "dagomem/modelo_tweets_2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Mapeo de etiquetas
label_map = {
    0: "enfado",
    1: "alegría",
    2: "optimismo",
    3: "tristeza"
}

# Ejemplos predefinidos
examples = [
    ["@user Interesting choice of words... Are you confirming that governments fund #terrorism? Bit of an open door, but still...", "enfado"],
    ["love to see them interacting 😸🙏 #itsbeensolong #laughing", "alegría"],
    ["Life is too short to hide your feelings. Don't be afraid to say what you feel.", "optimismo"],
    ["#GameOfThones how can you top that next week #heartbreaking", "tristeza"]
]


# Función de predicción
def predict_tags(text):
    # Tokenizar el texto de entrada
    encoded_input = tokenizer.encode_plus(
        text,
        padding="longest",
        truncation=True,
        return_tensors="pt"
    )

    # Realizar la predicción
    output = model(**encoded_input)

    # Obtener las etiquetas predichas
    predicted_label = output.logits.argmax(dim=1)

    return label_map[predicted_label.item()]

# Interfaz de usuario
gr.Interface(
    fn=predict_tags,
    inputs=gr.inputs.Textbox(lines=3, label="Texto"),
    outputs="text",
    title="Análisis de emociones de tweets (en inglés)",
    description="Escribe un tweet en inglés y el modelo te dirá qué emoción transmite.",
    examples=examples
    ).launch(share=False)