Adding new train_step logic to make things less confusing for users
Author: Rocketknight1
Pull request: yes

Change to `train_step` I discussed earlier with @gante, making the new TF approaches less confusing for users.

This PR is **extremely** experimental right now, and since it touches `train_step` it could break the entire TF side of the library, so please don't merge it until I'm done!

Comments:
_The documentation is not available anymore as the PR was closed or merged._

This should now be ready for review! The main changes are:

1) `model.fit()` with the internal loss is now tested correctly
2) The internal loss will now work no matter where you pass your labels, which should reduce user confusion by about 50%
3) Attempting to use Keras metrics with the internal loss will now throw an informative error, instead of just creating a huge mess like it used to

I see some failing tests which are possibly caused by an outdated version of TF in the Docker image that doesn't understand `steps_per_execution`. Will investigate tomorrow!

Tests are passing and I think we're ready for final review!

I have a couple of things I'm still unsure about, though:

@gante: I removed the `expand_1d` call in `train_step`. Keras by default expands 1D input arrays to 2D, e.g. shape (8,) becomes (8, 1). This is not what our models expect at all and the test revealed some failures, that were fixed by removing it. Can you think of anything that change might break, though?

@sgugger / @LysandreJik: Is there any way to get a list of all possible 'label' column names for our models? This probably seems like a strange question, but the reason I need it is that if the user passes a single tensor as the labels to Keras, I need to make that tensor a key in the input dict, which means figuring out which argument name is the 'label'. Right now I just hardcoded a list of common ones, and then I inspect the function signature to see which of those is present, but it would be a lot better if I had some kind of definitive list.

We don't have a definitive list and maintaining it is going to be difficult as well with new modalities arriving. Perhaps we should add an attribute `label_names` to all models to have this info stored somewhere? The `Trainer` would also benefit from it on the PyTorch side.

Wdyt @LysandreJik ?

I wouldn't mind having the models know their label names, otherwise inspecting the signature and looking for any argument that contains `label` should give the correct information as well. If doing that during inference/training, it might slow down the forward pass so if you feel like adding a `label_names` is cleaner that's fine for me.

The problem is that some models have labels without labels in their names (QA models, I'm looking at you!)

For TF models trained with Keras, the `call` (`forward`) method is usually only run once to build the graph, so expensive Python calls like `inspect` are totally okay!

That compilation step is really the source of a lot of different design decisions between the two frameworks - e.g. TF `einsum` takes a long time to figure out the optimal contraction during compilation and then saves it, so you can get very different runtimes as a result.

> The problem is that some models have labels without labels in their names (QA models, I'm looking at you!)

For consistency's sake, shouldn't we rename them to have `label` in their name? (with appropriate deprecation cycle of course)
Maybe a non-issue if no users have ever been misled, but if it makes everything clearer *and* helps programmatic handling of labels, then it might be worth it

I don't think making a change on the labels for the QA models is warranted as we can make an exception for those (double check the name "QuestionAnswering" and that it has `start_positions` and `end_positions` in its signature for instance). We can do the inspection once at init, so it's not a big deal in terms of performance.

> @gante: I removed the expand_1d call in train_step. Keras by default expands 1D input arrays to 2D, e.g. shape (8,) becomes (8, 1). This is not what our models expect at all and the test revealed some failures, that were fixed by removing it. Can you think of anything that change might break, though?

I can't think of any problem. It is also being tested, so if there are problems, we should be able to catch them quickly.

For reference, here are all of the arguments for any `forward()` method in the codebase on a subclass of `PreTrainedModel` that contain "label" in their name:

`{'sentence_order_label', 'mc_labels', 'next_sentence_label', 'obj_labels', 'mask_labels', 'sentence_image_labels', 'class_labels', 'labels', 'entity_labels', 'matched_label'}`

@sgugger @LysandreJik None of them look especially wrong to me, so I'm happy to assume this list + {"start_positions", "end_positions"} covers all of the possible labels and just do an `inspect` call if you think that's a good solution!

Added a utility in #16526 

Quick update: This PR is mostly ready, but I'm going to wait for `find_labels` to be merged, then rebase and use it here.

Rebased with `find_labels` and looks good in testing, so I'm going to merge!