import gradio as gr
from transformers import pipeline

def mask(text):
    mask_model = pipeline("fill-mask", model="google-bert/bert-base-uncased")
    output = mask_model(text)
    return output[0]['sequence']


# Gradio UI
examples=[['Today I went to [MASK] after I got out of bed.']]
ui = gr.Interface(
    fn=mask,
    inputs="textbox",
    outputs="textbox",
    submit_btn="Predict",
    title="BERT Test - Zircon.tech",
    description="Enter some text with [MASK] to let BERT guess the missing word!",
    examples=examples
    )
ui.launch(debug=True)