import gradio
from transformers import pipeline

unmasker = pipeline('fill-mask', model='bert-large-uncased-whole-word-masking')


def predict_words(text):
    results = unmasker(text)
    word_list = [f"{result['token_str']}: {result['score']:.2%}" for result in results]
    return word_list[:5]


inputs = [
    gradio.components.Textbox(label="Enter Text (Use [MASK] for the masked word; Use it only once)",
                              placeholder="Enter your text here."),
]

outputs = [
    gradio.components.Textbox(label="Prediction 1"),
    gradio.components.Textbox(label="Prediction 2"),
    gradio.components.Textbox(label="Prediction 3"),
    gradio.components.Textbox(label="Prediction 4"),
    gradio.components.Textbox(label="Prediction 5")
]

title = "Fill Mask (bert-large-uncased-whole-word-masking)"

gradio.Interface(fn=predict_words, inputs=inputs, outputs=outputs, title=title, allow_flagging="never",
                 css="footer{display:none !important}").launch()