Plume128k-onnx

Plume128k ONNX

ONNX export of projecte-aina/Plume128k, a 2B parameter translation model from Projecte AINA at the Barcelona Supercomputing Center (BSC). The model is Catalan-centric: it was trained only on Catalan-centric parallel data, and it covers 16 supervised directions plus 56 zero-shot directions.

Export task

Plume is a decoder-only model. Its model_type is gemma and its architecture is GemmaForCausalLM. It is a prompted translator, not an encoder-decoder sequence-to-sequence model.

The export therefore uses:

optimum-cli export onnx --model projecte-aina/Plume128k --task text-generation-with-past --no-post-process <outdir>

This is different from the dedicated MT models (Marian, M2M100, MADLAD, NLLB), which use text2text-generation-with-past and produce an encoder_model.onnx / decoder_model.onnx / decoder_with_past_model.onnx triple. A decoder-only export produces one model.onnx graph with KV-cache inputs and outputs.

Load it with ORTModelForCausalLM, not ORTModelForSeq2SeqLM.

scripts/export/quantize_causal_lm_int8.py from TigreGotico/linguonnx makes the int8 graph in int8/ from the fp32 graph:

python quantize_causal_lm_int8.py model.onnx int8 --weight-only --block-size 32 --accuracy-level 4

The script stores 8-bit weights in MatMulNBits nodes, in blocks of 32. It keeps the token embedding and the output projection (lm_head) in fp32. The int8 graph needs onnxruntime>=1.22.0. ONNX Runtime 1.18.1, 1.20.1 and 1.21.1 refuse 8-bit MatMulNBits when they load the graph.

Prompt template

This is the whole usability of the model. The template is:

<s> [SRC_LANG] SOURCE_TEXT \n[TGT_LANG]

Written as a Python format string:

prompt = "<s> [{}] {} \n[{}]".format(src_lang_code, sentence, tgt_lang_code)

Notes that matter:

• The literal <s> at the start is required. The tokenizer does not add a BOS token automatically. The string <s> maps to token id 0, which is the BOS token.
• [cat_Latn], [spa_Latn] and the other language tags are each a single token in the vocabulary. Write them exactly, with the square brackets.
• There is a space after <s>, a space before \n, and a real newline before the target tag.
• There is no chat template. Do not use apply_chat_template.

Language codes: cat_Latn, spa_Latn, eng_Latn, ita_Latn, eus_Latn, deu_Latn, por_Latn, glg_Latn, fra_Latn.

Usage

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForCausalLM

model_id = "TigreGotico/Plume128k-onnx"
tok = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForCausalLM.from_pretrained(model_id)          # fp32
# model = ORTModelForCausalLM.from_pretrained(
#     model_id, subfolder="int8", file_name="model_quantized.onnx")   # int8

def translate(sentence, src="eng_Latn", tgt="cat_Latn", max_new_tokens=64):
    prompt = "<s> [{}] {} \n[{}]".format(src, sentence, tgt)
    enc = tok(prompt, return_tensors="pt")
    enc.pop("token_type_ids", None)          # the tokenizer emits it; gemma does not accept it
    out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
                         num_beams=1, pad_token_id=tok.pad_token_id)
    return tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()

print(translate("The weather is very nice today in Barcelona."))
# El temps és molt agradable avui a Barcelona.

enc.pop("token_type_ids", None) is required. The tokenizer returns token_type_ids, and generate() raises ValueError: The following model_kwargs are not used by the model without it.

Files

Both graphs keep their weights in an external data file: model.onnx_data for fp32 and int8/model_quantized.onnx_data for int8. Each graph must be downloaded together with its data file. from_pretrained handles this for you. If you fetch files by hand, take model.onnx and model.onnx_data as a pair — model.onnx alone is only ~850 KB and holds no weights.

The int8 graph is named model_quantized.onnx, not model.onnx, on purpose. If both files were called model.onnx, from_pretrained(model_id) with no subfolder would find int8/model.onnx first and silently load the int8 weights when you asked for fp32. You must therefore pass file_name="model_quantized.onnx" together with subfolder="int8" to get int8.

Parity

Ten translation prompts over five directions (en→ca, es→ca, ca→en, ca→es, fr→ca, ca→de), greedy decoding (do_sample=False, num_beams=1, max_new_tokens=64), compared against the PyTorch original:

The int8 measurement uses 128 FLORES-200 devtest sentences: sentences 0 to 7 in each of the 16 supervised directions. Decoding is greedy (num_beams=1, max_new_tokens=128), through linguonnx on ONNX Runtime 1.31.1. On these 128 sentences, the fp32 graph gives the same tokens as PyTorch in all 128. chrF is the sacrebleu corpus chrF.

No int8 output repeats text until the token limit.

int8 is reported for information and is not a release gate.

Sample translations (fp32, greedy)

Licence and attribution

Apache 2.0, inherited from the base model.

The model was trained by Projecte AINA at the Barcelona Supercomputing Center (BSC). All credit for the model and its training data belongs to them. This repository only adds the ONNX export and the int8 quantization. See the original model card for training data, evaluation, and intended use.