Convert SentencePieceTokenizer and associated models to new assets pa…

…radigm (#1323)

* Convert SentencePiece tokenizer to save_assets/load_assets

* Convert albert to new assets paradigm

* Convert DebertaV3 to new assets paradigm

* Fix formatting issues

* Convert FNet to new assets paradigm

* Convert XLMRoberta to new assets paradigm

* Convert T5 Tokenizer to new assets paradigm

* Fix sentencepiece tokenizer config test

* Change set_vocabulary to set_proto

* Change proto to raw proto

* Change to proto_bytes

keras_nlp/models/albert/albert_classifier_test.py

@@ -42,7 +42,7 @@ def setUp(self):
             hidden_dim=2,
             embedding_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/albert/albert_masked_lm_preprocessor.py

@@ -131,29 +131,38 @@ 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,
+                "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/albert/albert_masked_lm_test.py

@@ -50,7 +50,7 @@ def setUp(self):
             hidden_dim=2,
             embedding_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/albert/albert_preprocessor.py

@@ -158,20 +158,28 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.tokenizer = tokenizer
+        self.truncate = truncate
+        self.sequence_length = sequence_length
+        self.packer = None
+
+    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=truncate,
-            sequence_length=sequence_length,
+            truncate=self.truncate,
+            sequence_length=self.sequence_length,
         )
+        self.built = True
 
     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

keras_nlp/models/albert/albert_tokenizer.py

@@ -87,25 +87,38 @@ class AlbertTokenizer(SentencePieceTokenizer):
     """
 
     def __init__(self, proto, **kwargs):
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.pad_token = "<pad>"
+        self.mask_token = "[MASK]"
+
         super().__init__(proto=proto, **kwargs)
 
-        # Check for necessary special tokens.
-        cls_token = "[CLS]"
-        sep_token = "[SEP]"
-        pad_token = "<pad>"
-        mask_token = "[MASK]"
-        for token in [cls_token, sep_token, pad_token, mask_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_proto(self, proto):
+        super().set_proto(proto)
+        if proto is not None:
+            for token in [
+                self.cls_token,
+                self.sep_token,
+                self.pad_token,
+                self.mask_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(self.cls_token)
+            self.sep_token_id = self.token_to_id(self.sep_token)
+            self.pad_token_id = self.token_to_id(self.pad_token)
+            self.mask_token_id = self.token_to_id(self.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/deberta_v3/deberta_v3_classifier_test.py

@@ -45,7 +45,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/deberta_v3/deberta_v3_masked_lm_preprocessor.py

@@ -133,28 +133,38 @@ def __init__(
             **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,
+                "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/deberta_v3/deberta_v3_masked_lm_test.py

@@ -48,7 +48,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/deberta_v3/deberta_v3_preprocessor.py

@@ -156,20 +156,28 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.tokenizer = tokenizer
+        self.truncate = truncate
+        self.sequence_length = sequence_length
+        self.packer = None
+
+    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=truncate,
-            sequence_length=sequence_length,
+            truncate=self.truncate,
+            sequence_length=self.sequence_length,
         )
+        self.built = True
 
     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

keras_nlp/models/deberta_v3/deberta_v3_tokenizer.py

@@ -93,33 +93,38 @@ class DebertaV3Tokenizer(SentencePieceTokenizer):
     """
 
     def __init__(self, proto, **kwargs):
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.pad_token = "[PAD]"
+        self.mask_token = "[MASK]"
+
         super().__init__(proto=proto, **kwargs)
 
-        # Check for necessary special tokens.
-        cls_token = "[CLS]"
-        sep_token = "[SEP]"
-        pad_token = "[PAD]"
-        mask_token = "[MASK]"
-
-        # We do not throw an error if `mask_token` is not present in the
-        # vocabulary.
-        for token in [cls_token, pad_token, sep_token]:
-            if token not in super().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)
-        # If the mask token is not in the vocabulary, add it to the end of the
-        # vocabulary.
-        if mask_token in super().get_vocabulary():
-            self.mask_token_id = super().token_to_id(mask_token)
+    def set_proto(self, proto):
+        super().set_proto(proto)
+        if proto is not None:
+            for token in [self.cls_token, self.pad_token, self.sep_token]:
+                if token not in super().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(self.cls_token)
+            self.sep_token_id = self.token_to_id(self.sep_token)
+            self.pad_token_id = self.token_to_id(self.pad_token)
+            # If the mask token is not in the vocabulary, add it to the end of the
+            # vocabulary.
+            if self.mask_token in super().get_vocabulary():
+                self.mask_token_id = super().token_to_id(self.mask_token)
+            else:
+                self.mask_token_id = super().vocabulary_size()
         else:
-            self.mask_token_id = super().vocabulary_size()
+            self.cls_token_id = None
+            self.sep_token_id = None
+            self.pad_token_id = None
+            self.mask_token_id = None
 
     def vocabulary_size(self):
         sentence_piece_size = super().vocabulary_size()

keras_nlp/models/f_net/f_net_classifier_test.py

@@ -40,7 +40,7 @@ def setUp(self):
             num_layers=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/f_net/f_net_masked_lm_preprocessor.py

@@ -136,29 +136,38 @@ 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,
+                "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/f_net/f_net_masked_lm_test.py

@@ -47,7 +47,7 @@ def setUp(self):
             num_layers=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,