Fix train_step, test_step and tests for CLIP

[Rocketknight1]

CLIP models were not being tested correctly with fit() because the test skipped models without a hf_compute_loss method. This skip was added to skip base models like TFBERTModel that do not have specific output heads and losses. However, it also skips models like CLIP that do not use compute_loss / hf_compute_loss methods.

The new test checks whether the model's return type dataclass has a loss key, which is a more reliable check. Enabling this reveals the bug in fit() for TFClip, so this PR also includes fixes to train_step and test_step for CLIP and models like it that require return_loss=True to be passed, but do not set it by default.

Draft for now because this will likely flush out other bugs or cause other problems!

Fixes #18670.

[HuggingFaceDocBuilderDev]

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

[Rocketknight1]

The tests passed. I'm as surprised as everyone else. It's ready for review!

[ydshieh]

Thanks @Rocketknight1 for quick fix. I am wondering if it makes sense to wrap the loss computation for TFCLIP into a hf_compute_loss?

I also see for TFSegformerForSemanticSegmentation, it defines hf_compute_loss under its own model class.

If doable, maybe it's good to add TFSemanticSegmentationLoss, TFcontrastiveLoss class etc.

Just want to hear some opinions. cc @gante @amyeroberts

[Rocketknight1]

@ydshieh I don't think that's necessary - the new check we use means we don't need to rely on hf_compute_loss being present anymore!

[Rocketknight1]

Just realized that using return instead of continue in the tests was skipping a lot of tests, which might have been why it was green so easily. Rerunning everything!

[gante]

Looks good 👍 And I agree, relying on a loss output is safer!

Will re-review after fixing errors (I see the latest commit, which replaces the return, broke some tests :D)

[Rocketknight1]

Further updates: Now that we're no longer incorrectly skipping tests, this turned up quite a few bugs! The main source of issues is that some more recent models are returning scalar losses, but both our test_loss_computation test and Keras fit() expect that the loss has some shape, even if that shape is (1,).

[amyeroberts]

Overall structure looks good to me. Leaving some comments as I was 👀 , but just nits.

[amyeroberts]

Can we assumed there's only one added label? Not sure if it matters, but if we can and it does, can we add an assert above?

[sgugger]

Left a nit, but LGTM! Thanks for working on this!