Hi Kay,
If you look into this tutorial: Solve GLUE tasks using BERT on TPU  |  Text  |  TensorFlow 47
You’ll get this exactly information to help you. To be more exact, in the make_bert_preprocess_model  function