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.',
    '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>.',
    'Kuchli yomg‘irlar tufayli bir qator <mask> kuchli sel oqishi kuzatildi.',
    'Oʻzbekiston iqtisodiyoti bozor <mask> bosqichma-bosqich oʻtadi, tashqi savdo siyosati import oʻrnini bosishga asoslangan.'
]

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

def df(arr):
    d = {}
    for val in arr:
        d[val['token_str']] = val['score']
    return d
    
def fn(text):
    arr = []
    for model in models:
        arr.append(df(pipeline("fill-mask", model=model)(text)))
    return arr[0], arr[1], arr[2], arr[3], arr[4]
    
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            output0 = gr.Label(label=models[0])
        with gr.Column():
            output1 = gr.Label(label=models[1])
        with gr.Column():
            output2 = gr.Label(label=models[2])
    with gr.Row():
        with gr.Column():
            output3 = gr.Label(label=models[3])
        with gr.Column():
            output4 = gr.Label(label=models[4])
        with gr.Column():
            input = gr.Textbox(label="Input", value=examples[0], lines=8, max_lines=8)
            btn = gr.Button("Check")
    gr.Examples(examples, fn=fn, inputs=[input], outputs=[output0, output1, output2, output3, output4], cache_examples=True, batch=True)
    btn.click(fn, inputs=[input], outputs=[output0, output1, output2, output3, output4])
if __name__ == "__main__":
    demo.queue().launch()