from transformers import pipeline
import gradio as gr

fill_mask = pipeline("fill-mask", model="dccuchile/bert-base-spanish-wwm-cased")

def complete_mask(sentence):
    if "[MASK]" not in sentence:
        return "⚠️ Por favor, incluye '[MASK]' en la oración."
    
    results = fill_mask(sentence)

    output = "\n".join(
        [f"🔹 **{res['token_str']}** (Probabilidad: {res['score']:.2%})" for res in results 
         if res['token_str'].strip() and res['token_str'] not in ["[UNK]", "."]]
    )
    
    return output if output else "⚠️ No se encontraron predicciones válidas."

with gr.Blocks() as demo:
    gr.Markdown("""
    # Modelo Fill-Mask en Español
    Esta aplicación permite completar una oración con la palabra que falta, representada como '[MASK]'. 
    El modelo sugerirá las mejores opciones de palabras para completar la oración, con sus respectivas probabilidades.
    """)

    gr.Interface(
        fn=complete_mask,
        inputs=gr.Textbox(label="Oración con [MASK]", placeholder="Ejemplo: La capital de España es [MASK]."),
        outputs=gr.Markdown(label="Predicciones"),
    )

    gr.Markdown("""
    ---
    Demostración de completado de palabras(máscaras) usando el modelo [dccuchile/bert-base-spanish-wwm-cased](https://huggingface.co/dccuchile/bert-base-spanish-wwm-cased).
    
    Desarrollado con ❤️ por [@srjosueaaron](https://www.instagram.com/srjosueaaron/).
    """)

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