Repository: keras-team/keras-nlp

Issue:
Register our BERT components as a serializable Keras objects
We would like to be able to annotate our new BERT models and components with `@keras.utils.register_keras_serializable(package="keras_nlp")`. This will allow the actual python object to restore after saving, which will be import to access things like `model.tokenizer` or `model.token_embedding` after restoring.


Discussion:
Doing this naively leads to some errors.

`E       TypeError: __init__() missing 2 required positional arguments: 'inputs' and 'outputs'`

So a little more investigation will be needed.

Resolving patch:
diff --git a/keras_nlp/models/bert/bert_models.py b/keras_nlp/models/bert/bert_models.py
--- a/keras_nlp/models/bert/bert_models.py
+++ b/keras_nlp/models/bert/bert_models.py
@@ -71,6 +71,7 @@ def _handle_pretrained_model_arguments(bert_variant, weights, vocabulary_size):
     return weights, vocabulary_size
 
 
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BertCustom(keras.Model):
     """BERT encoder network with custom hyperparmeters.
 
@@ -248,21 +249,22 @@ def __init__(
         self.cls_token_index = cls_token_index
 
     def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "vocabulary_size": self.vocabulary_size,
-                "hidden_dim": self.hidden_dim,
-                "intermediate_dim": self.intermediate_dim,
-                "num_layers": self.num_layers,
-                "num_heads": self.num_heads,
-                "max_sequence_length": self.max_sequence_length,
-                "num_segments": self.num_segments,
-                "dropout": self.dropout,
-                "cls_token_index": self.cls_token_index,
-            }
-        )
-        return config
+        return {
+            "vocabulary_size": self.vocabulary_size,
+            "hidden_dim": self.hidden_dim,
+            "intermediate_dim": self.intermediate_dim,
+            "num_layers": self.num_layers,
+            "num_heads": self.num_heads,
+            "max_sequence_length": self.max_sequence_length,
+            "num_segments": self.num_segments,
+            "dropout": self.dropout,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 
 MODEL_DOCSTRING = """Bert "{type}" architecture.
diff --git a/keras_nlp/models/bert/bert_tasks.py b/keras_nlp/models/bert/bert_tasks.py
--- a/keras_nlp/models/bert/bert_tasks.py
+++ b/keras_nlp/models/bert/bert_tasks.py
@@ -67,6 +67,7 @@
 """
 
 
+@keras.utils.register_keras_serializable(package="keras_nlp")
 class BertClassifier(keras.Model):
     def __init__(
         self,
@@ -104,6 +105,20 @@ def __init__(
         self.backbone = backbone
         self.num_classes = num_classes
 
+    def get_config(self):
+        return {
+            "backbone": keras.layers.serialize(self.backbone),
+            "num_classes": self.num_classes,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        if "backbone" in config:
+            config["backbone"] = keras.layers.deserialize(config["backbone"])
+        return cls(**config)
+
 
 setattr(
     BertClassifier,


Test patch:
diff --git a/keras_nlp/models/bert/bert_models_test.py b/keras_nlp/models/bert/bert_models_test.py
--- a/keras_nlp/models/bert/bert_models_test.py
+++ b/keras_nlp/models/bert/bert_models_test.py
@@ -120,12 +120,19 @@ def test_bert_base_vocab_error(self):
                 name="encoder",
             )
 
-    def test_saving_model(self):
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
         model_output = self.model(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "model")
-        self.model.save(save_path)
+        self.model.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, BertCustom)
+
+        # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(
             model_output["pooled_output"], restored_output["pooled_output"]
diff --git a/keras_nlp/models/bert/bert_tasks_test.py b/keras_nlp/models/bert/bert_tasks_test.py
--- a/keras_nlp/models/bert/bert_tasks_test.py
+++ b/keras_nlp/models/bert/bert_tasks_test.py
@@ -76,11 +76,18 @@ def test_bert_classifier_compile_batched_ds(self, jit_compile):
         self.classifier.compile(jit_compile=jit_compile)
         self.classifier.predict(self.input_dataset)
 
-    def test_saving_model(self):
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
         model_output = self.classifier(self.input_batch)
         save_path = os.path.join(self.get_temp_dir(), "model")
-        self.classifier.save(save_path)
+        self.classifier.save(save_path, save_format)
         restored_model = keras.models.load_model(save_path)
 
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, BertClassifier)
+
+        # Check that output matches.
         restored_output = restored_model(self.input_batch)
         self.assertAllClose(model_output, restored_output)