#!/usr/bin/env python3
"""Portable Python demo using the universal artifacts (tokenizer.json +
lang_tokens.json + ONNX), with no SMaLL-100-specific tokenizer class — the same
recipe every platform follows. optimum runs the greedy decode.

    pip install optimum[onnxruntime] tokenizers torch
    python examples/python/translate.py
"""
import json
import os

import torch
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from tokenizers import Tokenizer

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
tok = Tokenizer.from_file(os.path.join(ROOT, "tokenizer.json"))
LT = json.load(open(os.path.join(ROOT, "lang_tokens.json")))["lang_to_id"]
model = ORTModelForSeq2SeqLM.from_pretrained(
    ROOT, subfolder="onnx", use_merged=True, use_io_binding=False)


def translate(text: str, tgt: str) -> str:
    ids = [LT[tgt]] + tok.encode(text).ids          # tokenizer already appends </s>
    input_ids = torch.tensor([ids])
    out = model.generate(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
        num_beams=1, max_length=128,
    )
    return tok.decode(out[0].tolist(), skip_special_tokens=True)


if __name__ == "__main__":
    for text, tgt in [
        ("你好，请问最近的地铁站怎么走？", "en"),
        ("Excuse me, where can I find a pharmacy?", "zh"),
        ("この電車は空港に行きますか？", "ko"),
    ]:
        print(f"[{tgt}] {text}  ->  {translate(text, tgt)}")