import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

import shutil, os
for d in ["~/.cache/huggingface", "~/.cache/torch"]:
    path = os.path.expanduser(d)
    if os.path.exists(path):
        shutil.rmtree(path, ignore_errors=True)


# Tải model và tokenizer từ Hugging Face
model_name = "ntviet/mbart-hre-viet1.2"
# model = MBartForConditionalGeneration.from_pretrained(model_name)
model = MBartForConditionalGeneration.from_pretrained(model_name, torch_dtype="auto", low_cpu_mem_usage=True)

tokenizer = MBart50TokenizerFast.from_pretrained(model_name)


# Hàm dịch từ Hre sang Việt
def translate_hre_to_viet(text):
    # Tokenize văn bản đầu vào
    inputs = tokenizer(text, return_tensors="pt")

    # Thực hiện dịch với forced_bos_token_id để chỉ định đầu ra là tiếng Việt
    generated_tokens = model.generate(
        **inputs,
        forced_bos_token_id=tokenizer.lang_code_to_id["vi_VN"]
    )

    # Giải mã kết quả
    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    return translation[0]

# Tạo giao diện Gradio
demo = gr.Interface(
    fn=translate_hre_to_viet,
    inputs=gr.Textbox(lines=5, placeholder="Nhập văn bản tiếng Hre...", label="Tiếng Hre"),
    outputs=gr.Textbox(label="Tiếng Việt"),
    title="Dịch máy Hre - Việt",
    description="Mô hình dịch từ tiếng Hre sang tiếng Việt sử dụng mBart",
    examples=[
        ["Chin triêuq chin hring chin jât chin ngan chin hring chin jât chinh"],
        ["Vaiq jât triêuq chin hring ngan"]
    ],
    allow_flagging="never"
)

# Khởi chạy ứng dụng
demo.launch()