app.py

import gradio as gr
import torch
import sentencepiece as spm
from transformers import T5ForConditionalGeneration
from huggingface_hub import hf_hub_download

MODEL_ID = "BrandenTung/t5-zh-en-opus"

# Load model
model = T5ForConditionalGeneration.from_pretrained(MODEL_ID)
model.eval()

# Download tokenizer from model repo
spm_path = hf_hub_download(
    repo_id=MODEL_ID,
    filename="spm.model"
)

sp = spm.SentencePieceProcessor()
sp.load(spm_path)

PREFIX = "translate Chinese to English: "

def translate(text):
    text = (text or "").strip()
    if not text:
        return ""

    ids = sp.encode(PREFIX + text, out_type=int)
    input_ids = torch.tensor([ids], dtype=torch.long)

    with torch.no_grad():
        outputs = model.generate(input_ids, max_length=128, num_beams=5)

    return sp.decode(outputs[0].tolist())


iface = gr.Interface(
    fn=translate,
    inputs="text",
    outputs="text",
    title="ZH → EN Translator",
    description="T5-based Chinese to English translation model trained on OPUS-100"
)

iface.launch()