Rewrite TensorFlow train_step and test_step

[Rocketknight1]

Draft PR for a full rewrite of the TF train/test steps. I swear this will fix like 50% of our TF issues in one PR.

Current status:

• Correctly handles output mapping across most model classes for losses + metrics
• Keras metrics are back, even with the dummy loss. (!!!!)
• Keras metrics work correctly even for multi-output models (like QA)
• In most cases, users can pass tensors in either the input dict or the labels and the model will handle them correctly.
• No more errors when calling fit() when the model has nested output structure (e.g. the model outputting a past tuple)

What's left to do:

• Models with multiple unusual outputs that do not match label names may still have issues with metrics. This is relatively uncommon. We support adding a property to those classes to tell Keras what to do with the labels, but we haven't added it to any models yet. (None are failing in tests, so hopefully we won't need to worry too much about this!)
• Testing testing testing! I want to rerun all notebooks/examples and make sure the user experience is good.
• CI testing - We need to make sure we don't regress on any of this
• Discoverability: After this is merged we should update notebooks/examples to show off the cool new features, and document our TF workflow/philosophy somewhere that new users will find.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[Rocketknight1]

(Requesting reviews now that @gante is back)

[gante]

<3 This is great, Keras users will definitely feel more at home

I've added two comments: a suggestion (for potentially more organized code) and a question. Other than that, LGTM!

[gante]

Suggestion for discussion: Could these hardcoded defaults be part of the TFXXXLoss (e.g.)? The actual models inherit from these classes (e.g.), and thus we could add them on a per-loss basis, as opposed to having a big if/else in the train/test steps :D

[Rocketknight1]

I explored this when writing the PR! I think that would work in a lot of cases, but there are some models which have their own custom losses, and other models that define hf_compute_loss in the model class itself.

So I'm not sure if moving this to the Loss classes would be that easy, but for cleanliness, I can extract this to a method called something like infer_label_to_output_map() and just call that in train_step instead?

[gante]

Extracting to an external function sounds good 👍 (especially because it is reused between train and test)

[gante]

This converts y from a dictionary with one item to the value of that dictionary entry. Looking below, it seems like it should handle dicts correctly. What's happening here?

[Rocketknight1]

The reason I did it this way is to catch more cases, but I realize now I could have been a lot smarter about it. One sec!

[Rocketknight1]

Fixed! This code was added because the user often passes a dict where the key is "labels", which is not the name of any of the outputs. The correct thing to do for those models is to map the "labels" tensor to the first model output - I changed this line so that it checks the single key is called "labels" before doing so.

[Rocketknight1]

Fixed it up a little more - now we try to map by key name before falling back to mapping to the first output as a last resort.

[gante]

LGTM 👍