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):
    prompt = {"錯別字校正":"糾正句子中的錯字：", "文法校正":"糾正句子中的錯誤：", 
    "文本重構":"在不改動文意的情況下改寫句子：", "文本簡化":"在不改動文意的情況下改寫句子：", "整體校正":"修改句子的錯誤或使其更通順："}
    input_ids = tokenizer(prompt[function] + 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 Grammarly - 中文文本自動編輯器
    ### 貼上中文文章來使你的句子更順暢～
    Start typing below to see the correction.
    """
    )
    funt = gr.Radio(["錯別字校正", "文法校正", "文本重構", "文本簡化", "整體校正"])
    #設定輸入元件
    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()