Skip to main content

Custom CometKiwi training

This is a historical custom-estimator training example, distinct from native COMET/KIWI inference. The YAMLs contain commented choices for each training stage; they are not ready-to-run end-to-end fine-tuning configs. Select training.train_from, freezing, step count, and LoRA settings as described below, and supply the external dataset. Start with native COMET for pretrained scoring. Full custom training was not rerun during this documentation audit.

NOTE

This is NOT the exact replication of the Unbabel Comet Kiwi.

What is common: We use the same base model xlm-roberta-xl or xlm-roberta-xxl encoders We use the same dataset (1720-da.mlqe)

You can get the dataset to train the model below here: https://huggingface.co/eole-nlp/cometkiwi-xxl-eole/

What is different: wmt23-cometkiwi-da-xl(or xxl) use Layerwise attention which brings complexity without significant better accuracy we use Gelu instead of Tanh in the Estimator

Our scores for XL and XXL are in the same range when they are very different for Unbabel/wmt23-cometkiwi-XL or XXL

To make your life easier, run these commands from the recipe directory (here recipes/scoring/comet_custom).​

Retrieve and convert model​

Set environment variables​

export EOLE_MODEL_DIR=<where_to_store_models>

Download and convert the base model​

eole convert HF --model_dir facebook/xlm-roberta-xxl --output $EOLE_MODEL_DIR/xlm-roberta-xxl-eole

NOTE Pass --dtype fp16 explicitly if you want FP16 converted weights.

XXL is a 10.7G params model hence will save a 21.4GB file on disk (safetensors format) XL is a 3.5GB params model hence will save a 7.0GB file on disk

After conversion in $EOLE_MODEL_DIR you will get the following files: config.json model.00.safetensors sentencepiece.bpe.model vocab.json vocab.txt

The current HF converter accounts for tokenizer-added tokens. Do not blindly append added_vocab.txt to a newly converted vocabulary; compare its size with the saved embedding rows before changing artifacts.

Training cometkiwi​

Training will happen in two steps (see cometkiwi-xxl-eole.yaml file) FYI the trained model can be downloaded here: https://huggingface.co/eole-nlp/cometkiwi-xxl-eole/

STEP 1: We train from the converted xlm-roberta-xxl model but we keep everything frozen. We just add an Estimator layer that will be trained during 4000 steps To make sure we do not lose anything we will rename the pre-trained subfolder into step_4000_pretrain In this step we do NOT use LoRA but we use 4bit quant to make things easier and fit in a smaller cards

STEP 2: We release the encoder to make weights trainable hence we need to use LoRA since the model is big and gradients would not fit in consumer grade cards. We train during 40000 steps

For the two steps above, the training command is the same but make sure the yaml file is modified according to the instructions eole train --config cometkiwi-xxl-eole.yaml

After this step we need to merge the LoRA weights into the original model with the following command:

eole model lora --base_model "./cometkiwi-xxl-eole/step_4000_pretrain/" --lora_weights "./cometkiwi-xxl-eole/" --output "./cometkiwi-xxl-eole/merged"

Cometkiwi Inference​

Format the source / target file you want to score in the Comet format:

./combine.sh newstest2014.en newstest2014.de newstest2014.comb

Score the .comb file:

eole predict --config cometkiwi-xxl-inference.yaml --src newstest2014.comb --output newstest2014.scores --with_score

For now the scores are in the third column, so you can cut -f3 the output file