belumind/nllb-trilingual-en-vi-ja-onnx

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM

model_id = "sotalab/nllb-trilingual-en-vi-ja-onnx"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForSeq2SeqLM.from_pretrained(
    model_id,
    encoder_file_name="encoder_model_quantized.onnx",
    decoder_file_name="decoder_model_quantized.onnx",
    decoder_with_past_file_name="decoder_with_past_model_quantized.onnx",
)

def translate(text, src_lang, tgt_lang):
    tokenizer.src_lang = src_lang
    inputs = tokenizer(text, return_tensors="pt", max_length=256, truncation=True)
    tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
    
    outputs = model.generate(
        **inputs,
        forced_bos_token_id=tgt_lang_id,
        max_new_tokens=256,
        num_beams=1,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# English to Vietnamese
print(translate("Hello, how are you?", "eng_Latn", "vie_Latn"))
# Output: Chào, bạn khỏe không?

# English to Japanese
print(translate("Hello, how are you?", "eng_Latn", "jpn_Jpan"))
# Output: こんにちは、お元気ですか？

# Vietnamese to Japanese
print(translate("Tôi thích học tiếng Nhật.", "vie_Latn", "jpn_Jpan"))
# Output: 私は日本語を学ぶのが好きです。

def translate_batch(texts, src_lang, tgt_lang):
    tokenizer.src_lang = src_lang
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=256)
    tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang)
    
    outputs = model.generate(**inputs, forced_bos_token_id=tgt_lang_id, max_new_tokens=256, num_beams=1)
    return [tokenizer.decode(out, skip_special_tokens=True) for out in outputs]

texts = ["Good morning", "How are you?", "Thank you very much"]
results = translate_batch(texts, "eng_Latn", "vie_Latn")
for text, result in zip(texts, results):
    print(f"{text} → {result}")

class TrilingualTranslator:
    LANG_CODES = {"en": "eng_Latn", "vi": "vie_Latn", "ja": "jpn_Jpan"}
    
    def __init__(self, model_id="sotalab/nllb-trilingual-en-vi-ja-onnx"):
        from transformers import AutoTokenizer
        from optimum.onnxruntime import ORTModelForSeq2SeqLM
        
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = ORTModelForSeq2SeqLM.from_pretrained(
            model_id,
            encoder_file_name="encoder_model_quantized.onnx",
            decoder_file_name="decoder_model_quantized.onnx",
            decoder_with_past_file_name="decoder_with_past_model_quantized.onnx",
        )
    
    def translate(self, text, src="en", tgt="vi"):
        src_code = self.LANG_CODES[src]
        tgt_code = self.LANG_CODES[tgt]
        
        self.tokenizer.src_lang = src_code
        inputs = self.tokenizer(text, return_tensors="pt", max_length=256, truncation=True)
        tgt_lang_id = self.tokenizer.convert_tokens_to_ids(tgt_code)
        
        outputs = self.model.generate(**inputs, forced_bos_token_id=tgt_lang_id, max_new_tokens=256, num_beams=1)
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

# Usage
translator = TrilingualTranslator()
print(translator.translate("Hello world", "en", "vi"))
print(translator.translate("Hello world", "en", "ja"))

@misc{nllb-trilingual-2024,
  author = {SotaLab},
  title = {NLLB Trilingual Translation Model (EN-VI-JA) - INT8 ONNX},
  year = {2024},
  publisher = {Hugging Face},
  url = {https://huggingface.co/sotalab/nllb-trilingual-en-vi-ja-onnx}
}