import gradio as gr
from t5.t5_model import T5Model
from transformers import AutoTokenizer, T5ForConditionalGeneration
#tokenizer = AutoTokenizer.from_pretrained("CodeTed/traditional_CSC_t5")
#model = T5ForConditionalGeneration.from_pretrained("CodeTed/traditional_CSC_t5")
model = T5Model('t5', "CodeTed/Chinese_Spelling_Correction_T5", args={"eval_batch_size": 1}, cuda_device=-1, evaluate=True)

def cged_correction(sentence = '為了降低少子化，政府可以堆動獎勵生育的政策。'):
    for _ in range(3):
        outputs = model.predict(["糾正句子中的錯字：" + sentence + "_輸出句："])
        sentence = outputs[0]
    return outputs[0]

with gr.Blocks() as demo:
    gr.Markdown(
        """
    # 中文錯別字校正 - Chinese Spelling Correction
    ### Find Spelling Error and get the correction！
    Start typing below to see the correction.
    """
    )
    #設定輸入元件
    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, outputs=output)
demo.launch()