import gradio as gr

from optimum.onnxruntime import ORTModelForSeq2SeqLM
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("icon-it-tdtu/mt-vi-en-optimum")
model = ORTModelForSeq2SeqLM.from_pretrained("icon-it-tdtu/mt-vi-en-optimum")

def translate(text):
    inputs = tokenizer(text, return_tensors='pt')
    outputs = model.generate(**inputs)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

iface = gr.Interface(
    fn=translate,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter Vietnamese text to translate..."),
    outputs="text",
    theme="huggingface"
)
iface.launch()