import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# Tải mô hình và tokenizer
model_name = "VietAI/envit5-translation"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Chuyển mô hình sang GPU nếu có
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Hàm dịch ngôn ngữ với tùy chọn ngôn ngữ
def translate(text, language):
    prefix = "vi: " if language == "Tiếng Việt -> Tiếng Anh" else "en: "
    inputs = tokenizer(prefix + text, return_tensors="pt", padding=True).input_ids.to(device)
    outputs = model.generate(inputs, max_length=512)
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return translation

# Tạo giao diện với Gradio và thêm nút chọn ngôn ngữ
interface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Nhập văn bản cần dịch"),
        gr.Dropdown(choices=["Tiếng Việt -> Tiếng Anh", "Tiếng Anh -> Tiếng Việt"], label="Chọn ngôn ngữ")
    ],
    outputs="text",
    title="Ứng dụng dịch ngôn ngữ",
    description="Dịch từ tiếng Việt sang tiếng Anh và ngược lại."
)

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