Switch byte pair tokenizer to save_assets/load_assets (#1322)

As part of this work, we need to also switch all downstream
preprocessing layers to create packers on build (instead of on call).

keras_nlp/models/bart/bart_preprocessor.py

@@ -140,15 +140,23 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.tokenizer = tokenizer
+        self.encoder_sequence_length = encoder_sequence_length
+        self.decoder_sequence_length = decoder_sequence_length
+        self.encoder_packer = None
+        self.decoder_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.
 
         # TODO: Use `MultiSegmentPacker` instead of `StartEndPacker` once we
         # want to move to multi-segment packing and have improved
         # `MultiSegmentPacker`'s performance.
         self.encoder_packer = StartEndPacker(
-            start_value=tokenizer.start_token_id,
-            end_value=tokenizer.end_token_id,
-            pad_value=tokenizer.pad_token_id,
-            sequence_length=encoder_sequence_length,
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.encoder_sequence_length,
             return_padding_mask=True,
         )
 
@@ -161,19 +169,10 @@ def __init__(
             ],
             end_value=self.tokenizer.end_token_id,
             pad_value=self.tokenizer.pad_token_id,
-            sequence_length=decoder_sequence_length,
+            sequence_length=self.decoder_sequence_length,
             return_padding_mask=True,
         )
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "encoder_sequence_length": self.encoder_packer.sequence_length,
-                "decoder_sequence_length": self.decoder_packer.sequence_length,
-            }
-        )
-        return config
+        self.built = True
 
     def call(self, x, y=None, sample_weight=None):
         if not (
@@ -217,6 +216,16 @@ def call(self, x, y=None, sample_weight=None):
 
         return pack_x_y_sample_weight(x, y, sample_weight)
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "encoder_sequence_length": self.encoder_sequence_length,
+                "decoder_sequence_length": self.decoder_sequence_length,
+            }
+        )
+        return config
+
     @classproperty
     def tokenizer_cls(cls):
         return BartTokenizer

keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor.py

@@ -46,16 +46,6 @@ class BartSeq2SeqLMPreprocessor(BartPreprocessor):
         tokenizer: A `keras_nlp.models.BartTokenizer` instance.
         encoder_sequence_length: The length of the packed encoder inputs.
         decoder_sequence_length: The length of the packed decoder inputs.
-        truncate: string. The algorithm to truncate a list of batched segments
-            to fit within `sequence_length`. The value can be either
-            `round_robin` or `waterfall`:
-                - `"round_robin"`: Available space is assigned one token at a
-                    time in a round-robin fashion to the inputs that still need
-                    some, until the limit is reached.
-                - `"waterfall"`: The allocation of the budget is done using a
-                    "waterfall" algorithm that allocates quota in a
-                    left-to-right manner and fills up the buckets until we run
-                    out of budget. It supports an arbitrary number of segments.
 
     Call arguments:
         x: A dictionary with `encoder_text` and `decoder_text` as its keys.
@@ -139,7 +129,6 @@ def __init__(
         tokenizer,
         encoder_sequence_length,
         decoder_sequence_length,
-        truncate="round_robin",
         **kwargs
     ):
         # Since we truncate the last token from `decoder_token_ids`, we need to
@@ -156,16 +145,6 @@ def __init__(
         self._encoder_sequence_length = encoder_sequence_length
         self._decoder_sequence_length = decoder_sequence_length
 
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "encoder_sequence_length": self._encoder_sequence_length,
-                "decoder_sequence_length": self._decoder_sequence_length,
-            }
-        )
-        return config
-
     def call(self, x, y=None, sample_weight=None):
         if y is not None or sample_weight is not None:
             logging.warning(
@@ -191,10 +170,6 @@ def call(self, x, y=None, sample_weight=None):
         sample_weight = decoder_padding_mask[..., 1:]
         return pack_x_y_sample_weight(x, y, sample_weight)
 
-    @classproperty
-    def presets(cls):
-        return copy.deepcopy(backbone_presets)
-
     def generate_preprocess(
         self,
         x,
@@ -212,6 +187,9 @@ def generate_preprocess(
         the decoder sequence (as generation is expected to continue at the end
         of the inputted decoder prompt).
         """
+        if not self.built:
+            self.build(None)
+
         # If `sequence_length` is not provided, we use the default value.
         if sequence_length is None:
             sequence_length = self._decoder_sequence_length
@@ -262,6 +240,9 @@ def generate_postprocess(
         padding and start/end tokens, and then converting the integer sequence
         back to a string.
         """
+        if not self.built:
+            self.build(None)
+
         decoder_token_ids, decoder_padding_mask = (
             x["decoder_token_ids"],
             x["decoder_padding_mask"],
@@ -279,3 +260,17 @@ def generate_postprocess(
             decoder_token_ids, decoder_padding_mask
         )
         return self.tokenizer.detokenize(decoder_token_ids)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "encoder_sequence_length": self._encoder_sequence_length,
+                "decoder_sequence_length": self._decoder_sequence_length,
+            }
+        )
+        return config
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)

keras_nlp/models/bart/bart_seq_2_seq_lm_preprocessor_test.py

@@ -21,7 +21,7 @@
 from keras_nlp.tests.test_case import TestCase
 
 
-class BartPreprocessorTest(TestCase):
+class BartSeq2SeqLMPreprocessorTest(TestCase):
     def setUp(self):
         self.vocab = ["<s>", "<pad>", "</s>", "air", "Ġair", "plane", "Ġat"]
         self.vocab += ["port", "<mask>"]

keras_nlp/models/bart/bart_tokenizer.py

@@ -78,34 +78,45 @@ class BartTokenizer(BytePairTokenizer):
 
     def __init__(
         self,
-        vocabulary,
-        merges,
+        vocabulary=None,
+        merges=None,
         **kwargs,
     ):
-        # Special tokens.
-        start_token = "<s>"
-        pad_token = "<pad>"
-        end_token = "</s>"
+        self.start_token = "<s>"
+        self.pad_token = "<pad>"
+        self.end_token = "</s>"
 
         super().__init__(
             vocabulary=vocabulary,
             merges=merges,
-            unsplittable_tokens=[start_token, pad_token, end_token],
+            unsplittable_tokens=[
+                self.start_token,
+                self.pad_token,
+                self.end_token,
+            ],
             **kwargs,
         )
 
-        # Check whether special tokens are present in the vocabulary.
-        for token in [start_token, pad_token, end_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.start_token_id = self.token_to_id(start_token)
-        self.pad_token_id = self.token_to_id(pad_token)
-        self.end_token_id = self.token_to_id(end_token)
+    def set_vocabulary_and_merges(self, vocabulary, merges):
+        super().set_vocabulary_and_merges(vocabulary, merges)
+
+        if vocabulary is not None:
+            # Check for necessary special tokens.
+            for token in [self.start_token, self.pad_token, self.end_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.start_token_id = self.token_to_id(self.start_token)
+            self.pad_token_id = self.token_to_id(self.pad_token)
+            self.end_token_id = self.token_to_id(self.end_token)
+        else:
+            self.start_token_id = None
+            self.pad_token_id = None
+            self.end_token_id = None
 
     @classproperty
     def presets(cls):

keras_nlp/models/bert/bert_preprocessor.py

@@ -143,16 +143,6 @@ def __init__(
         self.truncate = truncate
         self.packer = None
 
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "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.
@@ -176,6 +166,16 @@ def call(self, x, y=None, sample_weight=None):
         }
         return pack_x_y_sample_weight(x, y, sample_weight)
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "truncate": self.truncate,
+            }
+        )
+        return config
+
     @classproperty
     def tokenizer_cls(cls):
         return BertTokenizer

keras_nlp/models/bert/bert_tokenizer.py

@@ -78,6 +78,10 @@ def __init__(
         lowercase=False,
         **kwargs,
     ):
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.pad_token = "[PAD]"
+        self.mask_token = "[MASK]"
         super().__init__(
             vocabulary=vocabulary,
             lowercase=lowercase,
@@ -89,22 +93,18 @@ def set_vocabulary(self, 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]:
+            for token in [self.cls_token, self.pad_token, self.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)
+            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

keras_nlp/models/distil_bert/distil_bert_tokenizer.py

@@ -76,6 +76,10 @@ def __init__(
         lowercase=False,
         **kwargs,
     ):
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.pad_token = "[PAD]"
+        self.mask_token = "[MASK]"
         super().__init__(
             vocabulary=vocabulary,
             lowercase=lowercase,
@@ -87,22 +91,18 @@ def set_vocabulary(self, 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.get_vocabulary():
+            for token in [self.cls_token, self.pad_token, self.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)
+            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

keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor.py

@@ -142,6 +142,9 @@ def generate_preprocess(
         the sequence (as generation is expected to continue at the end of the
         inputted prompt).
         """
+        if not self.built:
+            self.build(None)
+
         x = convert_inputs_to_list_of_tensor_segments(x)[0]
         x = self.tokenizer(x)
         token_ids, padding_mask = self.packer(
@@ -162,6 +165,9 @@ def generate_postprocess(
         padding and start/end tokens, and then converting the integer sequence
         back to a string.
         """
+        if not self.built:
+            self.build(None)
+
         token_ids, padding_mask = x["token_ids"], x["padding_mask"]
         token_ids = ops.convert_to_numpy(token_ids)
         padding_mask = ops.convert_to_numpy(padding_mask)

keras_nlp/models/gpt2/gpt2_causal_lm_test.py

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