import gradio as gr
from transformers import pipeline

# Load model (bert-base-cased)
unmasker = pipeline(
    "fill-mask",
    model="bert-base-cased"
)

# Mask token for bert-base-cased
MASK_TOKEN = "[MASK]"

def fill_mask(text, top_k):
    try:
        if MASK_TOKEN not in text:
            return f"⚠️ Please include the mask token: {MASK_TOKEN}"

        results = unmasker(text, top_k=top_k)

        output = []
        for r in results:
            output.append(f"{r['sequence']} (Score: {round(r['score'], 4)})")

        return "\n".join(output)

    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# 🧠 Mask Filling with BERT")
    gr.Markdown(f"Use `{MASK_TOKEN}` to let the model predict missing words.")

    text_input = gr.Textbox(
        label="Enter your sentence",
        value="This course will teach you all about [MASK] models."
    )

    top_k_input = gr.Slider(
        minimum=1,
        maximum=10,
        value=2,
        step=1,
        label="Number of predictions (top_k)"
    )

    output = gr.Textbox(label="Predictions")

    submit_btn = gr.Button("Run")

    submit_btn.click(
        fn=fill_mask,
        inputs=[text_input, top_k_input],
        outputs=output
    )

# Launch
if __name__ == "__main__":
    app.launch()