Test transformer loss decreases during training

[GargiGupta-io]

Description of the change

• Extend the small transformer integration test to train end-to-end for 15 epochs.
• Seed model initialization and assert that the initial and final losses are finite and that training decreases the loss.

Reference

Fixes #1276.

Colab Notebook

Not applicable; this is an integration-test-only change.

Validation

• The modified transformer model and training assertions pass on TensorFlow, JAX, and PyTorch.
• Local validation substituted deterministic token IDs for the tokenizer because tensorflow-text is unavailable in the local environment; CI exercises the full tokenizer path.
• Ruff formatting and lint checks pass.

Checklist

• I have added all the necessary test coverage for my change.
• I have verified that the integration test works with TensorFlow, JAX, and PyTorch.
• My PR is based on the latest changes of the main branch.
• I have followed the Keras Hub Model contribution guidelines (not applicable; no model is added).
• I have followed the Keras Hub API design guidelines (not applicable; no API is changed).
• I have signed the Contributor License Agreement.

[gemini-code-assist]

Code Review

This pull request updates the basic usage integration test to train the transformer model end-to-end and assert a finite, decreasing loss. Feedback was provided to increase the number of training epochs from 5 to 15 to ensure the test is robust and does not suffer from flakiness across different Keras backends.

[gemini-code-assist]

Asserting a strict decrease in loss after only 5 epochs (5 gradient steps) on a tiny dataset of 2 samples can be fragile and prone to flakiness. Due to differences in weight initialization and random number generation across Keras backends (TensorFlow, JAX, PyTorch), the initial weights and training trajectories will differ even with a set seed. With Adam's momentum initialization, 5 steps might not be enough to guarantee a strict decrease in loss on all backends.

Increasing the number of epochs to 15 or 20 ensures that the model has enough steps to consistently minimize the loss and overfit the 2 samples, making the test much more robust across all backends without adding noticeable overhead.

Suggested change

	model.fit(x, y, epochs=5, verbose=0)
	final_loss = model.evaluate(x, y, verbose=0)
	# Make sure we have a valid loss.
	self.assertGreater(loss, 0)
	# Make sure training produces a finite, decreasing loss.
	self.assertTrue(np.isfinite(initial_loss))
	self.assertTrue(np.isfinite(final_loss))
	self.assertLess(final_loss, initial_loss)
	model.fit(x, y, epochs=15, verbose=0)
	final_loss = model.evaluate(x, y, verbose=0)
	# Make sure training produces a finite, decreasing loss.
	self.assertTrue(np.isfinite(initial_loss))
	self.assertTrue(np.isfinite(final_loss))
	self.assertLess(final_loss, initial_loss)

References

1. Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. ^(link)