import gradio as gr
from transformers import pipeline
import numpy as np
import torch

BASE_MODEL = "muhtasham/TajBERTo"
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'.join(k)
    return final_string 
         
gradio_ui = gr.Interface(
    fn=mask_fill,
    title="Predicting masked words in Tajik Language 🇹🇯",
    description="Enter a a sentence to predict the masked word",
    inputs=[
        gr.inputs.Textbox(lines=3),
    ],
    outputs=[
        gr.outputs.Textbox(label="Answer"),
    ],
    examples=[
        ["Пойтахти <mask> Душанбе"],
    ],
    enable_queue=True,
    allow_screenshot=False,
    allow_flagging=False,
)
gradio_ui.launch(debug=True)