import gradio as gr
from transformers import MT5Tokenizer, MT5ForConditionalGeneration

title = "Dịch ngôn ngữ bản địa"
modelMuVi = 'mr4/trans-mu-vi'
modelViMu = 'mr4/trans-vi-mu'
examplesMu = [
    'Ủn cha̒w tưở da̭.',
    'Mẩi zơ̒ hơ̒i za?',
    'Hwa tḙp chôổng ma̭ng. Hwa zói chôổng pổ.'
]
examplesVi = [
    'Em chào anh ạ.',
    'Mấy giờ rồi?',
    'Hoa đẹp giống mẹ, Hoa giỏi giống bô.'
]
cache = 'cache'
default_language = 'Tiếng Mường -> Tiếng Việt phổ thông'
languages = [
    default_language,
    'Tiếng Việt phổ thông -> Tiếng Mường',
]

def translate(text, language):
    match language:
        case 0:
            model_name = modelMuVi
        case 1:
            model_name = modelViMu
        case _:
            return ['', 'Không hỗ trợ ngôn ngữ hiện tại!']
    tokenizer = MT5Tokenizer.from_pretrained(model_name, cache_dir=cache)
    model = MT5ForConditionalGeneration.from_pretrained(model_name, cache_dir=cache)
    inputs = tokenizer(text, return_tensors="pt", max_length=512, padding=True, truncation=True)
    translated = model.generate(**inputs)
    return [tokenizer.decode(translated[0], skip_special_tokens=True), model_name]

with gr.Blocks() as demo:
    with gr.Row():
        gr.HTML(
            value=f"<h1 style='text-align: center; font-weight: bold;'>{title}</h1>")
    with gr.Row():
        source_language_dropdown = gr.Dropdown(elem_id="mr4llt-language", choices=languages, value=default_language, label="Chọn ngôn ngữ cần dịch", type='index')
    with gr.Row():
        with gr.Column():
            input_textbox = gr.Text(elem_id="mr4llt-input", placeholder="Nhập nội dung cần dịch", label="Nội dung")
        with gr.Column():
            translated_textbox = gr.Text(elem_id="mr4llt-output", placeholder="", label="Bản dịch")
    btn = gr.Button("Dịch")
    lbl = gr.HTML("")
    btn.click(translate, inputs=[input_textbox, source_language_dropdown], outputs=[translated_textbox, lbl])
    gr.Examples(examples=examplesMu, inputs=input_textbox, label="Ví dụ tiếng Mường")
    gr.Examples(examples=examplesVi, inputs=input_textbox, label="Ví dụ tiếng Việt phổ thông")
    with gr.Row():
        gr.HTML(value="""
<b>GHI CHÚ:</b><br />
Sản phẩm đang trong quá trình thử nghiệm chỉ có thể dịch được các <u>câu ngắn</u> và thời gian dịch <u>sẽ lâu</u>.<br />
AI hiện tại có mức độ hiểu biết chỉ dừng ở mức tương đương một đứa <i>trẻ mầm non</i>.<br />
Do dữ liệu được sử dụng trong quá trình đào tạo còn ít, nên đôi khi sẽ xuất hiện hiện tượng <u>dịch nhưng lại thành phiên âm</u> vì thế mọi người chú ý kiểm tra lại kết quả.
        """)
if __name__ == "__main__":
    demo.launch()