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

# Cargar tu modelo desde HuggingFace
MODEL_NAME = "antonn-dromundo/InclusioCheck-BETO-HateSpeech"

print("🔄 Cargando modelo...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Mover a CPU (los Spaces gratuitos usan CPU)
device = torch.device("cpu")
model.to(device)
model.eval()

print("✅ Modelo cargado exitosamente")

# Etiquetas
LABELS = {
    0: "✅ Lenguaje Inclusivo",
    1: "⚠️ Discurso de Odio/Excluyente"
}

def analizar_texto(texto):
    """
    Analiza un texto y devuelve si es inclusivo o contiene discurso de odio
    """
    if not texto or len(texto.strip()) == 0:
        return {
            "⚠️ Error": "Por favor ingresa un texto para analizar"
        }
    
    # Tokenizar
    inputs = tokenizer(
        texto,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True
    )
    
    # Mover inputs a CPU
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Predicción
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=-1)
        predicted_class = torch.argmax(probs, dim=-1).item()
        confidence = probs[0][predicted_class].item()
    
    # Preparar resultado
    label = LABELS[predicted_class]
    
    # Crear diccionario de resultados para Gradio
    resultado = {
        LABELS[0]: float(probs[0][0]),
        LABELS[1]: float(probs[0][1])
    }
    
    # Mensaje adicional
    if predicted_class == 1:
        mensaje = f"⚠️ Este texto contiene lenguaje potencialmente excluyente u ofensivo (confianza: {confidence:.1%})"
    else:
        mensaje = f"✅ Este texto parece inclusivo y respetuoso (confianza: {confidence:.1%})"
    
    return resultado, mensaje

# Ejemplos para demostración
ejemplos = [
    ["Me encanta trabajar en equipo con personas diversas"],
    ["Todos los inmigrantes son delincuentes"],
    ["La diversidad es nuestra mayor fortaleza"],
    ["Manejas bien para ser mujer"],
    ["Respeto y valoro las opiniones de todas las personas"]
]

# Crear interfaz
demo = gr.Interface(
    fn=analizar_texto,
    inputs=gr.Textbox(
        label="Escribe el texto a analizar",
        placeholder="Ejemplo: Me gusta trabajar con personas de diferentes culturas...",
        lines=3
    ),
    outputs=[
        gr.Label(label="Análisis del Texto", num_top_classes=2),
        gr.Textbox(label="Resultado", lines=2)
    ],
    title="🔍 SinOdio - Detector de Lenguaje Inclusivo",
    description="""
    **Analiza textos en español para detectar lenguaje excluyente o discurso de odio.**
    
    Este modelo está basado en BETO (BERT en español) y fue entrenado para identificar
    contenido que pueda ser ofensivo, discriminatorio o excluyente.
    
    ⚡ Powered by transformers y entrenado con dataset de hate speech en español.
    """,
    examples=ejemplos,
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

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