from transformers import pipeline
from transformers import MT5Tokenizer
from transformers import MT5ForConditionalGeneration
import streamlit as st

model_name = "jamie613/mt5_correct_puntuation"
model = MT5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = MT5Tokenizer.from_pretrained(model_name)

pipe = pipeline("text2text-generation", model = model, tokenizer = tokenizer, max_length = 300)
text = st.text_area('請輸入要訂正標點符號的中文文句。最長300字。')

if text:
  out = pipe(text)
  st.write(out)