Run keras saving tests on nightly and fix RobertaClassifier test

.github/workflows/nightly.yml

@@ -12,10 +12,12 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.8
+      - name: Set up Python 3.9
         uses: actions/setup-python@v1
+        # TODO: we probably want 3.8 here, but are facing an issue with core
+        # keras.
         with:
-          python-version: 3.8
+          python-version: 3.9
       - name: Get pip cache dir
         id: pip-cache
         run: |

keras_nlp/conftest.py

@@ -79,8 +79,9 @@ def pytest_collection_modifyitems(config, items):
     skip_xla = pytest.mark.skipif(
         sys.platform == "darwin", reason="XLA unsupported on MacOS."
     )
+    # Run Keras saving tests on 2.12 stable, nightlies and later releases.
     skip_keras_saving_test = pytest.mark.skipif(
-        version.parse(tf.__version__) < version.parse("2.12"),
+        version.parse(tf.__version__) < version.parse("2.12.0-dev0"),
         reason="keras_v3 format requires tf > 2.12.",
     )
     skip_large = pytest.mark.skipif(

keras_nlp/models/gpt2/gpt2_causal_lm_test.py

@@ -144,11 +144,11 @@ def test_gpt2_causal_lm_generate(self, jit_compile):
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format, filename):
+    def test_saved_model(self, save_format, filename):
         keras.utils.set_random_seed(42)
         model_output = self.causal_lm.predict(self.raw_batch)
         save_path = os.path.join(self.get_temp_dir(), filename)
-        self.causal_lm.save(save_path, save_format)
+        self.causal_lm.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.

keras_nlp/models/roberta/roberta_classifier_test.py

@@ -125,10 +125,10 @@ def test_roberta_classifier_fit_no_preprocessing(self, jit_compile):
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
     )
-    def test_saving_model(self, save_format, filename):
+    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(), filename)
-        self.classifier.save(save_path, save_format)
+        self.classifier.save(save_path, save_format=save_format)
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.