File size: 1,484 Bytes
a0d8498 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | #!/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)}")
|