from transformers import pipeline
import gradio as gr

classifier = pipeline(
    "zero-shot-classification",
    model="Recognai/bert-base-spanish-wwm-cased-xnli"
)

def zero_shot_classification(text, labels):
    candidate_labels = [label.strip() for label in labels.split(",")]  # Convertir en lista
    result = classifier(text, candidate_labels)
    
    output = "\n".join(
        [f"{label}: {score:.2f}" for label, score in zip(result['labels'], result['scores'])]
    )
    return output

with gr.Blocks() as demo:
    gr.Markdown("""
    # Clasificación Zero-Shot en Español
    Esta aplicación permite clasificar un texto en diferentes categorías de tu elección. 
    Solo debes proporcionar el texto y las etiquetas que deseas evaluar, y el modelo asignará probabilidades a cada una de ellas.
    """)

    gr.Interface(
        fn=zero_shot_classification,
        inputs=[
            gr.Textbox(label="Texto a clasificar", placeholder="Escribe el texto aquí..."),
            gr.Textbox(label="Etiquetas de clasificación", placeholder="Ejemplo: cultura, sociedad, economia, salud, deportes")
        ],
        outputs=gr.Textbox(label="Resultados"),
    )

    gr.Markdown("""
    ---
    Demostración de clasificación Zero-Shot usando el modelo [Recognai/bert-base-spanish-wwm-cased-xnli](https://huggingface.co./Recognai/bert-base-spanish-wwm-cased-xnli).
    
    Desarrollado con ❤️ por [@srjosueaaron](https://www.instagram.com/srjosueaaron/).
    """)

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