Start landing code for Kaggle integration (#1320)

* Demo preset directories

* Address comments

* Move packer and masker to build and debug tests

* Address comments

* Fix maskedLM preprocessor test

* Fix remaining tests

* Fix serialization tests on tf backend

---------

Co-authored-by: Neel Kovelamudi <[email protected]>

keras_nlp/layers/modeling/token_and_position_embedding.py

@@ -122,7 +122,7 @@ def get_config(self):
                 ),
                 "tie_weights": self.token_embedding.tie_weights,
                 "mask_zero": self.token_embedding.mask_zero,
-            },
+            }
         )
         return config
 

keras_nlp/layers/modeling/token_and_position_embedding_test.py

@@ -14,7 +14,6 @@
 
 import numpy as np
 
-from keras_nlp.backend import keras
 from keras_nlp.backend import ops
 from keras_nlp.backend import random
 from keras_nlp.layers.modeling.token_and_position_embedding import (
@@ -31,7 +30,7 @@ def test_layer_behaviors(self):
                 "vocabulary_size": 5,
                 "sequence_length": 4,
                 "embedding_dim": 3,
-                "embeddings_initializer": keras.initializers.Constant(1.0),
+                "embeddings_initializer": "ones",
             },
             input_data=random.randint(minval=0, maxval=5, shape=(2, 4)),
             expected_output_shape=(2, 4, 3),

keras_nlp/layers/preprocessing/preprocessing_layer.py

@@ -29,7 +29,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self._convert_input_args = False
         self._allow_non_tensor_positional_args = True
-        self.built = True
+        # Most pre-preprocessing has no build.
+        if not hasattr(self, "build"):
+            self.built = True
+
+    def get_build_config(self):
+        return None
 
     def __call__(self, *args, **kwargs):
         # Always place on CPU for preprocessing, to avoid expensive back and

keras_nlp/models/backbone.py

@@ -15,6 +15,8 @@
 import os
 
 from keras_nlp.backend import keras
+from keras_nlp.utils.preset_utils import check_preset_class
+from keras_nlp.utils.preset_utils import load_from_preset
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
@@ -66,6 +68,31 @@ def from_config(cls, config):
     def presets(cls):
         return {}
 
+    @classmethod
+    def _legacy_from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        filename = os.path.basename(metadata["weights_url"])
+        weights = keras.utils.get_file(
+            filename,
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model
+
     @classmethod
     def from_preset(
         cls,
@@ -94,35 +121,17 @@ def from_preset(
         )
         ```
         """
-
-        if not cls.presets:
-            raise NotImplementedError(
-                "No presets have been created for this class."
-            )
-
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if not load_weights:
-            return model
-
-        filename = os.path.basename(metadata["weights_url"])
-        weights = keras.utils.get_file(
-            filename,
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
+        # TODO: delete me!
+        if preset in cls.presets:
+            return cls._legacy_from_preset(preset, **kwargs)
+
+        check_preset_class(preset, cls)
+        return load_from_preset(
+            preset,
+            load_weights=load_weights,
+            config_overrides=kwargs,
         )
 
-        model.load_weights(weights)
-        return model
-
     def __init_subclass__(cls, **kwargs):
         # Use __init_subclass__ to setup a correct docstring for from_preset.
         super().__init_subclass__(**kwargs)

keras_nlp/models/bert/bert_classifier_test.py

@@ -36,7 +36,7 @@ def setUp(self):
             num_heads=2,
             hidden_dim=2,
             intermediate_dim=4,
-            max_sequence_length=self.preprocessor.packer.sequence_length,
+            max_sequence_length=self.preprocessor.sequence_length,
         )
         self.init_kwargs = {
             "preprocessor": self.preprocessor,

keras_nlp/models/bert/bert_masked_lm_preprocessor.py

@@ -134,33 +134,30 @@ def __init__(
             truncate=truncate,
             **kwargs,
         )
-
+        self.mask_selection_rate = mask_selection_rate
+        self.mask_selection_length = mask_selection_length
+        self.mask_token_rate = mask_token_rate
+        self.random_token_rate = random_token_rate
+        self.masker = None
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # Defer masker creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
         self.masker = MaskedLMMaskGenerator(
-            mask_selection_rate=mask_selection_rate,
-            mask_selection_length=mask_selection_length,
-            mask_token_rate=mask_token_rate,
-            random_token_rate=random_token_rate,
-            vocabulary_size=tokenizer.vocabulary_size(),
-            mask_token_id=tokenizer.mask_token_id,
+            mask_selection_rate=self.mask_selection_rate,
+            mask_selection_length=self.mask_selection_length,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+            vocabulary_size=self.tokenizer.vocabulary_size(),
+            mask_token_id=self.tokenizer.mask_token_id,
             unselectable_token_ids=[
-                tokenizer.cls_token_id,
-                tokenizer.sep_token_id,
-                tokenizer.pad_token_id,
+                self.tokenizer.cls_token_id,
+                self.tokenizer.sep_token_id,
+                self.tokenizer.pad_token_id,
             ],
         )
 
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "mask_selection_rate": self.masker.mask_selection_rate,
-                "mask_selection_length": self.masker.mask_selection_length,
-                "mask_token_rate": self.masker.mask_token_rate,
-                "random_token_rate": self.masker.random_token_rate,
-            }
-        )
-        return config
-
     def call(self, x, y=None, sample_weight=None):
         if y is not None or sample_weight is not None:
             logging.warning(
@@ -187,3 +184,15 @@ def call(self, x, y=None, sample_weight=None):
         y = masker_outputs["mask_ids"]
         sample_weight = masker_outputs["mask_weights"]
         return pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config

keras_nlp/models/bert/bert_masked_lm_test.py

@@ -43,7 +43,7 @@ def setUp(self):
             num_heads=2,
             hidden_dim=2,
             intermediate_dim=4,
-            max_sequence_length=self.preprocessor.packer.sequence_length,
+            max_sequence_length=self.preprocessor.sequence_length,
         )
         self.init_kwargs = {
             "preprocessor": self.preprocessor,

keras_nlp/models/bert/bert_preprocessor.py

@@ -139,24 +139,32 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.tokenizer = tokenizer
-        self.packer = MultiSegmentPacker(
-            start_value=self.tokenizer.cls_token_id,
-            end_value=self.tokenizer.sep_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            truncate=truncate,
-            sequence_length=sequence_length,
-        )
+        self.sequence_length = sequence_length
+        self.truncate = truncate
+        self.packer = None
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
-                "sequence_length": self.packer.sequence_length,
-                "truncate": self.packer.truncate,
+                "sequence_length": self.sequence_length,
+                "truncate": self.truncate,
             }
         )
         return config
 
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.cls_token_id,
+            end_value=self.tokenizer.sep_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            truncate=self.truncate,
+            sequence_length=self.sequence_length,
+        )
+        self.built = True
+
     def call(self, x, y=None, sample_weight=None):
         x = convert_inputs_to_list_of_tensor_segments(x)
         x = [self.tokenizer(segment) for segment in x]

keras_nlp/models/bert/bert_tokenizer.py

@@ -74,7 +74,7 @@ class BertTokenizer(WordPieceTokenizer):
 
     def __init__(
         self,
-        vocabulary,
+        vocabulary=None,
         lowercase=False,
         **kwargs,
     ):
@@ -84,23 +84,32 @@ def __init__(
             **kwargs,
         )
 
-        # Check for necessary special tokens.
-        cls_token = "[CLS]"
-        sep_token = "[SEP]"
-        pad_token = "[PAD]"
-        mask_token = "[MASK]"
-        for token in [cls_token, pad_token, sep_token]:
-            if token not in self.get_vocabulary():
-                raise ValueError(
-                    f"Cannot find token `'{token}'` in the provided "
-                    f"`vocabulary`. Please provide `'{token}'` in your "
-                    "`vocabulary` or use a pretrained `vocabulary` name."
-                )
-
-        self.cls_token_id = self.token_to_id(cls_token)
-        self.sep_token_id = self.token_to_id(sep_token)
-        self.pad_token_id = self.token_to_id(pad_token)
-        self.mask_token_id = self.token_to_id(mask_token)
+    def set_vocabulary(self, vocabulary):
+        super().set_vocabulary(vocabulary)
+
+        if vocabulary is not None:
+            # Check for necessary special tokens.
+            cls_token = "[CLS]"
+            sep_token = "[SEP]"
+            pad_token = "[PAD]"
+            mask_token = "[MASK]"
+            for token in [cls_token, pad_token, sep_token]:
+                if token not in self.vocabulary:
+                    raise ValueError(
+                        f"Cannot find token `'{token}'` in the provided "
+                        f"`vocabulary`. Please provide `'{token}'` in your "
+                        "`vocabulary` or use a pretrained `vocabulary` name."
+                    )
+
+            self.cls_token_id = self.token_to_id(cls_token)
+            self.sep_token_id = self.token_to_id(sep_token)
+            self.pad_token_id = self.token_to_id(pad_token)
+            self.mask_token_id = self.token_to_id(mask_token)
+        else:
+            self.cls_token_id = None
+            self.sep_token_id = None
+            self.pad_token_id = None
+            self.mask_token_id = None
 
     @classproperty
     def presets(cls):

keras_nlp/models/distil_bert/distil_bert_classifier_test.py

@@ -42,7 +42,7 @@ def setUp(self):
             num_heads=2,
             hidden_dim=2,
             intermediate_dim=4,
-            max_sequence_length=self.preprocessor.packer.sequence_length,
+            max_sequence_length=self.preprocessor.sequence_length,
         )
         self.init_kwargs = {
             "preprocessor": self.preprocessor,

keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py

@@ -136,33 +136,30 @@ def __init__(
             truncate=truncate,
             **kwargs,
         )
-
+        self.mask_selection_rate = mask_selection_rate
+        self.mask_selection_length = mask_selection_length
+        self.mask_token_rate = mask_token_rate
+        self.random_token_rate = random_token_rate
+        self.masker = None
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # Defer masker creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
         self.masker = MaskedLMMaskGenerator(
-            mask_selection_rate=mask_selection_rate,
-            mask_selection_length=mask_selection_length,
-            mask_token_rate=mask_token_rate,
-            random_token_rate=random_token_rate,
-            vocabulary_size=tokenizer.vocabulary_size(),
-            mask_token_id=tokenizer.mask_token_id,
+            mask_selection_rate=self.mask_selection_rate,
+            mask_selection_length=self.mask_selection_length,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+            vocabulary_size=self.tokenizer.vocabulary_size(),
+            mask_token_id=self.tokenizer.mask_token_id,
             unselectable_token_ids=[
-                tokenizer.cls_token_id,
-                tokenizer.sep_token_id,
-                tokenizer.pad_token_id,
+                self.tokenizer.cls_token_id,
+                self.tokenizer.sep_token_id,
+                self.tokenizer.pad_token_id,
             ],
         )
 
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "mask_selection_rate": self.masker.mask_selection_rate,
-                "mask_selection_length": self.masker.mask_selection_length,
-                "mask_token_rate": self.masker.mask_token_rate,
-                "random_token_rate": self.masker.random_token_rate,
-            }
-        )
-        return config
-
     def call(self, x, y=None, sample_weight=None):
         if y is not None or sample_weight is not None:
             logging.warning(
@@ -183,3 +180,15 @@ def call(self, x, y=None, sample_weight=None):
         y = masker_outputs["mask_ids"]
         sample_weight = masker_outputs["mask_weights"]
         return pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config