import gradio as gr
from transformers import pipeline
from normalizer import cleaning

pipe = pipeline("fill-mask", model="HooshvareLab/albert-fa-zwnj-base-v2")


def greet(text):
    text = cleaning(text)
    results = pipe(text)
    result_str = ''

    for result in results:
        result_str += result['token_str'] + \
            ' (' + str(result['score']) + ')' + \
            '\n' + result['sequence'] + '\n\n'

    return result_str


demo = gr.Interface(fn=greet, inputs=gr.Textbox(label='input text'), outputs=gr.Textbox(label="output text:"),
                    allow_flagging='never',
                    examples=['آسمان [MASK] است.',
                              'تلفن [MASK] می‌خورد.',
                              'بچه‌ها مادر و پدر خود را [MASK] دارند.',
                              '[MASK] یک گیاه است.',
                              ])


demo.launch()