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()