from transformers import pipeline
import gradio as gr

# Initialize the pipeline once
mask_filler = pipeline(
    "fill-mask", 
    model="kaiku03/MLM-distilbert-base-uncased-finetuned-game_titles_accelerate"
)

def fill_mask(prompt):
    output = mask_filler(prompt)
    
    # Extract the top 3 sequences and their scores
    top_sequences = {i['sequence']: i['score'] for i in output[:10]}

    # Return the dictionary
    return top_sequences


# gradio app
# Define example inputs and outputs
examples = [
    "league of [MASK].",
    "[MASK] Raider",
    "Grand [MASK] Auto ",
]

iface = gr.Interface(
    fn=fill_mask,
    inputs='text',
    outputs=gr.Label(num_top_classes=4, label="Predictions"),
    examples=[
        [ex] for ex in examples
    ],
    title='Game Titles Prediction',
    description='This demo performs auto-completion of game titles based on surrounding text.',
    article='All done by Kaiku, documentation at https://medium.com/p/b09b45dd7ce8/edit',

)

iface.launch()