CodeTed's picture
Update app.py
86848fa verified
import gradio as gr
from t5.t5_model import T5Model
from transformers import AutoTokenizer, T5ForConditionalGeneration
#tokenizer = AutoTokenizer.from_pretrained("CodeTed/CGEDit")
#model = T5ForConditionalGeneration.from_pretrained("CodeTed/CGEDit")
model = T5Model('t5', "CodeTed/CGEDit", args={"eval_batch_size": 1}, cuda_device=-1, evaluate=True)
def cged_correction(sentence, function):
prompt = {"錯別字校正":"糾正句子中的錯字:", "文法校正":"糾正句子中的錯誤:",
"文本重構":"在不改動文意的情況下改寫句子:", "文本簡化":"在不改動文意的情況下改寫句子:", "整體校正":"修改句子的錯誤或使其更通順:"}
#input_ids = tokenizer(prompt[function] + sentence, return_tensors="pt").input_ids
for _ in range(3):
output = model.predict([prompt[function] + sentence + "_輸出句:"])
sentence = output[0]
#edited_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return output[0]
with gr.Blocks() as demo:
gr.Markdown(
"""
# Chinese Grammarly - 中文文本自動編輯器
### 貼上中文文章來使你的句子更順暢~
Start typing below to see the correction.
"""
)
funt = gr.Radio(["錯別字校正", "文法校正", "文本重構", "文本簡化", "整體校正"], label="Correction Type")
#設定輸入元件
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()