import gradio as gr
from transformers import pipeline

# Cargar tu modelo
classifier = pipeline("text-classification", model="marcostrfn/bart-base-spanish-nli-taller-test")

def predict_nli(premise, hypothesis):
    """
    Predice la relación de inferencia entre premisa e hipótesis
    """
    # Formato para NLI: "premisa [SEP] hipótesis"
    text = f"{premise} [SEP] {hypothesis}"
    
    results = classifier(text)
    
    # Formatear resultados como diccionario para Gradio
    return {result["label"]: result["score"] for result in results}

# Crear interfaz Gradio
demo = gr.Interface(
    fn=predict_nli,
    inputs=[
        gr.Textbox(label="Premisa", placeholder="Una persona monta un caballo"),
        gr.Textbox(label="Hipótesis", placeholder="Una persona está sobre un animal")
    ],
    outputs=gr.Label(label="Predicción", num_top_classes=3),
    title="BART Spanish NLI - Inferencia en Lenguaje Natural",
    description="Modelo BART fine-tuned para Natural Language Inference en español. Determina si la hipótesis es consecuencia (entailment), contradice (contradiction) o es neutral respecto a la premisa.",
    examples=[
        ["Un hombre está tocando la guitarra", "Un músico está tocando un instrumento"],
        ["El gato está durmiendo en el sofá", "El perro está corriendo en el jardín"],
        ["María estudia filosofía en la universidad", "María es estudiante"]
    ]
)

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