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.
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 onemodel.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:
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.
1from transformers import AutoTokenizer
2from optimum.onnxruntime import ORTModelForCausalLM
34model_id ="TigreGotico/Plume128k-onnx"5tok = AutoTokenizer.from_pretrained(model_id)6model = ORTModelForCausalLM.from_pretrained(model_id)# fp327# model = ORTModelForCausalLM.from_pretrained(8# model_id, subfolder="int8", file_name="model_quantized.onnx") # int8910deftranslate(sentence, src="eng_Latn", tgt="cat_Latn", max_new_tokens=64):11 prompt ="<s> [{}] {} \n[{}]".format(src, sentence, tgt)12 enc = tok(prompt, return_tensors="pt")13 enc.pop("token_type_ids",None)# the tokenizer emits it; gemma does not accept it14 out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,15 num_beams=1, pad_token_id=tok.pad_token_id)16return tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()1718print(translate("The weather is very nice today in Barcelona."))19# 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.
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:
export
exact match vs PyTorch
fp32
10/10 (100%)
int8 dynamic
3/10 (30%)
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)
direction
source
output
en→ca
The weather is very nice today in Barcelona.
El temps és molt agradable avui a Barcelona.
en→ca
She opened the window and looked at the sea.
Va obrir la finestra i va mirar el mar.
es→ca
Ayer se fue, tomo sus cosas y se puso a navegar.
Ahir se'n va anar, va agafar les seves coses i es va posar a navegar.
es→ca
El tren llega a las ocho de la manana.
El tren arriba a les vuit del matí.
ca→en
Els nens juguen al parc cada tarda.
Children play in the park every afternoon.
ca→en
M'agrada molt llegir llibres d'historia.
I really like reading history books.
ca→es
Demà anirem a la platja amb els meus amics.
Mañana vamos a la playa con mis amigos.
ca→es
La biblioteca tanca a les set del vespre.
La biblioteca cierra a las siete de la tarde.
fr→ca
Le chat dort sur le canapé du salon.
El gat dorm al sofà de la sala.
ca→de
Aquesta ciutat té una història molt llarga.
Diese Stadt hat eine lange Geschichte.
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.
