keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/layers/preprocessing/random_swap.py

@@ -21,6 +21,7 @@ from keras_hub.src.layers.preprocessing.preprocessing_layer import (
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
 from keras_hub.src.utils.tensor_utils import is_int_dtype
 from keras_hub.src.utils.tensor_utils import is_string_dtype
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
@@ -67,55 +68,55 @@ class RandomSwap(PreprocessingLayer):
 
     Word level usage.
     >>> keras.utils.set_random_seed(1337)
-    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
-    >>> augmenter=keras_hub.layers.RandomSwap(rate=0.4, seed=42)
-    >>> augmented=augmenter(inputs)
-    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string,
-    numpy=array([b'like I Hey', b'and Keras Tensorflow'], dtype=object)>
+    >>> x = ["Hey I like", "Keras and Tensorflow"]
+    >>> x = list(map(lambda x: x.split(), x))
+    >>> augmenter = keras_hub.layers.RandomSwap(rate=0.4, seed=42)
+    >>> y = augmenter(x)
+    >>> list(map(lambda y: " ".join(y), y))
+    ['like I Hey', 'and Keras Tensorflow']
 
     Character level usage.
     >>> keras.utils.set_random_seed(1337)
-    >>> inputs=tf.strings.unicode_split(["Hey Dude", "Speed Up"], "UTF-8")
-    >>> augmenter=keras_hub.layers.RandomSwap(rate=0.4, seed=42)
-    >>> augmented=augmenter(inputs)
-    >>> tf.strings.reduce_join(augmented, axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string,
-    numpy=array([b'deD yuHe', b'SUede pp'], dtype=object)>
+    >>> x = ["Hey Dude", "Speed Up"]
+    >>> x = list(map(lambda x: list(x), x))
+    >>> augmenter = keras_hub.layers.RandomSwap(rate=0.4, seed=42)
+    >>> y = augmenter(x)
+    >>> list(map(lambda y: "".join(y), y))
+    ['deD yuHe', 'SUede pp']
 
     Usage with skip_list.
     >>> keras.utils.set_random_seed(1337)
-    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
-    >>> augmenter=keras_hub.layers.RandomSwap(rate=0.4,
+    >>> x = ["Hey I like", "Keras and Tensorflow"]
+    >>> x = list(map(lambda x: x.split(), x))
+    >>> augmenter = keras_hub.layers.RandomSwap(rate=0.4,
     ...     skip_list=["Keras"], seed=42)
-    >>> augmented=augmenter(inputs)
-    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string,
-    numpy=array([b'like I Hey', b'Keras and Tensorflow'], dtype=object)>
+    >>> y = augmenter(x)
+    >>> list(map(lambda y: " ".join(y), y))
+    ['like I Hey', 'Keras and Tensorflow']
 
     Usage with skip_fn.
     >>> def skip_fn(word):
     ...     return tf.strings.regex_full_match(word, r"[I, a].*")
     >>> keras.utils.set_random_seed(1337)
-    >>> inputs=tf.strings.split(["Hey I like", "Keras and Tensorflow"])
-    >>> augmenter=keras_hub.layers.RandomSwap(rate=0.9, max_swaps=3,
+    >>> x = ["Hey I like", "Keras and Tensorflow"]
+    >>> x = list(map(lambda x: x.split(), x))
+    >>> augmenter = keras_hub.layers.RandomSwap(rate=0.9, max_swaps=3,
     ...     skip_fn=skip_fn, seed=11)
-    >>> augmented=augmenter(inputs)
-    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string,
-    numpy=array([b'like I Hey', b'Keras and Tensorflow'], dtype=object)>
+    >>> y = augmenter(x)
+    >>> list(map(lambda y: " ".join(y), y))
+    ['like I Hey', 'Keras and Tensorflow']
 
     Usage with skip_py_fn.
     >>> def skip_py_fn(word):
     ...     return len(word) < 4
     >>> keras.utils.set_random_seed(1337)
-    >>> inputs=tf.strings.split(["He was drifting along", "With the wind"])
-    >>> augmenter=keras_hub.layers.RandomSwap(rate=0.8, max_swaps=2,
+    >>> x = ["He was drifting along", "With the wind"]
+    >>> x = list(map(lambda x: x.split(), x))
+    >>> augmenter = keras_hub.layers.RandomSwap(rate=0.8, max_swaps=2,
     ...     skip_py_fn=skip_py_fn, seed=15)
-    >>> augmented=augmenter(inputs)
-    >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'He was along drifting',
-    b'wind the With'], dtype=object)>
+    >>> y = augmenter(x)
+    >>> list(map(lambda y: " ".join(y), y))
+    ['He was along drifting', 'wind the With']
     """
 
     def __init__(
@@ -166,8 +167,9 @@ class RandomSwap(PreprocessingLayer):
                 default_value=False,
             )
 
+    @preprocessing_function
     def call(self, inputs):
-        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
 
         skip_masks = None
         if self.skip_list:

keras_hub/src/layers/preprocessing/resizing_image_converter.py

@@ -0,0 +1,101 @@
+# Copyright 2024 The KerasHub 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 keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.layers.ResizingImageConverter")
+class ResizingImageConverter(ImageConverter):
+    """An `ImageConverter` that simply resizes the input image.
+
+    The `ResizingImageConverter` is a subclass of `ImageConverter` for models
+    that simply need to resize image tensors before using them for modeling.
+    The layer will take as input a raw image tensor (batched or unbatched) in the
+    channels last or channels first format, and output a resize tensor.
+
+    Args:
+        height: Integer, the height of the output shape.
+        width: Integer, the width of the output shape.
+        crop_to_aspect_ratio: If `True`, resize the images without aspect
+            ratio distortion. When the original aspect ratio differs
+            from the target aspect ratio, the output image will be
+            cropped so as to return the
+            largest possible window in the image (of size `(height, width)`)
+            that matches the target aspect ratio. By default
+            (`crop_to_aspect_ratio=False`), aspect ratio may not be preserved.
+        interpolation: String, the interpolation method.
+            Supports `"bilinear"`, `"nearest"`, `"bicubic"`,
+            `"lanczos3"`, `"lanczos5"`. Defaults to `"bilinear"`.
+        data_format: String, either `"channels_last"` or `"channels_first"`.
+            The ordering of the dimensions in the inputs. `"channels_last"`
+            corresponds to inputs with shape `(batch, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, width)`. It defaults to the
+            `image_data_format` value found in your Keras config file at
+            `~/.keras/keras.json`. If you never set it, then it will be
+            `"channels_last"`.
+
+    Examples:
+    ```python
+    # Resize images for `"pali_gemma_3b_224"`.
+    converter = keras_hub.layers.ImageConverter.from_preset("pali_gemma_3b_224")
+    converter(np.ones(2, 512, 512, 3)) # Output shape: (2, 224, 224, 3)
+    # Resize images for `"pali_gemma_3b_224"`.
+    converter = keras_hub.layers.ImageConverter.from_preset("pali_gemma_3b_448")
+    converter(np.ones(2, 512, 512, 3)) # Output shape: (2, 448, 448, 3)
+    ```
+    """
+
+    def __init__(
+        self,
+        height,
+        width,
+        crop_to_aspect_ratio=True,
+        interpolation="bilinear",
+        data_format=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        # By default, we just do a simple resize. Any model can subclass this
+        # layer for preprocessing of a raw image to a model image input.
+        self.resizing = keras.layers.Resizing(
+            height=height,
+            width=width,
+            crop_to_aspect_ratio=crop_to_aspect_ratio,
+            interpolation=interpolation,
+            data_format=data_format,
+        )
+
+    def image_size(self):
+        """Returns the preprocessed size of a single image."""
+        return (self.resizing.height, self.resizing.width)
+
+    @preprocessing_function
+    def call(self, inputs):
+        return self.resizing(inputs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "height": self.resizing.height,
+                "width": self.resizing.width,
+                "interpolation": self.resizing.interpolation,
+                "crop_to_aspect_ratio": self.resizing.crop_to_aspect_ratio,
+            }
+        )
+        return config

keras_hub/src/layers/preprocessing/start_end_packer.py

@@ -18,6 +18,7 @@ from keras_hub.src.layers.preprocessing.preprocessing_layer import (
     PreprocessingLayer,
 )
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
@@ -154,6 +155,7 @@ class StartEndPacker(PreprocessingLayer):
         self.pad_value = pad_value
         self.return_padding_mask = return_padding_mask
 
+    @preprocessing_function
     def call(
         self,
         inputs,
@@ -161,8 +163,7 @@ class StartEndPacker(PreprocessingLayer):
         add_start_value=True,
         add_end_value=True,
     ):
-        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
-
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
         x = inputs  # Intermediate result.
 
         batch_size = tf.shape(x)[0]

keras_hub/src/models/albert/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.albert.albert_backbone import AlbertBackbone
 from keras_hub.src.models.albert.albert_presets import backbone_presets
-from keras_hub.src.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (AlbertBackbone, AlbertTokenizer))
+register_presets(backbone_presets, AlbertBackbone)

keras_hub/src/models/albert/albert_masked_lm_preprocessor.py

@@ -12,18 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras
-from absl import logging
-
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
-    MaskedLMMaskGenerator,
-)
-from keras_hub.src.models.albert.albert_preprocessor import AlbertPreprocessor
+from keras_hub.src.models.albert.albert_backbone import AlbertBackbone
+from keras_hub.src.models.albert.albert_tokenizer import AlbertTokenizer
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
 
 
 @keras_hub_export("keras_hub.models.AlbertMaskedLMPreprocessor")
-class AlbertMaskedLMPreprocessor(AlbertPreprocessor):
+class AlbertMaskedLMPreprocessor(MaskedLMPreprocessor):
     """ALBERT preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -114,81 +110,5 @@ class AlbertMaskedLMPreprocessor(AlbertPreprocessor):
     ```
     """
 
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        mask_selection_rate=0.15,
-        mask_selection_length=96,
-        mask_token_rate=0.8,
-        random_token_rate=0.1,
-        **kwargs,
-    ):
-        super().__init__(
-            tokenizer,
-            sequence_length=sequence_length,
-            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=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=[
-                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.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
-
-    def call(self, x, y=None, sample_weight=None):
-        if y is not None or sample_weight is not None:
-            logging.warning(
-                f"{self.__class__.__name__} generates `y` and `sample_weight` "
-                "based on your input data, but your data already contains `y` "
-                "or `sample_weight`. Your `y` and `sample_weight` will be "
-                "ignored."
-            )
-
-        x = super().call(x)
-        token_ids, segment_ids, padding_mask = (
-            x["token_ids"],
-            x["segment_ids"],
-            x["padding_mask"],
-        )
-        masker_outputs = self.masker(token_ids)
-        x = {
-            "token_ids": masker_outputs["token_ids"],
-            "segment_ids": segment_ids,
-            "padding_mask": padding_mask,
-            "mask_positions": masker_outputs["mask_positions"],
-        }
-        y = masker_outputs["mask_ids"]
-        sample_weight = masker_outputs["mask_weights"]
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+    backbone_cls = AlbertBackbone
+    tokenizer_cls = AlbertTokenizer

keras_hub/src/models/albert/{albert_classifier.py → albert_text_classifier.py}

@@ -19,12 +19,19 @@ from keras_hub.src.models.albert.albert_backbone import AlbertBackbone
 from keras_hub.src.models.albert.albert_backbone import (
     albert_kernel_initializer,
 )
-from keras_hub.src.models.albert.albert_preprocessor import AlbertPreprocessor
-from keras_hub.src.models.classifier import Classifier
+from keras_hub.src.models.albert.albert_text_classifier_preprocessor import (
+    AlbertTextClassifierPreprocessor,
+)
+from keras_hub.src.models.text_classifier import TextClassifier
 
 
-@keras_hub_export("keras_hub.models.AlbertClassifier")
-class AlbertClassifier(Classifier):
+@keras_hub_export(
+    [
+        "keras_hub.models.AlbertTextClassifier",
+        "keras_hub.models.AlbertClassifier",
+    ]
+)
+class AlbertTextClassifier(TextClassifier):
     """An end-to-end ALBERT model for classification tasks
 
     This model attaches a classification head to a `keras_hub.model.AlbertBackbone`
@@ -43,7 +50,7 @@ class AlbertClassifier(Classifier):
     Args:
         backbone: A `keras_hub.models.AlertBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.AlbertPreprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.AlbertTextClassifierPreprocessor` or `None`. If
             `None`, this model will not apply preprocessing, and inputs should
             be preprocessed before calling the model.
         activation: Optional `str` or callable. The
@@ -61,7 +68,7 @@ class AlbertClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier.
-    classifier = keras_hub.models.AlbertClassifier.from_preset(
+    classifier = keras_hub.models.AlbertTextClassifier.from_preset(
         "albert_base_en_uncased",
         num_classes=4,
     )
@@ -90,7 +97,7 @@ class AlbertClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.AlbertClassifier.from_preset(
+    classifier = keras_hub.models.AlbertTextClassifier.from_preset(
         "albert_base_en_uncased",
         num_classes=4,
         preprocessor=None,
@@ -123,7 +130,7 @@ class AlbertClassifier(Classifier):
     tokenizer = keras_hub.models.AlbertTokenizer(
         proto=bytes_io.getvalue(),
     )
-    preprocessor = keras_hub.models.AlbertPreprocessor(
+    preprocessor = keras_hub.models.AlbertTextClassifierPreprocessor(
         tokenizer=tokenizer,
         sequence_length=128,
     )
@@ -136,7 +143,7 @@ class AlbertClassifier(Classifier):
         intermediate_dim=512,
         max_sequence_length=128,
     )
-    classifier = keras_hub.models.AlbertClassifier(
+    classifier = keras_hub.models.AlbertTextClassifier(
         backbone=backbone,
         preprocessor=preprocessor,
         num_classes=4,
@@ -146,7 +153,7 @@ class AlbertClassifier(Classifier):
     """
 
     backbone_cls = AlbertBackbone
-    preprocessor_cls = AlbertPreprocessor
+    preprocessor_cls = AlbertTextClassifierPreprocessor
 
     def __init__(
         self,
@@ -189,6 +196,23 @@ class AlbertClassifier(Classifier):
         self.activation = keras.activations.get(activation)
         self.dropout = dropout
 
+    def compile(
+        self,
+        optimizer="auto",
+        loss="auto",
+        *,
+        metrics="auto",
+        **kwargs,
+    ):
+        if optimizer == "auto":
+            optimizer = keras.optimizers.Adam(1e-5)
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            metrics=metrics,
+            **kwargs,
+        )
+
     def get_config(self):
         config = super().get_config()
         config.update(

keras_hub/src/models/albert/{albert_preprocessor.py → albert_text_classifier_preprocessor.py}

@@ -12,21 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras
-
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.multi_segment_packer import (
-    MultiSegmentPacker,
-)
+from keras_hub.src.models.albert.albert_backbone import AlbertBackbone
 from keras_hub.src.models.albert.albert_tokenizer import AlbertTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
+from keras_hub.src.models.text_classifier_preprocessor import (
+    TextClassifierPreprocessor,
 )
 
 
-@keras_hub_export("keras_hub.models.AlbertPreprocessor")
-class AlbertPreprocessor(Preprocessor):
+@keras_hub_export(
+    [
+        "keras_hub.models.AlbertTextClassifierPreprocessor",
+        "keras_hub.models.AlbertPreprocessor",
+    ]
+)
+class AlbertTextClassifierPreprocessor(TextClassifierPreprocessor):
     """An ALBERT preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
@@ -72,7 +72,7 @@ class AlbertPreprocessor(Preprocessor):
     Examples:
     Directly calling the layer on data.
     ```python
-    preprocessor = keras_hub.models.AlbertPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "albert_base_en_uncased"
     )
 
@@ -109,13 +109,13 @@ class AlbertPreprocessor(Preprocessor):
     tokenizer = keras_hub.models.AlbertTokenizer(
         proto=bytes_io.getvalue(),
     )
-    preprocessor = keras_hub.models.AlbertPreprocessor(tokenizer)
+    preprocessor = keras_hub.models.AlbertTextClassifierPreprocessor(tokenizer)
     preprocessor("The quick brown fox jumped.")
     ```
 
     Mapping with `tf.data.Dataset`.
     ```python
-    preprocessor = keras_hub.models.AlbertPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "albert_base_en_uncased"
     )
 
@@ -146,61 +146,5 @@ class AlbertPreprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = AlbertBackbone
     tokenizer_cls = AlbertTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.tokenizer = tokenizer
-        self.packer = None
-        self.truncate = truncate
-        self.sequence_length = sequence_length
-
-    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 get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "truncate": self.truncate,
-            }
-        )
-        return config
-
-    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]
-        token_ids, segment_ids = self.packer(x)
-        x = {
-            "token_ids": token_ids,
-            "segment_ids": segment_ids,
-            "padding_mask": token_ids != self.tokenizer.pad_token_id,
-        }
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    @property
-    def sequence_length(self):
-        """The padded length of model input sequences."""
-        return self._sequence_length
-
-    @sequence_length.setter
-    def sequence_length(self, value):
-        self._sequence_length = value
-        if self.packer is not None:
-            self.packer.sequence_length = value

keras_hub/src/models/albert/albert_tokenizer.py

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.albert.albert_backbone import AlbertBackbone
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
 
 
-@keras_hub_export("keras_hub.models.AlbertTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.AlbertTokenizer",
+        "keras_hub.models.AlbertTokenizer",
+    ]
+)
 class AlbertTokenizer(SentencePieceTokenizer):
     """ALBERT tokenizer layer based on SentencePiece.
 
@@ -28,10 +34,6 @@ class AlbertTokenizer(SentencePieceTokenizer):
     ALBERT models and provides a `from_preset()` method to automatically
     download a matching vocabulary for a ALBERT preset.
 
-    This tokenizer does not provide truncation or padding of inputs. It can be
-    combined with a `keras_hub.models.AlbertPreprocessor` layer for input
-    packing.
-
     If input is a batch of strings (rank > 0), the layer will output a
     `tf.RaggedTensor` where the last dimension of the output is ragged.
 
@@ -84,36 +86,15 @@ class AlbertTokenizer(SentencePieceTokenizer):
     ```
     """
 
-    def __init__(self, proto, **kwargs):
-        self.cls_token = "[CLS]"
-        self.sep_token = "[SEP]"
-        self.pad_token = "<pad>"
-        self.mask_token = "[MASK]"
+    backbone_cls = AlbertBackbone
 
+    def __init__(self, proto, **kwargs):
+        self._add_special_token("[CLS]", "cls_token")
+        self._add_special_token("[SEP]", "sep_token")
+        self._add_special_token("<pad>", "pad_token")
+        self._add_special_token("[MASK]", "mask_token")
+        # Also add `tokenizer.start_token` and `tokenizer.end_token` for
+        # compatibility with other tokenizers.
+        self._add_special_token("[CLS]", "start_token")
+        self._add_special_token("[SEP]", "end_token")
         super().__init__(proto=proto, **kwargs)
-
-    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