Does anyone have a Google Colab notebook they can share? (If possible, one that has been adapted to the new model upload)
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.