Bert Fine-Tuning on SQuAD with native TF2
I'm trying to run Bert fine tuning on SQuAD with Tensorflow's model.fit function.  I was able to run the tutorial on sequence classification from here https://huggingface.co/transformers/training.html#tensorflow, but a similar set up does not seem to work for question answering.  I am able to get it working with synthetic data, but not real data.  I can never seem to get the data in the right format.  Here is the script I am trying to use

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = TFBertForQuestionAnswering.from_pretrained('bert-base-uncased')

data = tfds.load('squad')
examples = SquadV1Processor().get_examples_from_dataset(data)
train_dataset = squad_convert_examples_to_features(examples, tokenizer, max_seq_length=328, doc_stride=128, max_query_length=32, is_training=True, return_dataset='tf')

optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer=optimizer, loss=loss)
model.fit(train_dataset, epochs=2, steps_per_epoch=1)

I've tried many other things, but this is the closest I got to something working.  The problem is that the dataset has keys that aren't recognized.  Specifically, I get this error

ValueError: Found unexpected keys that do not correspond to any Model output: dict_keys(['start_positions', 'end_positions', 'cls_index', 'p_mask', 'is_impossible']). Expected: ['output_1', 'output_2'] 

I might be able to find someway to manually remove these keys, but I feel like there should be some canonical way of getting this to work correctly.  I would appreciate any suggestions.  Thanks in advance.
Hello !

You have an example here https://colab.research.google.com/github/huggingface/datasets/blob/master/notebooks/Overview.ipynb<|||||>Perfect! Thank you, I did not see that.  I think I'm all good now.<|||||>I want to have a fine tuning on my own dataset,  but I don't know how to deal with the data format. Any suggestion will be appreciated. <|||||>@mug2mag this post is closed, please open another one with more detail on what you would like to do. Thanks.