import gradio as gr
from transformers import pipeline

BASE_MODEL = "nlpaueb/legal-bert-base-uncased"

mask_filler = pipeline("fill-mask", model=BASE_MODEL)  

def  mask_fill(text):
    k = []
    preds = mask_filler(text) 
    for pred in preds:
        k.append(pred["sequence"])
    final_string = '\n\n'.join(k)
    return final_string 
      
gradio_ui = gr.Interface(
    fn=mask_fill,
    title="Predicting masked words in legal text",
    description="Enter a a sentence to predict the masked word",
    inputs=[
        gr.inputs.Textbox(lines=3),
    ],
    outputs=[
        gr.outputs.Textbox(label="Answer"),
    ],
    examples=[
        ["The applicant submitted that her husband was subjected to treatment amounting to [MASK] whilst in the   custody of police."],
    ],
    enable_queue=True,
    allow_screenshot=False,
    allow_flagging=False,
)
gradio_ui.launch(debug=True)