import gradio as gr
from transformers import pipeline

examples = [
    'Alisher Navoiy – ulug‘ o‘zbek va boshqa turkiy xalqlarning <mask>, mutafakkiri va davlat arbobi bo‘lgan.',
    'Oʻzbekistonning poytaxti <mask> shahri boʻlib, davlat tili oʻzbek tili hisoblanadi.',
    'Oʻzbekiston iqtisodiyoti bozor <mask> bosqichma-bosqich oʻtadi, tashqi savdo siyosati import oʻrnini bosishga asoslangan.',
    'Kuchli yomg‘irlar tufayli bir qator <mask> kuchli sel oqishi kuzatildi.',
    'Registon maydoni - tarixda shaharning ilm-fan, siyosat va <mask> markazi boʻlgan.',
    'Venera - Quyosh tizimidagi o‘z o‘qi atrofida soat sohasi farqli ravishda aylanadigan yagona <mask>.'
]

models = [
    "sinonimayzer/UzRoBERTa-v1", 
    "tahrirchi/tahrirchi-bert-base", 
    "rifkat/uztext-3Gb-BPE-Roberta"
]

def df(arr):
    d = {}
    for val in arr:
        d[val['token_str']] = val['score']
    return d
    
def fn(text):
    result = ()
    for model in models:
        result = result + df(pipeline("fill-mask", model=model))
    return result

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            output0 = gr.Label(label=models[0])
            input = gr.Textbox(label="Input", lines=10, value=examples[0])
        with gr.Column():
            output1 = gr.Label(label=models[1])
            output2 = gr.Label(label=models[2])
    btn = gr.Button("Check")
    btn.click(fn, inputs=[input], outputs=[output0, output1, output2])
    gr.Examples(examples, fn=fn, inputs=[input], outputs=[output0, output1, output2], cache_examples=True, batch=True)
    
if __name__ == "__main__":
    demo.queue().launch()