Hi, not exactly what you requested (ie. not the latest Huggingface version),
but it should give the main idea with a bonus on working with TPU:
Kaggle notebook on implementation of MLM finetuning on TPUs with a custom training loop in TensorFlow 2.
https://www.kaggle.com/riblidezso/finetune-xlm-roberta-on-jigsaw-test-data-with-mlm 11