Using tf.data.Dataset with Keras on a TPU

Question

I am training a model with Keras which constitutes of a Huggingface RoBERTa model as a backbone with a downstream task of span prediction and binary prediction for text.