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


def cged_correction(sentence = '為了降低少子化，政府可以堆動獎勵生育的政策。'):
    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(
        """
    # Flip Text!
    Start typing below to see the output.
    """
    )
    #設定輸入元件
    sent = gr.Textbox(label="Sentence")
    # 設定輸出元件
    output = gr.Textbox(label="Result")
    #設定按鈕
    greet_btn = gr.Button("Correction")
    #設定按鈕點選事件
    greet_btn.click(fn=cged_correction, inputs=sent, outputs=output)
demo.launch()