import gradio as gr
from transformers import pipeline

examples = [
    'Kuchli yomg‘irlar tufayli bir qator <mask> kuchli sel oqishi kuzatildi.',
    'Shu munosabat bilan O‘zbekiston Prezidenti global inqiroz sharoitida savdo-iqtisodiy hamkorlikni <mask> va hududlararo aloqalarni rivojlantirishning muhim masalalariga to‘xtalib o‘tdi.',
]

def df(arr):
    d = {}
    for val in arr:
        d[val['token_str']] = val['score']
    return d
    
def fn(text):
    uzroberta_v1 = pipeline("fill-mask", model="sinonimayzer/UzRoBERTa-v1")
    roberta_1_8_v2 = pipeline("fill-mask", model="sinonimayzer/roberta-1.8-v2")
    return df(uzroberta_v1(text)), df(roberta_1_8_v2(text))

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(label="Input")
        with gr.Column():
            output1 = gr.Label(label="UzRoBERTa-v1")
            output2 = gr.Label(label="roberta-1.8-v2")
    
    btn = gr.Button("Generate")
    btn.click(fn, inputs=[inp], outputs=[output1, output2])
    gr.Examples(examples, inputs=[inp])

if __name__ == "__main__":
    demo.queue().launch(share=True)