import gradio as gr
from transformers import AutoTokenizer, T5ForConditionalGeneration
tokenizer = AutoTokenizer.from_pretrained("CodeTed/CGEDit")
model = T5ForConditionalGeneration.from_pretrained("CodeTed/CGEDit")


def cged_correction(sentence, function):
    input_ids = tokenizer('糾正句子裡的錯字：' + sentence, return_tensors="pt").input_ids
    outputs = model.generate(input_ids, max_length=200)
    edited_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return edited_text

with gr.Blocks() as demo:
    gr.Markdown(
        """
    # 中文錯別字校正 - Chinese Spelling Correction
    ### Find Spelling Error and get the correction！
    Start typing below to see the correction.
    """
    )
    funt = gr.Radio(["add", "subtract", "multiply", "divide"])
    #設定輸入元件
    sent = gr.Textbox(label="Sentence", placeholder="input the sentence")
    # 設定輸出元件
    output = gr.Textbox(label="Result", placeholder="correction")
    #設定按鈕
    greet_btn = gr.Button("Correction")
    #設定按鈕點選事件
    greet_btn.click(fn=cged_correction, inputs=[sent, funt], outputs=output)
demo.launch()