import gradio as gr
from transformers import MarianMTModel, MarianTokenizer

MODEL_ID = "vncgabriel/es2zh-mbart"

# carrega tokenizer e modelo Marian
tokenizer = MarianTokenizer.from_pretrained(MODEL_ID)
model     = MarianMTModel.from_pretrained(MODEL_ID)

def traduzir_es2zh(texto, max_length=100):
    # tokeniza o input
    batch = tokenizer(texto, return_tensors="pt")
    # gera a tradução
    gen = model.generate(
        **batch,
        max_length=max_length,
        num_beams=4,
        early_stopping=True
    )
    # decodifica e remove tokens especiais
    return tokenizer.decode(gen[0], skip_special_tokens=True)

iface = gr.Interface(
    fn=traduzir_es2zh,
    inputs=[
        gr.Textbox(lines=3, label="Texto en Español"),
        gr.Slider(20, 300, value=100, step=10, label="Tamaño máximo")
    ],
    outputs=gr.Textbox(label="Traducciõn Chino"),
    title="ES → ZH com es2zh-mbart",
)

if __name__ == "__main__":
    iface.launch()