How to build a custom question-answering head?
Author: muhammed-shihebi
Pull request: no

Using the ```TFBertForQuestionAnswering.from_pretrained()``` function, we get a predefined head on top of BERT together with a loss function that are suitable for this task. 

My question is how to create a custom head without relying on ```TFAutoModelForQuestionAnswering.from_pretrained()```. 

I want to do this because there is no place where the architecture of the head is explained clearly. By reading the code [here][1] we can see the architecture they are using, but I can't be sure I understand their code 100%. 

Starting from https://stackoverflow.com/questions/69025750/how-to-fine-tune-huggingface-bert-model-for-text-classificationis is good. However, it covers only the classification task, which is much simpler. 

```'start_positions'``` and ```'end_positions'``` are created following [this][2] tutorial.

So far, I've got the following:

```
train_dataset
# Dataset({
#     features: ['input_ids', 'token_type_ids', 'attention_mask', 'start_positions', 'end_positions'],
#     num_rows: 99205
# })
train_dataset.set_format(type='tensorflow', columns=['input_ids', 'token_type_ids', 'attention_mask'])
features = {x: train_dataset[x] for x in ['input_ids', 'token_type_ids', 'attention_mask']}
labels = [train_dataset[x] for x in ['start_positions', 'end_positions']]
labels = np.array(labels).T
tfdataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(16)

input_ids = tf.keras.layers.Input(shape=(256,), dtype=tf.int32, name='input_ids')
token_type_ids = tf.keras.layers.Input(shape=(256,), dtype=tf.int32, name='token_type_ids')
attention_mask = tf.keras.layers.Input((256,), dtype=tf.int32, name='attention_mask')


bert = TFAutoModel.from_pretrained("bert-base-multilingual-cased")
output = bert([input_ids, token_type_ids, attention_mask]).last_hidden_state
output = tf.keras.layers.Dense(2, name="qa_outputs")(output)
model = tf.keras.models.Model(inputs=[input_ids, token_type_ids, attention_mask], outputs=output)


num_train_epochs = 3
num_train_steps = len(tfdataset) * num_train_epochs
optimizer, schedule = create_optimizer(
   init_lr=2e-5,
   num_warmup_steps=0,
   num_train_steps=num_train_steps,
   weight_decay_rate=0.01
)

def qa_loss(labels, logits):
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction=tf.keras.losses.Reduction.NONE
    )
    start_loss = loss_fn(labels[0], logits[0])
    end_loss = loss_fn(labels[1], logits[1])
    return (start_loss + end_loss) / 2.0


model.compile(
    loss=loss_fn,
    optimizer=optimizer
)

model.fit(tfdataset, epochs=num_train_epochs)
```
And I am getting the following error: 
```
ValueError: `labels.shape` must equal `logits.shape` except for the last dimension. Received: labels.shape=(2,) and logits.shape=(256, 2)
```

It is complaining about the shape of the labels. This should not happen since I am using ```SparseCategoricalCrossentropy``` loss. 


  [1]: https://github.com/huggingface/transformers/blob/198c335d219a5eb4d3f124fdd1ce1a9cd9f78a9b/src/transformers/models/bert/modeling_tf_bert.py#L2065
  [2]: https://huggingface.co/course/chapter7/7?fw=pt

Comments:
Hello, thanks for opening an issue! We try to keep the github issues for bugs/feature requests.
Could you ask your question on the [forum](https://discuss.huggingface.co) instead?

Thanks!

Got it!