TigreGotico/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.

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 the fp32 and the int8 graphs are over the 2 GB protobuf limit, so each keeps its weights in an external data file model.onnx_data. Each `model.onnx` must be downloaded together with the `model.onnx_data` beside it. 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:

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

Use the fp32 export unless you must have the smaller file. int8 dynamic quantization hurts this 128k-vocabulary model much more than it hurts the 32k one. Most int8 outputs stay usable but drift ("La meteorologia és molt agradable" instead of "El temps és molt agradable", dropped capitalisation), and one prompt degenerates into a repetition loop: es→ca on "Ayer se fue, tomo sus cosas y se puso a navegar." returns "L'ha deixat, ha deixat, ha deixat, ha deixat, ha deixat." against a correct fp32 output. The larger embedding and output projection are the likely cause. If you need a small Plume, prefer TigreGotico/Plume32k-onnx int8, which holds up far better.

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.