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 results


demo = gr.Interface(fn=greet, inputs=gr.Textbox(label='input text'), outputs=gr.Textbox(label="output text:"),
                    allow_flagging='never')


demo.launch()