Add tests for Presets workflow, Add Metadata (#1326)

* Add metadata and Albert preset utils test

* Add Bart bytepiece preset workflow test

* Add BERT WordPiece preset workflow test

* Parameterize tests, switch to classifier, address comments

* Address comments and nits

* Fix formatting

* Add large test marker

keras_nlp/tokenizers/sentence_piece_tokenizer.py

@@ -38,7 +38,7 @@
     tf_text = None
 
 
-VOCAB_FILENAME = "vocabulary.txt"
+VOCAB_FILENAME = "vocabulary.spm"
 
 
 @keras_nlp_export("keras_nlp.tokenizers.SentencePieceTokenizer")
@@ -132,7 +132,7 @@ def __init__(
 
     def save_assets(self, dir_path):
         path = os.path.join(dir_path, VOCAB_FILENAME)
-        with open(path, "w") as file:
+        with open(path, "wb") as file:
             file.write(self.proto)
 
     def load_assets(self, dir_path):

keras_nlp/utils/preset_utils.py

@@ -100,19 +100,26 @@ def save_to_preset(
     # Include references to weights and assets.
     config["assets"] = assets
     config["weights"] = weights_filename if save_weights else None
-    recursive_pop(config, "config_config")
+    recursive_pop(config, "compile_config")
     recursive_pop(config, "build_config")
     with open(config_path, "w") as config_file:
         config_file.write(json.dumps(config, indent=4))
 
+    from keras_nlp import __version__ as keras_nlp_version
+
+    keras_version = keras.version() if hasattr(keras, "version") else None
+
     # Save any associated metadata.
-    metadata = {
-        # TODO: save keras version and keras-nlp version.
-        "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
-    }
-    metadata_path = os.path.join(preset, "metadata.json")
-    with open(metadata_path, "w") as metadata_file:
-        metadata_file.write(json.dumps(metadata, indent=4))
+    if config_filename == "config.json":
+        metadata = {
+            "keras_version": keras_version,
+            "keras_nlp_version": keras_nlp_version,
+            "parameter_count": layer.count_params(),
+            "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
+        }
+        metadata_path = os.path.join(preset, "metadata.json")
+        with open(metadata_path, "w") as metadata_file:
+            metadata_file.write(json.dumps(metadata, indent=4))
 
 
 def load_from_preset(

keras_nlp/utils/preset_utils_test.py

@@ -0,0 +1,89 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+
+import pytest
+from absl.testing import parameterized
+
+from keras_nlp.models import AlbertClassifier
+from keras_nlp.models import BertClassifier
+from keras_nlp.models import RobertaClassifier
+from keras_nlp.tests.test_case import TestCase
+from keras_nlp.utils import preset_utils
+
+
+class PresetUtilsTest(TestCase):
+    @parameterized.parameters(
+        (AlbertClassifier, "albert_base_en_uncased", "sentencepiece"),
+        (RobertaClassifier, "roberta_base_en", "bytepair"),
+        (BertClassifier, "bert_tiny_en_uncased", "wordpiece"),
+    )
+    @pytest.mark.keras_3_only
+    @pytest.mark.large
+    def test_preset_saving(self, cls, preset_name, tokenizer_type):
+        save_dir = self.get_temp_dir()
+        model = cls.from_preset(preset_name, num_classes=2)
+        preset_utils.save_to_preset(model, save_dir)
+
+        if tokenizer_type == "bytepair":
+            vocab_filename = "assets/tokenizer/vocabulary.json"
+            expected_assets = [
+                "assets/tokenizer/vocabulary.json",
+                "assets/tokenizer/merges.txt",
+            ]
+        elif tokenizer_type == "sentencepiece":
+            vocab_filename = "assets/tokenizer/vocabulary.spm"
+            expected_assets = ["assets/tokenizer/vocabulary.spm"]
+        else:
+            vocab_filename = "assets/tokenizer/vocabulary.txt"
+            expected_assets = ["assets/tokenizer/vocabulary.txt"]
+
+        # Check existence of files
+        self.assertTrue(os.path.exists(os.path.join(save_dir, vocab_filename)))
+        self.assertTrue(os.path.exists(os.path.join(save_dir, "config.json")))
+        self.assertTrue(
+            os.path.exists(os.path.join(save_dir, "model.weights.h5"))
+        )
+        self.assertTrue(os.path.exists(os.path.join(save_dir, "metadata.json")))
+
+        # Check the model config (`config.json`)
+        config_json = open(os.path.join(save_dir, "config.json"), "r").read()
+        self.assertTrue(
+            "build_config" not in config_json
+        )  # Test on raw json to include nested keys
+        self.assertTrue(
+            "compile_config" not in config_json
+        )  # Test on raw json to include nested keys
+        config = json.loads(config_json)
+        self.assertAllEqual(config["assets"], expected_assets)
+        self.assertEqual(config["weights"], "model.weights.h5")
+
+        # Try loading the model from preset directory
+        restored_model = preset_utils.load_from_preset(save_dir)
+
+        train_data = (
+            ["the quick brown fox.", "the slow brown fox."],  # Features.
+        )
+        model_input_data = model.preprocessor(*train_data)
+        restored_model_input_data = restored_model.preprocessor(*train_data)
+
+        # Check that saved vocab is equal to the original preset vocab
+        self.assertAllClose(model_input_data, restored_model_input_data)
+
+        # Check model outputs
+        self.assertAllEqual(
+            model(model_input_data), restored_model(restored_model_input_data)
+        )