import gradio as gr
import pandas as pd
from transformers import pipeline
from transformers.pipelines.base import PipelineException

fill_mask = pipeline("fill-mask", model="Keyurjotaniya007/bert-large-cased-wikitext-mlm-3.0", device=-1)

def predict_mask(sentence: str, top_k: int):
    mask = fill_mask.tokenizer.mask_token

    sentence = sentence.replace("[MASK]", mask)

    if mask not in sentence:
        return pd.DataFrame(
            [["Error: please include `[MASK]` in your sentence.", 0.0]],
            columns=["Sequence", "Score"]
        )

    try:
        preds = fill_mask(sentence, top_k=top_k)
    except PipelineException as e:
        return pd.DataFrame([[f"Error: {str(e)}", 0.0]],
                            columns=["Sequence", "Score"])

    rows = [[p["sequence"], round(p["score"], 3)] for p in preds]
    return pd.DataFrame(rows, columns=["Sequence", "Score"])

with gr.Blocks(title="Masked Language Modeling") as demo:
    gr.Markdown(
        "# Masked Language Modeling\n"
        "Enter a sentence with one `[MASK]` token and see the top-K completions."
    )

    with gr.Row():
        sentence = gr.Textbox(
            lines=2,
            placeholder="e.g. The Great Wall of [MASK] is visible from space.",
            label="Input Sentence"
        )
        top_k = gr.Slider(
            minimum=1, maximum=10, step=1, value=5,
            label="K Predictions[Min=1 & Max=10]"
        )

    predict_btn = gr.Button("Evaluate [MASK] Words", variant="primary")

    results_df = gr.Dataframe(
        headers=["Sequence", "Score"],
        datatype=["str", "number"],
        wrap=True,
        interactive=False,
        label="Predictions"
    )

    predict_btn.click(
        fn=predict_mask,
        inputs=[sentence, top_k],
        outputs=results_df
    )

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