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

keras_hub/src/models/text_classifier_preprocessor.py

@@ -0,0 +1,138 @@
+# 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.multi_segment_packer import (
+    MultiSegmentPacker,
+)
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.models.TextClassifierPreprocessor")
+class TextClassifierPreprocessor(Preprocessor):
+    """Base class for text classification preprocessing layers.
+
+    `TextClassifierPreprocessor` tasks wrap a `keras_hub.tokenizer.Tokenizer` to
+    create a preprocessing layer for text classification tasks. It is intended
+    to be paired with a `keras_hub.models.TextClassifier` task.
+
+    All `TextClassifierPreprocessor` take inputs three ordered inputs, `x`, `y`,
+    and `sample_weight`. `x`, the first input, should always be included. It can
+    be a single string, a batch of strings, or a tuple of batches of string
+    segments that should be combined into a single sequence. See examples below.
+    `y` and `sample_weight` are optional inputs that will be passed through
+    unaltered. Usually, `y` will be the classification label, and
+    `sample_weight` will not be provided.
+
+    The layer will output either `x`, an `(x, y)` tuple if labels were provided,
+    or an `(x, y, sample_weight)` tuple if labels and sample weight were
+    provided. `x` will be a dictionary with tokenized input, the exact contents
+    of the dictionary will depend on the model being used.
+
+    All `TextClassifierPreprocessor` tasks include a `from_preset()` constructor
+    which can be used to load a pre-trained config and vocabularies. You can
+    call the `from_preset()` constructor directly on this base class, in which
+    case the correct class for you model will be automatically instantiated.
+
+    Examples.
+    ```python
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
+        "bert_base_en_uncased",
+        sequence_length=256, # Optional.
+    )
+
+    # Tokenize and pad/truncate a single sentence.
+    x = "The quick brown fox jumped."
+    x = preprocessor(x)
+
+    # Tokenize and pad/truncate a labeled sentence.
+    x, y = "The quick brown fox jumped.", 1
+    x, y = preprocessor(x, y)
+
+    # Tokenize and pad/truncate a batch of labeled sentences.
+    x, y = ["The quick brown fox jumped.", "Call me Ishmael."], [1, 0]
+    x, y = preprocessor(x, y)
+
+    # Tokenize and combine a batch of labeled sentence pairs.
+    first = ["The quick brown fox jumped.", "Call me Ishmael."]
+    second = ["The fox tripped.", "Oh look, a whale."]
+    labels = [1, 0]
+    x, y = (first, second), labels
+    x, y = preprocessor(x, y)
+
+    # Use a `tf.data.Dataset`.
+    ds = tf.data.Dataset.from_tensor_slices(((first, second), labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tokenizer = tokenizer
+        self.packer = None
+        self.sequence_length = sequence_length
+        self.truncate = truncate
+
+    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.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            truncate=self.truncate,
+            sequence_length=self.sequence_length,
+        )
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        x = x if isinstance(x, tuple) else (x,)
+        x = tuple(self.tokenizer(segment) for segment in x)
+        token_ids, segment_ids = self.packer(x)
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": token_ids != self.tokenizer.pad_token_id,
+            "segment_ids": segment_ids,
+        }
+        return keras.utils.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
+
+    @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/whisper/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.whisper.whisper_backbone import WhisperBackbone
 from keras_hub.src.models.whisper.whisper_presets import backbone_presets
-from keras_hub.src.models.whisper.whisper_tokenizer import WhisperTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (WhisperBackbone, WhisperTokenizer))
+register_presets(backbone_presets, WhisperBackbone)

keras_hub/src/models/whisper/{whisper_audio_feature_extractor.py → whisper_audio_converter.py}

@@ -15,24 +15,19 @@
 
 import numpy as np
 
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.audio_converter import AudioConverter
+from keras_hub.src.models.whisper.whisper_backbone import WhisperBackbone
+
 try:
     import tensorflow as tf
 except ImportError:
-    raise ImportError(
-        "To use `keras_hub`, please install Tensorflow: `pip install tensorflow`. "
-        "The TensorFlow package is required for data preprocessing with any backend."
-    )
-
-from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.layers.preprocessing.preprocessing_layer import (
-    PreprocessingLayer,
-)
+    tf = None
 
 
-@keras_hub_export("keras_hub.models.WhisperAudioFeatureExtractor")
-class WhisperAudioFeatureExtractor(PreprocessingLayer):
-    """
-    Whisper audio feature extractor layer.
+@keras_hub_export("keras_hub.layers.WhisperAudioConverter")
+class WhisperAudioConverter(AudioConverter):
+    """Whisper audio converter layer.
 
     This layer takes in a batch of audio tensors, and computes the log-mel
     spectrogram features for each audio tensor.
@@ -55,22 +50,25 @@ class WhisperAudioFeatureExtractor(PreprocessingLayer):
             `max_audio_length * sampling_rate`. Defaults to `30`.
 
     Examples:
-
     ```python
     audio_tensor = tf.ones((8000,), dtype="float32")
 
     # Compute the log-mel spectrogram.
-    whisper_audio_feature_extractor = keras_hub.models.WhisperAudioFeatureExtractor()
-    whisper_audio_feature_extractor(audio_tensor)
+    audio_converter = keras_hub.models.WhisperAudioConverter.from_preset(
+        "whisper_base_en",
+    )
+    audio_converter(audio_tensor)
 
     # Compute the log-mel spectrogram for a batch of audio tensors.
     audio_tensor_1 = tf.ones((8000,), dtype="float32")
-    audio_tensor_2 = tf.ones((10000,), dtype="float32"
+    audio_tensor_2 = tf.ones((10000,), dtype="float32")
     audio_tensor = tf.ragged.stack([audio_tensor_1, audio_tensor_2], axis=0)
-    whisper_audio_feature_extractor(audio_tensor)
+    audio_converter(audio_tensor)
     ```
     """
 
+    backbone_cls = WhisperBackbone
+
     def __init__(
         self,
         num_mels=80,
@@ -97,6 +95,10 @@ class WhisperAudioFeatureExtractor(PreprocessingLayer):
         # `(num_fft_bins // 2 + 1, num_mels).`
         self.mel_filters = self._get_mel_filters()
 
+    def audio_shape(self):
+        """Returns the preprocessed size of a single audio sample."""
+        return (self.max_audio_length, self.num_mels)
+
     def _get_mel_filters(self):
         """
         Adapted from Hugging Face

keras_hub/src/models/whisper/whisper_backbone.py

@@ -24,7 +24,6 @@ from keras_hub.src.layers.modeling.token_and_position_embedding import (
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.models.whisper.whisper_decoder import WhisperDecoder
 from keras_hub.src.models.whisper.whisper_encoder import WhisperEncoder
-from keras_hub.src.utils.tensor_utils import assert_tf_backend
 
 
 def whisper_kernel_initializer(stddev=0.02):
@@ -117,8 +116,6 @@ class WhisperBackbone(Backbone):
         dtype=None,
         **kwargs,
     ):
-        assert_tf_backend(self.__class__.__name__)
-
         # === Layers ===
         self.encoder_conv_layer_1 = keras.layers.Conv1D(
             filters=hidden_dim,

keras_hub/src/models/whisper/whisper_presets.py

@@ -25,7 +25,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_en/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_en/3",
     },
     "whisper_base_en": {
         "metadata": {
@@ -38,7 +38,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_en/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_en/3",
     },
     "whisper_small_en": {
         "metadata": {
@@ -51,7 +51,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_en/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_en/3",
     },
     "whisper_medium_en": {
         "metadata": {
@@ -64,7 +64,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_en/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_en/3",
     },
     "whisper_tiny_multi": {
         "metadata": {
@@ -77,7 +77,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_multi/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_multi/3",
     },
     "whisper_base_multi": {
         "metadata": {
@@ -90,7 +90,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_multi/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_multi/3",
     },
     "whisper_small_multi": {
         "metadata": {
@@ -103,7 +103,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_multi/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_multi/3",
     },
     "whisper_medium_multi": {
         "metadata": {
@@ -116,7 +116,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_multi/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_multi/3",
     },
     "whisper_large_multi": {
         "metadata": {
@@ -129,7 +129,7 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi/3",
     },
     "whisper_large_multi_v2": {
         "metadata": {
@@ -143,6 +143,6 @@ backbone_presets = {
             "path": "whisper",
             "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi_v2/2",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi_v2/3",
     },
 }

keras_hub/src/models/whisper/whisper_tokenizer.py

@@ -15,6 +15,7 @@
 import json
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.whisper.whisper_backbone import WhisperBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
@@ -25,7 +26,12 @@ def _load_dict(dict_or_path):
     return dict_or_path
 
 
-@keras_hub_export("keras_hub.models.WhisperTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.WhisperTokenizer",
+        "keras_hub.models.WhisperTokenizer",
+    ]
+)
 class WhisperTokenizer(BytePairTokenizer):
     """Whisper text tokenizer using Byte-Pair Encoding subword segmentation.
 
@@ -47,6 +53,8 @@ class WhisperTokenizer(BytePairTokenizer):
             tokenizer.
     """
 
+    backbone_cls = WhisperBackbone
+
     def __init__(
         self,
         vocabulary=None,
@@ -94,20 +102,22 @@ class WhisperTokenizer(BytePairTokenizer):
         self.translate_token_id = special_tokens[self.translate_token]
         self.transcribe_token_id = special_tokens[self.transcribe_token]
 
-        self.special_tokens = special_tokens
+        self._special_token_dict = special_tokens
         self.language_tokens = language_tokens
-
-        # TODO: Add language tokens to `unsplittable_tokens` once we figure
-        # out the performance issue with a large list.
-        unsplittable_tokens = list(special_tokens.keys())
-
         super().__init__(
             vocabulary=vocabulary,
             merges=merges,
-            unsplittable_tokens=unsplittable_tokens,
             **kwargs,
         )
 
+    @property
+    def special_tokens(self):
+        return list(self._special_token_dict.keys())
+
+    @property
+    def special_token_ids(self):
+        return list(self._special_token_dict.values())
+
     def save_assets(self, dir_path):
         # TODO: whisper is currently mutating it's vocabulary before passing
         # it to the super class, so we need to restore the unmutated vocabulary
@@ -140,7 +150,7 @@ class WhisperTokenizer(BytePairTokenizer):
                 self.translate_token,
                 self.transcribe_token,
             ]:
-                vocabulary[token] = self.special_tokens[token]
+                vocabulary[token] = self._special_token_dict[token]
         else:
             self._initial_vocabulary = None
 
@@ -148,15 +158,9 @@ class WhisperTokenizer(BytePairTokenizer):
 
     def get_config(self):
         config = super().get_config()
-
-        # In the constructor, we pass the list of special tokens to the
-        # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
-        # delete it from the config here.
-        del config["unsplittable_tokens"]
-
         config.update(
             {
-                "special_tokens": self.special_tokens,
+                "special_tokens": self._special_token_dict,
                 "language_tokens": self.language_tokens,
             }
         )

keras_hub/src/models/xlm_roberta/__init__.py

@@ -18,9 +18,6 @@ from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
 from keras_hub.src.models.xlm_roberta.xlm_roberta_presets import (
     backbone_presets,
 )
-from keras_hub.src.models.xlm_roberta.xlm_roberta_tokenizer import (
-    XLMRobertaTokenizer,
-)
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (XLMRobertaBackbone, XLMRobertaTokenizer))
+register_presets(backbone_presets, XLMRobertaBackbone)

keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm_preprocessor.py

@@ -13,19 +13,23 @@
 # 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.layers.preprocessing.multi_segment_packer import (
+    MultiSegmentPacker,
 )
-from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
-    XLMRobertaPreprocessor,
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
 )
+from keras_hub.src.models.xlm_roberta.xlm_roberta_tokenizer import (
+    XLMRobertaTokenizer,
+)
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.XLMRobertaMaskedLMPreprocessor")
-class XLMRobertaMaskedLMPreprocessor(XLMRobertaPreprocessor):
+class XLMRobertaMaskedLMPreprocessor(MaskedLMPreprocessor):
     """XLM-RoBERTa preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -120,76 +124,26 @@ class XLMRobertaMaskedLMPreprocessor(XLMRobertaPreprocessor):
     ```
     """
 
-    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
+    backbone_cls = XLMRobertaBackbone
+    tokenizer_cls = XLMRobertaTokenizer
 
     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.start_token_id,
-                self.tokenizer.end_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,
-            }
+        # Roberta is doubles up the sep token, so we override build.
+        self.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            sep_value=[self.tokenizer.end_token_id] * 2,
+            pad_value=self.tokenizer.pad_token_id,
+            truncate=self.truncate,
+            sequence_length=self.sequence_length,
         )
-        return config
+        self.built = True
 
+    @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)

keras_hub/src/models/xlm_roberta/{xlm_roberta_classifier.py → xlm_roberta_text_classifier.py}

@@ -16,20 +16,25 @@
 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.roberta.roberta_backbone import (
     roberta_kernel_initializer,
 )
+from keras_hub.src.models.text_classifier import TextClassifier
 from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
     XLMRobertaBackbone,
 )
-from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
-    XLMRobertaPreprocessor,
+from keras_hub.src.models.xlm_roberta.xlm_roberta_text_classifier_preprocessor import (
+    XLMRobertaTextClassifierPreprocessor,
 )
 
 
-@keras_hub_export("keras_hub.models.XLMRobertaClassifier")
-class XLMRobertaClassifier(Classifier):
+@keras_hub_export(
+    [
+        "keras_hub.models.XLMRobertaTextClassifier",
+        "keras_hub.models.XLMRobertaClassifier",
+    ]
+)
+class XLMRobertaTextClassifier(TextClassifier):
     """An end-to-end XLM-RoBERTa model for classification tasks.
 
     This model attaches a classification head to a
@@ -50,7 +55,7 @@ class XLMRobertaClassifier(Classifier):
     Args:
         backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.XLMRobertaPreprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.XLMRobertaTextClassifierPreprocessor` 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 activation function to use
@@ -68,7 +73,7 @@ class XLMRobertaClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier.
-    classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
+    classifier = keras_hub.models.XLMRobertaTextClassifier.from_preset(
         "xlm_roberta_base_multi",
         num_classes=4,
     )
@@ -96,7 +101,7 @@ class XLMRobertaClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
+    classifier = keras_hub.models.XLMRobertaTextClassifier.from_preset(
         "xlm_roberta_base_multi",
         num_classes=4,
         preprocessor=None,
@@ -128,7 +133,7 @@ class XLMRobertaClassifier(Classifier):
     tokenizer = keras_hub.models.XLMRobertaTokenizer(
         proto=proto
     )
-    preprocessor = keras_hub.models.XLMRobertaPreprocessor(
+    preprocessor = keras_hub.models.XLMRobertaTextClassifierPreprocessor(
         tokenizer,
         sequence_length=128,
     )
@@ -140,7 +145,7 @@ class XLMRobertaClassifier(Classifier):
         intermediate_dim=512,
         max_sequence_length=128,
     )
-    classifier = keras_hub.models.XLMRobertaClassifier(
+    classifier = keras_hub.models.XLMRobertaTextClassifier(
         backbone=backbone,
         preprocessor=preprocessor,
         num_classes=4,
@@ -150,7 +155,7 @@ class XLMRobertaClassifier(Classifier):
     """
 
     backbone_cls = XLMRobertaBackbone
-    preprocessor_cls = XLMRobertaPreprocessor
+    preprocessor_cls = XLMRobertaTextClassifierPreprocessor
 
     def __init__(
         self,