north_2_cen / app.py
minjibi's picture
Update app.py
bae9ab4
import gradio as gr
import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from transformers import (
MT5ForConditionalGeneration,
MT5TokenizerFast,
)
model = MT5ForConditionalGeneration.from_pretrained(
"minjibi/north_to_cen",
return_dict=True,
)
tokenizer = MT5TokenizerFast.from_pretrained(
"minjibi/north_to_cen"
)
def generate_text(input_text):
input_ids = tokenizer.encode(input_text, return_tensors="pt")
output = model.generate(input_ids, max_new_tokens = len(input_text.split(' ')) * 2, num_beams = 2, early_stopping=True, length_penalty = -5.0)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
return generated_text
examples = [
["ห้าม เตียว ไป เตียว มา หรือ อู้ กั๋น ขณะ ตี้ ทำ วัตร สวด มนต์ นั่ง สมาธิ"],
["แม่น แม่น แล้ว ตาง พิพิธภัณฑ์ เปิ้น เปิด หื้อ เข้า ไป ผ่อ ต๋อน กี่ โมง ก๋า"],
["ตาง วัด บ่ อนุญาต หื้อ เอา ไฟแช็ก ของ วัด ปิ๊ก ไป บ้าน เน่อ ครับ"],
["กิ๋น ข้าว กับ หยัง"]
]
iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", examples=examples)
iface.launch()