Fine-tune MLM in Roberta custom loss (additional component)

I want to fine-tune Roberta on MLM task on my own data.
However, for each word, I also have an additional vector with 10 elements.
So whenever I predict a masked token, I want the loss to be:
aMLM loss + bvector_prediction_loss.

How can I do it? didnt find any example or tutorial?

Any ideas?

You can simply overwrite the Trainer loss: Specify Loss for Trainer / TrainingArguments - #2 by nielsr

@nielsr Let’s say my labels has two parts: label1 and label2.
What will be the best way to pass label1 and label2 to compute_loss ?
I add another columns to the dataset, such that now my columns are:

    features: ['input_ids', 'attention_mask', 'labels', 'labels2'],
    num_rows: 563

But after I init the trainer, if I put a breakpoint in the 1st row of compute_loss(), I see that there is no ‘labels2’ column in inputs