import gradio as gr
from transformers import pipeline

fill_mask = pipeline("fill-mask", model="SamJoshua/EsperBERTo")


def predict(text: str):
    if text.count("<mask>") != 1:
        raise gr.Error("Please include exactly one <mask> token in the text.")
    return {r["token_str"].strip(): r["score"] for r in fill_mask(text, top_k=5)}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Esperanto text with one <mask>", value="La suno <mask>."),
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    examples=[["La suno <mask>."], ["Jen la komenco de bela <mask>."]],
    title="EsperBERTo fill-mask",
    cache_examples=False,
)

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