import gradio as gr
import torch
from transformers import pipeline

model_ckpt = "HenryAI/KerasBERTv1"
fill_mask = pipeline("fill-mask", model=model_ckpt)  


def  cloze_task(text):
    preds = fill_mask(text)
    return preds[0]["sequence"]
    
    [{'sequence': 'from tensorflow.keras import layers', 'score': 0.9447882771492004, 'token': 455, 'token_str': ' layers'}, {'sequence': 'from tensorflow.keras import optimizers', 'score': 0.010178953409194946, 'token': 9110, 'token_str': ' optimizers'}, {'sequence': 'from tensorflow.keras import regularizers', 'score': 0.008282472379505634, 'token': 14453, 'token_str': ' regularizers'}, {'sequence': 'from tensorflow.keras import losses', 'score': 0.004894345533102751, 'token': 3114, 'token_str': ' losses'}, {'sequence': 'from tensorflow.keras import Model', 'score': 0.003724579466506839, 'token': 2663, 'token_str': ' Model'}]
     
description="Let's see if BERT has learnt how to write Keras!"
title="Keras BERT v1",
interface = gr.Interface(fn= cloze_task,
    title= title, description = description,
    inputs=[gr.inputs.Textbox(lines=3)],
    outputs=[gr.outputs.Textbox(label="Answer"),],
    examples=[["from tensorflow.keras import <mask>"],],
    enable_queue=True
)
interface.launch(debug=True)