Repository: keras-team/keras-nlp

Issue:
Fix saved model testing
We have a silly bug for a few of our saved model tests where we are calling

```
model.save(save_path, save_format)
```

Instead of

```
model.save(save_path, save_format=save_format)
```

Which means we aren't properly the h5 format. If you fix this issue, you actually hit another issue with h5 saving.

`FAILED keras_nlp/models/bert/bert_tasks_test.py::BertClassifierTest::test_saving_model_save_format_h5 - ValueError: Unknown layer: BertClassifier. Please ensure this object is passed to the `custom_objects` argument.`


Discussion:

Resolving patch:
diff --git a/keras_nlp/conftest.py b/keras_nlp/conftest.py
--- a/keras_nlp/conftest.py
+++ b/keras_nlp/conftest.py
@@ -14,6 +14,8 @@
 import sys
 
 import pytest
+import tensorflow as tf
+from packaging import version
 
 
 def pytest_addoption(parser):
@@ -50,6 +52,10 @@ def pytest_collection_modifyitems(config, items):
     skip_xla = pytest.mark.skipif(
         sys.platform == "darwin", reason="XLA unsupported on MacOS."
     )
+    skip_keras_saving_test = pytest.mark.skipif(
+        version.parse(tf.__version__) < version.parse("2.12"),
+        reason="keras_v3 format requires tf > 2.12.",
+    )
     skip_large = pytest.mark.skipif(
         not run_large_tests, reason="need --run_large option to run"
     )
@@ -59,6 +65,8 @@ def pytest_collection_modifyitems(config, items):
     for item in items:
         if "jit_compile_true" in item.name:
             item.add_marker(skip_xla)
+        if "keras_format" in item.name:
+            item.add_marker(skip_keras_saving_test)
         if "large" in item.keywords:
             item.add_marker(skip_large)
         if "extra_large" in item.keywords:


Test patch:
diff --git a/keras_nlp/layers/fnet_encoder_test.py b/keras_nlp/layers/fnet_encoder_test.py
--- a/keras_nlp/layers/fnet_encoder_test.py
+++ b/keras_nlp/layers/fnet_encoder_test.py
@@ -117,8 +117,11 @@ def test_checkpointing_fnet_encoder(self):
         encoder2_output = encoder2(data)
         self.assertAllClose(encoder1_output, encoder2_output)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         model = keras.Sequential(
             [
                 keras.Input(shape=(4, 6)),
@@ -129,8 +132,8 @@ def test_save_model(self, format):
         )
         data = tf.random.uniform(shape=[2, 4, 6])
         model(data)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(data)
diff --git a/keras_nlp/layers/multi_segment_packer_test.py b/keras_nlp/layers/multi_segment_packer_test.py
--- a/keras_nlp/layers/multi_segment_packer_test.py
+++ b/keras_nlp/layers/multi_segment_packer_test.py
@@ -161,8 +161,11 @@ def test_config(self):
             cloned_packer([seq1, seq2]),
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
         seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
         packer = MultiSegmentPacker(
@@ -174,8 +177,8 @@ def test_saving(self, format):
         )
         outputs = packer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model((seq1, seq2)),
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -303,8 +303,11 @@ def test_get_config_and_from_config(self):
             restored_encoder.get_config(), {**config, **expected_config_subset}
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         max_sequence_length = 4
         feature_size = 6
         test_layer = position_embedding.PositionEmbedding(
@@ -317,8 +320,8 @@ def test_save_model(self, format):
         data = tf.zeros(shape=[2, max_sequence_length, feature_size])
         model(data)
 
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -127,8 +127,11 @@ def test_mask_propagation(self):
         outputs = test_layer(input_data)
         self.assertAllEqual(outputs._keras_mask, mask)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         vocabulary_size = 5
         sequence_length = 4
         embedding_dim = 3
@@ -144,8 +147,8 @@ def test_save_model(self, format):
         data = tf.zeros(shape=[2, sequence_length])
         model(data)
 
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model.predict(data)
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -265,8 +265,11 @@ def test_mask_propagation_without_cross_attention(self):
         outputs = decoder(decoder_sequence)
         self.assertAllEqual(outputs._keras_mask, mask)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         encoder_input = keras.Input(shape=[4, 6])
         decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
@@ -282,16 +285,19 @@ def test_save_model(self, format):
         encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model([decoder_sequence, encoder_sequence])
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
 
         loaded_model = keras.models.load_model(path)
         model_output = model([decoder_sequence, encoder_sequence])
         loaded_model_output = loaded_model([decoder_sequence, encoder_sequence])
         self.assertAllClose(model_output, loaded_model_output)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model_without_cross_attention(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model_without_cross_attention(self, save_format, filename):
         decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
@@ -305,8 +311,8 @@ def test_save_model_without_cross_attention(self, format):
         )
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         model(decoder_sequence)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         loaded_model = keras.models.load_model(path)
 
         model_output = model(decoder_sequence)
diff --git a/keras_nlp/layers/transformer_encoder_test.py b/keras_nlp/layers/transformer_encoder_test.py
--- a/keras_nlp/layers/transformer_encoder_test.py
+++ b/keras_nlp/layers/transformer_encoder_test.py
@@ -158,8 +158,11 @@ def test_checkpointing_transformer_encoder(self):
         encoder2_output = encoder2(data)
         self.assertAllClose(encoder1_output, encoder2_output)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_save_model(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         model = keras.Sequential(
             [
                 keras.Input(shape=(4, 6)),
@@ -172,8 +175,8 @@ def test_save_model(self, format):
         )
         data = tf.random.uniform(shape=[2, 4, 6])
         model_output = model(data)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
 
         loaded_model = keras.models.load_model(path)
         loaded_model_output = loaded_model(data)
diff --git a/keras_nlp/models/bert/bert_backbone_test.py b/keras_nlp/models/bert/bert_backbone_test.py
--- a/keras_nlp/models/bert/bert_backbone_test.py
+++ b/keras_nlp/models/bert/bert_backbone_test.py
@@ -85,12 +85,13 @@ def test_compile_batched_ds(self, jit_compile):
         self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/bert/bert_classifier_test.py b/keras_nlp/models/bert/bert_classifier_test.py
--- a/keras_nlp/models/bert/bert_classifier_test.py
+++ b/keras_nlp/models/bert/bert_classifier_test.py
@@ -105,12 +105,13 @@ def test_bert_classifier_fit_no_preprocessing(self, jit_compile):
         self.classifier_no_preprocessing.fit(self.preprocessed_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.classifier.predict(self.raw_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.classifier.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.classifier.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/bert/bert_preprocessor_test.py b/keras_nlp/models/bert/bert_preprocessor_test.py
--- a/keras_nlp/models/bert/bert_preprocessor_test.py
+++ b/keras_nlp/models/bert/bert_preprocessor_test.py
@@ -112,14 +112,15 @@ def test_errors_for_2d_list_input(self):
             self.preprocessor(ambiguous_input)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["THE QUICK BROWN FOX."])
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/bert/bert_tokenizer_test.py b/keras_nlp/models/bert/bert_tokenizer_test.py
--- a/keras_nlp/models/bert/bert_tokenizer_test.py
+++ b/keras_nlp/models/bert/bert_tokenizer_test.py
@@ -57,15 +57,16 @@ def test_vocabulary_size(self):
         self.assertEqual(tokenizer.vocabulary_size(), 13)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["THE QUICK BROWN FOX."])
         tokenizer = BertTokenizer(vocabulary=self.vocab)
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py b/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
--- a/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone_test.py
@@ -77,12 +77,13 @@ def test_distilbert_base_compile_batched_ds(self, jit_compile):
         self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
--- a/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier_test.py
@@ -70,13 +70,14 @@ def test_distilbert_classifier_compile_batched_ds(self, jit_compile):
         model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = DistilBertClassifier(self.model, 4, 128, name="classifier")
         model_output = model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py b/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
--- a/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_preprocessor_test.py
@@ -102,14 +102,15 @@ def test_errors_for_2d_list_input(self):
             self.preprocessor(ambiguous_input)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["THE QUICK BROWN FOX."])
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
--- a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
@@ -59,15 +59,16 @@ def test_vocabulary_size(self):
         self.assertEqual(tokenizer.vocabulary_size(), 13)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["THE QUICK BROWN FOX."])
         tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/gpt2/gpt2_backbone_test.py b/keras_nlp/models/gpt2/gpt2_backbone_test.py
--- a/keras_nlp/models/gpt2/gpt2_backbone_test.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone_test.py
@@ -79,12 +79,13 @@ def test_gpt2_compile_batched_ds(self, jit_compile):
         self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
--- a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
+++ b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
@@ -67,16 +67,17 @@ def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 10)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant([" airplane at airport"])
 
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.tokenizer(inputs)
         model = keras.Model(inputs, outputs)
 
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
 
         restored_model = keras.models.load_model(path)
diff --git a/keras_nlp/models/roberta/roberta_backbone_test.py b/keras_nlp/models/roberta/roberta_backbone_test.py
--- a/keras_nlp/models/roberta/roberta_backbone_test.py
+++ b/keras_nlp/models/roberta/roberta_backbone_test.py
@@ -84,12 +84,13 @@ def test_variable_sequence_length_call_roberta(self):
             )
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/roberta/roberta_classifier_test.py b/keras_nlp/models/roberta/roberta_classifier_test.py
--- a/keras_nlp/models/roberta/roberta_classifier_test.py
+++ b/keras_nlp/models/roberta/roberta_classifier_test.py
@@ -68,13 +68,14 @@ def test_roberta_classifier_compile_batched_ds(self, jit_compile):
         model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = RobertaClassifier(self.model, 4, 128, name="classifier")
         model_output = model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/roberta/roberta_preprocessor_test.py b/keras_nlp/models/roberta/roberta_preprocessor_test.py
--- a/keras_nlp/models/roberta/roberta_preprocessor_test.py
+++ b/keras_nlp/models/roberta/roberta_preprocessor_test.py
@@ -139,16 +139,17 @@ def test_errors_for_2d_list_input(self):
             self.preprocessor(ambiguous_input)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant([" airplane at airport"])
 
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
 
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
 
         restored_model = keras.models.load_model(path)
diff --git a/keras_nlp/models/roberta/roberta_tokenizer_test.py b/keras_nlp/models/roberta/roberta_tokenizer_test.py
--- a/keras_nlp/models/roberta/roberta_tokenizer_test.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer_test.py
@@ -66,16 +66,17 @@ def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 12)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant([" airplane at airport"])
 
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.tokenizer(inputs)
         model = keras.Model(inputs, outputs)
 
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
 
         restored_model = keras.models.load_model(path)
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_backbone_test.py
@@ -84,12 +84,13 @@ def test_variable_sequence_length_call_xlm_roberta(self):
             )
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model_output = self.model(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier_test.py
@@ -70,13 +70,14 @@ def test_xlm_roberta_classifier_compile_batched_ds(self, jit_compile):
         classifier.predict(self.input_dataset)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         classifier = XLMRobertaClassifier(self.model, 4, 128, name="classifier")
         classifier_output = classifier(self.input_batch)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        classifier.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        classifier.save(save_path, save_format=save_format)
         restored_classifier = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessor_test.py
@@ -133,14 +133,15 @@ def test_errors_for_2d_list_input(self):
             self.preprocessor(ambiguous_input)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown fox"])
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.preprocessor(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_tokenizer_test.py
@@ -97,16 +97,17 @@ def test_token_to_id(self):
         self.assertEqual(self.tokenizer.token_to_id("▁round"), 10)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown fox"])
 
         inputs = keras.Input(dtype="string", shape=())
         outputs = self.tokenizer(inputs)
         model = keras.Model(inputs, outputs)
 
-        path = os.path.join(self.get_temp_dir(), "model")
+        path = os.path.join(self.get_temp_dir(), filename)
         model.save(path, save_format=save_format)
 
         restored_model = keras.models.load_model(path)
diff --git a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
--- a/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_pair_tokenizer_test.py
@@ -135,15 +135,18 @@ def test_config(self):
             cloned_tokenizer(input_data),
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["the quick brown whale."])
         tokenizer = self.tokenizer
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
--- a/keras_nlp/tokenizers/byte_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -244,8 +244,11 @@ def test_config(self):
         }
         self.assertEqual(tokenizer.get_config(), exp_config)
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["this is fun"])
 
         tokenizer = ByteTokenizer(
@@ -258,8 +261,8 @@ def test_saving(self, format):
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
--- a/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_test.py
@@ -185,8 +185,11 @@ def test_config(self):
             cloned_tokenizer(input_data),
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         filepath = os.path.join(self.get_temp_dir(), "model.txt")
         input_data = tf.constant(["the quick brown whale."])
         with tf.io.gfile.GFile(filepath, "wb") as file:
@@ -197,8 +200,8 @@ def test_saving(self, format):
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py b/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
--- a/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
+++ b/keras_nlp/tokenizers/unicode_codepoint_tokenizer_test.py
@@ -344,8 +344,11 @@ def test_config(self):
             exp_config_different_encoding,
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["ninjas and samurais", "time travel"])
 
         tokenizer = UnicodeCodepointTokenizer(
@@ -359,8 +362,8 @@ def test_saving(self, format):
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -199,8 +199,11 @@ def test_config(self):
             cloned_tokenizer(input_data),
         )
 
-    @parameterized.named_parameters(("tf_format", "tf"), ("h5_format", "h5"))
-    def test_saving(self, format):
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
         input_data = tf.constant(["quick brOWN whale"])
         vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
         tokenizer = WordPieceTokenizer(
@@ -213,8 +216,8 @@ def test_saving(self, format):
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)
         model = keras.Model(inputs, outputs)
-        path = os.path.join(self.get_temp_dir(), "model")
-        model.save(path, save_format=format)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
         restored_model = keras.models.load_model(path)
         self.assertAllEqual(
             model(input_data),
diff --git a/keras_nlp/utils/pipeline_model_test.py b/keras_nlp/utils/pipeline_model_test.py
--- a/keras_nlp/utils/pipeline_model_test.py
+++ b/keras_nlp/utils/pipeline_model_test.py
@@ -133,14 +133,15 @@ def test_on_batch(self):
         model.predict_on_batch(x=x)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = NoopPipeline()
         x = tf.random.uniform((8, 5))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(
             save_path, custom_objects={"NoopPipeline": NoopPipeline}
         )
@@ -248,14 +249,15 @@ def test_on_batch_no_preprocessing(self):
         model.predict_on_batch(x=x)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = FeaturePipeline()
         x = tf.strings.as_string(tf.random.uniform((8, 5)))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(
             save_path, custom_objects={"FeaturePipeline": FeaturePipeline}
         )
@@ -356,14 +358,15 @@ def test_on_batch_no_preprocessing(self):
         model.predict_on_batch(x=x)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = LabelPipeline()
         x = tf.random.uniform((8, 5))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(
             save_path, custom_objects={"LabelPipeline": LabelPipeline}
         )
@@ -447,14 +450,15 @@ def test_on_batch_no_preprocessing(self):
         model.predict_on_batch(x=x)
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = DataPipeline()
         data = tf.strings.as_string(tf.random.uniform((8, 1)))
         model_output = model.predict(data)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(
             save_path, custom_objects={"DataPipeline": DataPipeline}
         )
@@ -495,14 +499,15 @@ def test_fit_no_preprocessing(self):
         model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
     @parameterized.named_parameters(
-        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format):
+    def test_saved_model(self, save_format, filename):
         model = FunctionalPipeline()
         x = tf.strings.as_string(tf.random.uniform((8, 5)))
         model_output = model.predict(x)
-        save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path, save_format)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        model.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(
             save_path, custom_objects={"FunctionalPipeline": FunctionalPipeline}
         )