keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.15.0.dev20240911134614__py3-none-any.whl

keras_hub/src/models/deberta_v3/{deberta_v3_preprocessor.py → deberta_v3_text_classifier_preprocessor.py}

@@ -12,24 +12,28 @@
 # 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.deberta_v3.deberta_v3_backbone import (
+    DebertaV3Backbone,
 )
 from keras_hub.src.models.deberta_v3.deberta_v3_tokenizer import (
     DebertaV3Tokenizer,
 )
-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,
 )
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
-@keras_hub_export("keras_hub.models.DebertaV3Preprocessor")
-class DebertaV3Preprocessor(Preprocessor):
+@keras_hub_export(
+    [
+        "keras_hub.models.DebertaV3TextClassifierPreprocessor",
+        "keras_hub.models.DebertaV3Preprocessor",
+    ]
+)
+class DebertaV3TextClassifierPreprocessor(TextClassifierPreprocessor):
     """A DeBERTa preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
@@ -74,7 +78,7 @@ class DebertaV3Preprocessor(Preprocessor):
     Examples:
     Directly calling the layer on data.
     ```python
-    preprocessor = keras_hub.models.DebertaV3Preprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "deberta_v3_base_en"
     )
 
@@ -110,13 +114,15 @@ class DebertaV3Preprocessor(Preprocessor):
     tokenizer = keras_hub.models.DebertaV3Tokenizer(
         proto=bytes_io.getvalue(),
     )
-    preprocessor = keras_hub.models.DebertaV3Preprocessor(tokenizer)
+    preprocessor = keras_hub.models.DebertaV3TextClassifierPreprocessor(
+        tokenizer
+    )
     preprocessor("The quick brown fox jumped.")
     ```
 
     Mapping with `tf.data.Dataset`.
     ```python
-    preprocessor = keras_hub.models.DebertaV3Preprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "deberta_v3_base_en"
     )
 
@@ -147,60 +153,13 @@ class DebertaV3Preprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = DebertaV3Backbone
     tokenizer_cls = DebertaV3Tokenizer
 
-    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
-
+    @preprocessing_function
     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, _ = self.packer(x)
-        x = {
-            "token_ids": token_ids,
-            "padding_mask": token_ids != self.tokenizer.pad_token_id,
-        }
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        # Backbone has no segment ID input.
+        del x["segment_ids"]
         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/deberta_v3/deberta_v3_tokenizer.py

@@ -14,6 +14,9 @@
 
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.deberta_v3.deberta_v3_backbone import (
+    DebertaV3Backbone,
+)
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
@@ -24,7 +27,12 @@ except ImportError:
     tf = None
 
 
-@keras_hub_export("keras_hub.models.DebertaV3Tokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.DebertaV3Tokenizer",
+        "keras_hub.models.DebertaV3Tokenizer",
+    ]
+)
 class DebertaV3Tokenizer(SentencePieceTokenizer):
     """DeBERTa tokenizer layer based on SentencePiece.
 
@@ -34,10 +42,6 @@ class DebertaV3Tokenizer(SentencePieceTokenizer):
     DeBERTa models and provides a `from_preset()` method to automatically
     download a matching vocabulary for a DeBERTa preset.
 
-    This tokenizer does not provide truncation or padding of inputs. It can be
-    combined with a `keras_hub.models.DebertaV3Preprocessor` 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.
 
@@ -94,38 +98,37 @@ class DebertaV3Tokenizer(SentencePieceTokenizer):
     ```
     """
 
+    backbone_cls = DebertaV3Backbone
+
     def __init__(self, proto, **kwargs):
-        self.cls_token = "[CLS]"
-        self.sep_token = "[SEP]"
-        self.pad_token = "[PAD]"
+        self._add_special_token("[CLS]", "cls_token")
+        self._add_special_token("[SEP]", "sep_token")
+        self._add_special_token("[PAD]", "pad_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")
+        # Handle mask separately as it's not always in the vocab.
         self.mask_token = "[MASK]"
-
+        self.mask_token_id = None
         super().__init__(proto=proto, **kwargs)
 
+    @property
+    def special_tokens(self):
+        return super().special_tokens + [self.mask_token]
+
+    @property
+    def special_token_ids(self):
+        return super().special_token_ids + [self.mask_token_id]
+
     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.cls_token_id = None
-            self.sep_token_id = None
-            self.pad_token_id = None
             self.mask_token_id = None
 
     def vocabulary_size(self):
@@ -136,6 +139,8 @@ class DebertaV3Tokenizer(SentencePieceTokenizer):
 
     def get_vocabulary(self):
         sentence_piece_vocabulary = super().get_vocabulary()
+        if self.mask_token_id is None:
+            return sentence_piece_vocabulary
         if self.mask_token_id < super().vocabulary_size():
             return sentence_piece_vocabulary
         return sentence_piece_vocabulary + ["[MASK]"]

keras_hub/src/models/densenet/densenet_backbone.py

@@ -14,14 +14,13 @@
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.feature_pyramid_backbone import FeaturePyramidBackbone
 
-BN_AXIS = 3
 BN_EPSILON = 1.001e-5
 
 
 @keras_hub_export("keras_hub.models.DenseNetBackbone")
-class DenseNetBackbone(Backbone):
+class DenseNetBackbone(FeaturePyramidBackbone):
     """Instantiates the DenseNet architecture.
 
     This class implements a DenseNet backbone as described in
@@ -35,7 +34,7 @@ class DenseNetBackbone(Backbone):
         include_rescaling: bool, whether to rescale the inputs. If set
             to `True`, inputs will be passed through a `Rescaling(1/255.0)`
             layer. Defaults to `True`.
-        image_shape: optional shape tuple, defaults to (224, 224, 3).
+        image_shape: optional shape tuple, defaults to (None, None, 3).
         compression_ratio: float, compression rate at transition layers,
             defaults to 0.5.
         growth_rate: int, number of filters added by each dense block,
@@ -62,12 +61,14 @@ class DenseNetBackbone(Backbone):
         self,
         stackwise_num_repeats,
         include_rescaling=True,
-        image_shape=(224, 224, 3),
+        image_shape=(None, None, 3),
         compression_ratio=0.5,
         growth_rate=32,
         **kwargs,
     ):
         # === Functional Model ===
+        data_format = keras.config.image_data_format()
+        channel_axis = -1 if data_format == "channels_last" else 1
         image_input = keras.layers.Input(shape=image_shape)
 
         x = image_input
@@ -75,37 +76,47 @@ class DenseNetBackbone(Backbone):
             x = keras.layers.Rescaling(1 / 255.0)(x)
 
         x = keras.layers.Conv2D(
-            64, 7, strides=2, use_bias=False, padding="same", name="conv1_conv"
+            64,
+            7,
+            strides=2,
+            use_bias=False,
+            padding="same",
+            data_format=data_format,
+            name="conv1_conv",
         )(x)
         x = keras.layers.BatchNormalization(
-            axis=BN_AXIS, epsilon=BN_EPSILON, name="conv1_bn"
+            axis=channel_axis, epsilon=BN_EPSILON, name="conv1_bn"
         )(x)
         x = keras.layers.Activation("relu", name="conv1_relu")(x)
         x = keras.layers.MaxPooling2D(
-            3, strides=2, padding="same", name="pool1"
+            3, strides=2, padding="same", data_format=data_format, name="pool1"
         )(x)
 
+        pyramid_outputs = {}
         for stack_index in range(len(stackwise_num_repeats) - 1):
             index = stack_index + 2
             x = apply_dense_block(
                 x,
+                channel_axis,
                 stackwise_num_repeats[stack_index],
                 growth_rate,
                 name=f"conv{index}",
             )
+            pyramid_outputs[f"P{index}"] = x
             x = apply_transition_block(
-                x, compression_ratio, name=f"pool{index}"
+                x, channel_axis, compression_ratio, name=f"pool{index}"
             )
 
         x = apply_dense_block(
             x,
+            channel_axis,
             stackwise_num_repeats[-1],
             growth_rate,
             name=f"conv{len(stackwise_num_repeats) + 1}",
         )
-
+        pyramid_outputs[f"P{len(stackwise_num_repeats) +1}"] = x
         x = keras.layers.BatchNormalization(
-            axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
+            axis=channel_axis, epsilon=BN_EPSILON, name="bn"
         )(x)
         x = keras.layers.Activation("relu", name="relu")(x)
 
@@ -117,6 +128,7 @@ class DenseNetBackbone(Backbone):
         self.compression_ratio = compression_ratio
         self.growth_rate = growth_rate
         self.image_shape = image_shape
+        self.pyramid_outputs = pyramid_outputs
 
     def get_config(self):
         config = super().get_config()
@@ -132,7 +144,7 @@ class DenseNetBackbone(Backbone):
         return config
 
 
-def apply_dense_block(x, num_repeats, growth_rate, name=None):
+def apply_dense_block(x, channel_axis, num_repeats, growth_rate, name=None):
     """A dense block.
 
     Args:
@@ -145,11 +157,13 @@ def apply_dense_block(x, num_repeats, growth_rate, name=None):
         name = f"dense_block_{keras.backend.get_uid('dense_block')}"
 
     for i in range(num_repeats):
-        x = apply_conv_block(x, growth_rate, name=f"{name}_block_{i}")
+        x = apply_conv_block(
+            x, channel_axis, growth_rate, name=f"{name}_block_{i}"
+        )
     return x
 
 
-def apply_transition_block(x, compression_ratio, name=None):
+def apply_transition_block(x, channel_axis, compression_ratio, name=None):
     """A transition block.
 
     Args:
@@ -157,24 +171,28 @@ def apply_transition_block(x, compression_ratio, name=None):
       compression_ratio: float, compression rate at transition layers.
       name: string, block label.
     """
+    data_format = keras.config.image_data_format()
     if name is None:
         name = f"transition_block_{keras.backend.get_uid('transition_block')}"
 
     x = keras.layers.BatchNormalization(
-        axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_bn"
+        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_bn"
     )(x)
     x = keras.layers.Activation("relu", name=f"{name}_relu")(x)
     x = keras.layers.Conv2D(
-        int(x.shape[BN_AXIS] * compression_ratio),
+        int(x.shape[channel_axis] * compression_ratio),
         1,
         use_bias=False,
+        data_format=data_format,
         name=f"{name}_conv",
     )(x)
-    x = keras.layers.AveragePooling2D(2, strides=2, name=f"{name}_pool")(x)
+    x = keras.layers.AveragePooling2D(
+        2, strides=2, data_format=data_format, name=f"{name}_pool"
+    )(x)
     return x
 
 
-def apply_conv_block(x, growth_rate, name=None):
+def apply_conv_block(x, channel_axis, growth_rate, name=None):
     """A building block for a dense block.
 
     Args:
@@ -182,19 +200,24 @@ def apply_conv_block(x, growth_rate, name=None):
       growth_rate: int, number of filters added by each dense block.
       name: string, block label.
     """
+    data_format = keras.config.image_data_format()
     if name is None:
         name = f"conv_block_{keras.backend.get_uid('conv_block')}"
 
     shortcut = x
     x = keras.layers.BatchNormalization(
-        axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_0_bn"
+        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_0_bn"
     )(x)
     x = keras.layers.Activation("relu", name=f"{name}_0_relu")(x)
     x = keras.layers.Conv2D(
-        4 * growth_rate, 1, use_bias=False, name=f"{name}_1_conv"
+        4 * growth_rate,
+        1,
+        use_bias=False,
+        data_format=data_format,
+        name=f"{name}_1_conv",
     )(x)
     x = keras.layers.BatchNormalization(
-        axis=BN_AXIS, epsilon=BN_EPSILON, name=f"{name}_1_bn"
+        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_1_bn"
     )(x)
     x = keras.layers.Activation("relu", name=f"{name}_1_relu")(x)
     x = keras.layers.Conv2D(
@@ -202,9 +225,10 @@ def apply_conv_block(x, growth_rate, name=None):
         3,
         padding="same",
         use_bias=False,
+        data_format=data_format,
         name=f"{name}_2_conv",
     )(x)
-    x = keras.layers.Concatenate(axis=BN_AXIS, name=f"{name}_concat")(
+    x = keras.layers.Concatenate(axis=channel_axis, name=f"{name}_concat")(
         [shortcut, x]
     )
     return x

keras_hub/src/models/distil_bert/__init__.py

@@ -18,9 +18,6 @@ from keras_hub.src.models.distil_bert.distil_bert_backbone import (
 from keras_hub.src.models.distil_bert.distil_bert_presets import (
     backbone_presets,
 )
-from keras_hub.src.models.distil_bert.distil_bert_tokenizer import (
-    DistilBertTokenizer,
-)
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (DistilBertBackbone, DistilBertTokenizer))
+register_presets(backbone_presets, DistilBertBackbone)

keras_hub/src/models/distil_bert/distil_bert_masked_lm_preprocessor.py

@@ -13,19 +13,20 @@
 # 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.distil_bert.distil_bert_backbone import (
+    DistilBertBackbone,
 )
-from keras_hub.src.models.distil_bert.distil_bert_preprocessor import (
-    DistilBertPreprocessor,
+from keras_hub.src.models.distil_bert.distil_bert_tokenizer import (
+    DistilBertTokenizer,
 )
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.DistilBertMaskedLMPreprocessor")
-class DistilBertMaskedLMPreprocessor(DistilBertPreprocessor):
+class DistilBertMaskedLMPreprocessor(MaskedLMPreprocessor):
     """DistilBERT preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -119,76 +120,13 @@ class DistilBertMaskedLMPreprocessor(DistilBertPreprocessor):
     ```
     """
 
-    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,
-            ],
-        )
+    backbone_cls = DistilBertBackbone
+    tokenizer_cls = DistilBertTokenizer
 
+    @preprocessing_function
     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, padding_mask = x["token_ids"], x["padding_mask"]
-        masker_outputs = self.masker(token_ids)
-        x = {
-            "token_ids": masker_outputs["token_ids"],
-            "padding_mask": padding_mask,
-            "mask_positions": masker_outputs["mask_positions"],
-        }
-        y = masker_outputs["mask_ids"]
-        sample_weight = masker_outputs["mask_weights"]
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        # Backbone has no segment ID input.
+        del x["segment_ids"]
         return keras.utils.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_hub/src/models/distil_bert/{distil_bert_classifier.py → distil_bert_text_classifier.py}

@@ -16,20 +16,20 @@
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.classifier import Classifier
 from keras_hub.src.models.distil_bert.distil_bert_backbone import (
     DistilBertBackbone,
 )
 from keras_hub.src.models.distil_bert.distil_bert_backbone import (
     distilbert_kernel_initializer,
 )
-from keras_hub.src.models.distil_bert.distil_bert_preprocessor import (
-    DistilBertPreprocessor,
+from keras_hub.src.models.distil_bert.distil_bert_text_classifier_preprocessor import (
+    DistilBertTextClassifierPreprocessor,
 )
+from keras_hub.src.models.text_classifier import TextClassifier
 
 
-@keras_hub_export("keras_hub.models.DistilBertClassifier")
-class DistilBertClassifier(Classifier):
+@keras_hub_export("keras_hub.models.DistilBertTextClassifier")
+class DistilBertTextClassifier(TextClassifier):
     """An end-to-end DistilBERT model for classification tasks.
 
     This model attaches a classification head to a
@@ -50,7 +50,7 @@ class DistilBertClassifier(Classifier):
     Args:
         backbone: A `keras_hub.models.DistilBert` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.DistilBertPreprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.DistilBertTextClassifierPreprocessor` 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
@@ -69,12 +69,12 @@ class DistilBertClassifier(Classifier):
     labels = [0, 3]
 
     # Use a shorter sequence length.
-    preprocessor = keras_hub.models.DistilBertPreprocessor.from_preset(
+    preprocessor = keras_hub.models.DistilBertTextClassifierPreprocessor.from_preset(
         "distil_bert_base_en_uncased",
         sequence_length=128,
     )
     # Pretrained classifier.
-    classifier = keras_hub.models.DistilBertClassifier.from_preset(
+    classifier = keras_hub.models.DistilBertTextClassifier.from_preset(
         "distil_bert_base_en_uncased",
         num_classes=4,
         preprocessor=preprocessor,
@@ -102,7 +102,7 @@ class DistilBertClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.DistilBertClassifier.from_preset(
+    classifier = keras_hub.models.DistilBertTextClassifier.from_preset(
         "distil_bert_base_en_uncased",
         num_classes=4,
         preprocessor=None,
@@ -119,7 +119,7 @@ class DistilBertClassifier(Classifier):
     tokenizer = keras_hub.models.DistilBertTokenizer(
         vocabulary=vocab,
     )
-    preprocessor = keras_hub.models.DistilBertPreprocessor(
+    preprocessor = keras_hub.models.DistilBertTextClassifierPreprocessor(
         tokenizer=tokenizer,
         sequence_length=128,
     )
@@ -131,7 +131,7 @@ class DistilBertClassifier(Classifier):
         intermediate_dim=512,
         max_sequence_length=128,
     )
-    classifier = keras_hub.models.DistilBertClassifier(
+    classifier = keras_hub.models.DistilBertTextClassifier(
         backbone=backbone,
         preprocessor=preprocessor,
         num_classes=4,
@@ -141,7 +141,7 @@ class DistilBertClassifier(Classifier):
     """
 
     backbone_cls = DistilBertBackbone
-    preprocessor_cls = DistilBertPreprocessor
+    preprocessor_cls = DistilBertTextClassifierPreprocessor
 
     def __init__(
         self,