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

keras_hub/src/models/backbone.py

@@ -20,18 +20,11 @@ from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import assert_quantization_support
 from keras_hub.src.utils.preset_utils import CONFIG_FILE
 from keras_hub.src.utils.preset_utils import MODEL_WEIGHTS_FILE
-from keras_hub.src.utils.preset_utils import check_config_class
-from keras_hub.src.utils.preset_utils import check_format
-from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import list_presets
-from keras_hub.src.utils.preset_utils import list_subclasses
-from keras_hub.src.utils.preset_utils import load_serialized_object
+from keras_hub.src.utils.preset_utils import builtin_presets
+from keras_hub.src.utils.preset_utils import get_preset_loader
 from keras_hub.src.utils.preset_utils import save_metadata
 from keras_hub.src.utils.preset_utils import save_serialized_object
 from keras_hub.src.utils.python_utils import classproperty
-from keras_hub.src.utils.timm.convert import load_timm_backbone
-from keras_hub.src.utils.transformers.convert import load_transformers_backbone
 
 
 @keras_hub_export("keras_hub.models.Backbone")
@@ -147,11 +140,8 @@ class Backbone(keras.Model):
 
     @classproperty
     def presets(cls):
-        """List built-in presets for a `Task` subclass."""
-        presets = list_presets(cls)
-        for subclass in list_subclasses(cls):
-            presets.update(subclass.presets)
-        return presets
+        """List built-in presets for a `Backbone` subclass."""
+        return builtin_presets(cls)
 
     @classmethod
     def from_preset(
@@ -166,7 +156,7 @@ class Backbone(keras.Model):
         to save and load a pre-trained model. The `preset` can be passed as a
         one of:
 
-        1. a built in preset identifier like `'bert_base_en'`
+        1. a built-in preset identifier like `'bert_base_en'`
         2. a Kaggle Models handle like `'kaggle://user/bert/keras/bert_base_en'`
         3. a Hugging Face handle like `'hf://user/bert_base_en'`
         4. a path to a local preset directory like `'./bert_base_en'`
@@ -181,7 +171,7 @@ class Backbone(keras.Model):
         all built-in presets available on the class.
 
         Args:
-            preset: string. A built in preset identifier, a Kaggle Models
+            preset: string. A built-in preset identifier, a Kaggle Models
                 handle, a Hugging Face handle, or a path to a local directory.
             load_weights: bool. If `True`, the weights will be loaded into the
                 model architecture. If `False`, the weights will be randomly
@@ -201,27 +191,15 @@ class Backbone(keras.Model):
         )
         ```
         """
-        format = check_format(preset)
-
-        if format == "transformers":
-            return load_transformers_backbone(cls, preset, load_weights)
-        elif format == "timm":
-            return load_timm_backbone(cls, preset, load_weights, **kwargs)
-
-        preset_cls = check_config_class(preset)
-        if not issubclass(preset_cls, cls):
+        loader = get_preset_loader(preset)
+        backbone_cls = loader.check_backbone_class()
+        if not issubclass(backbone_cls, cls):
             raise ValueError(
-                f"Preset has type `{preset_cls.__name__}` which is not a "
+                f"Saved preset has type `{backbone_cls.__name__}` which is not "
                 f"a subclass of calling class `{cls.__name__}`. Call "
-                f"`from_preset` directly on `{preset_cls.__name__}` instead."
+                f"`from_preset` directly on `{backbone_cls.__name__}` instead."
             )
-
-        backbone = load_serialized_object(preset, CONFIG_FILE, **kwargs)
-        if load_weights:
-            jax_memory_cleanup(backbone)
-            backbone.load_weights(get_file(preset, MODEL_WEIGHTS_FILE))
-
-        return backbone
+        return loader.load_backbone(backbone_cls, load_weights, **kwargs)
 
     def save_to_preset(self, preset_dir):
         """Save backbone to a preset directory.

keras_hub/src/models/bart/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.bart.bart_backbone import BartBackbone
 from keras_hub.src.models.bart.bart_presets import backbone_presets
-from keras_hub.src.models.bart.bart_tokenizer import BartTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (BartBackbone, BartTokenizer))
+register_presets(backbone_presets, BartBackbone)

keras_hub/src/models/bart/bart_preprocessor.py

@@ -17,11 +17,10 @@ import keras
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.bart.bart_backbone import BartBackbone
 from keras_hub.src.models.bart.bart_tokenizer import BartTokenizer
 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.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.models.BartPreprocessor")
@@ -129,6 +128,7 @@ class BartPreprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = BartBackbone
     tokenizer_cls = BartTokenizer
 
     def __init__(
@@ -174,6 +174,7 @@ class BartPreprocessor(Preprocessor):
         )
         self.built = True
 
+    @preprocessing_function
     def call(
         self,
         x,
@@ -200,26 +201,13 @@ class BartPreprocessor(Preprocessor):
         if decoder_sequence_length is None:
             decoder_sequence_length = self.decoder_sequence_length
 
-        encoder_text = x["encoder_text"]
-        decoder_text = x["decoder_text"]
-
-        encoder_text = convert_inputs_to_list_of_tensor_segments(encoder_text)
-        decoder_text = convert_inputs_to_list_of_tensor_segments(decoder_text)
-
-        if len(encoder_text) > 1 or len(decoder_text) > 1:
-            raise ValueError(
-                '`BARTPreprocessor` requires both `"encoder_text"` and '
-                f'`"decoder_text"` to contain only one segment, but received '
-                f"{len(encoder_text)} and {len(decoder_text)}, respectively."
-            )
-
-        encoder_inputs = self.tokenizer(encoder_text[0])
+        encoder_inputs = self.tokenizer(x["encoder_text"])
         encoder_token_ids, encoder_padding_mask = self.encoder_packer(
             encoder_inputs,
             sequence_length=encoder_sequence_length,
         )
 
-        decoder_inputs = self.tokenizer(decoder_text[0])
+        decoder_inputs = self.tokenizer(x["decoder_text"])
         decoder_token_ids, decoder_padding_mask = self.decoder_packer(
             decoder_inputs,
             sequence_length=decoder_sequence_length,

keras_hub/src/models/bart/bart_seq_2_seq_lm_preprocessor.py

@@ -13,24 +13,15 @@
 # limitations under the License.
 
 
-import keras
-from absl import logging
-
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.bart.bart_preprocessor import BartPreprocessor
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
-from keras_hub.src.utils.tensor_utils import strip_to_ragged
-
-try:
-    import tensorflow as tf
-except ImportError:
-    tf = None
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.bart.bart_backbone import BartBackbone
+from keras_hub.src.models.bart.bart_tokenizer import BartTokenizer
+from keras_hub.src.models.seq_2_seq_lm_preprocessor import Seq2SeqLMPreprocessor
 
 
 @keras_hub_export("keras_hub.models.BartSeq2SeqLMPreprocessor")
-class BartSeq2SeqLMPreprocessor(BartPreprocessor):
+class BartSeq2SeqLMPreprocessor(Seq2SeqLMPreprocessor):
     """BART Seq2Seq LM preprocessor.
 
     This layer is used as preprocessor for seq2seq tasks using the BART model.
@@ -125,138 +116,20 @@ class BartSeq2SeqLMPreprocessor(BartPreprocessor):
     ```
     """
 
-    def call(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        *,
-        encoder_sequence_length=None,
-        decoder_sequence_length=None,
-        # `sequence_length` is an alias for `decoder_sequence_length`
-        sequence_length=None,
-    ):
-        if y is not None or sample_weight is not None:
-            logging.warning(
-                "`BartSeq2SeqLMPreprocessor` infers `y` and `sample_weight` "
-                "from the provided input data, i.e., `x`. However, non-`None`"
-                "values have been passed for `y` or `sample_weight` or both. "
-                "These values will be ignored."
-            )
-
-        if encoder_sequence_length is None:
-            encoder_sequence_length = self.encoder_sequence_length
-        decoder_sequence_length = decoder_sequence_length or sequence_length
-        if decoder_sequence_length is None:
-            decoder_sequence_length = self.decoder_sequence_length
-
-        x = super().call(
-            x,
-            encoder_sequence_length=encoder_sequence_length,
-            decoder_sequence_length=decoder_sequence_length + 1,
-        )
-        decoder_token_ids = x.pop("decoder_token_ids")
-        decoder_padding_mask = x.pop("decoder_padding_mask")
-
-        # The last token does not have a next token. Hence, we truncate it.
-        x = {
-            **x,
-            "decoder_token_ids": decoder_token_ids[..., :-1],
-            "decoder_padding_mask": decoder_padding_mask[..., :-1],
-        }
-        # Target `y` will be the decoder input sequence shifted one step to the
-        # left (i.e., the next token).
-        y = decoder_token_ids[..., 1:]
-        sample_weight = decoder_padding_mask[..., 1:]
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    def generate_preprocess(
-        self,
-        x,
-        *,
-        encoder_sequence_length=None,
-        # `sequence_length` is an alias for `decoder_sequence_length`
-        decoder_sequence_length=None,
-        sequence_length=None,
-    ):
-        """Convert encoder and decoder input strings to integer token inputs for generation.
-
-        Similar to calling the layer for training, this method takes in a dict
-        containing `"encoder_text"` and `"decoder_text"`, with strings or tensor
-        strings for values, tokenizes and packs the input, and computes a
-        padding mask masking all inputs not filled in with a padded value.
-
-        Unlike calling the layer for training, this method does not compute
-        labels and will never append a tokenizer.end_token_id to the end of
-        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 isinstance(x, dict):
-            encoder_text = x["encoder_text"]
-            decoder_text = x["decoder_text"]
-        else:
-            encoder_text = x
-            # Initialize empty prompt for the decoder.
-            decoder_text = tf.fill((tf.shape(encoder_text)[0],), "")
-
-        if encoder_sequence_length is None:
-            encoder_sequence_length = self.encoder_sequence_length
-        decoder_sequence_length = decoder_sequence_length or sequence_length
-        if decoder_sequence_length is None:
-            decoder_sequence_length = self.decoder_sequence_length
-
-        # Tokenize and pack the encoder inputs.
-        # TODO: Remove `[0]` once we have shifted to `MultiSegmentPacker`.
-        encoder_text = convert_inputs_to_list_of_tensor_segments(encoder_text)[
-            0
-        ]
-        encoder_token_ids = self.tokenizer(encoder_text)
-        encoder_token_ids, encoder_padding_mask = self.encoder_packer(
-            encoder_token_ids,
-            sequence_length=encoder_sequence_length,
-        )
-
-        # Tokenize and pack the decoder inputs.
-        decoder_text = convert_inputs_to_list_of_tensor_segments(decoder_text)[
-            0
-        ]
-        decoder_token_ids = self.tokenizer(decoder_text)
-        decoder_token_ids, decoder_padding_mask = self.decoder_packer(
-            decoder_token_ids,
-            sequence_length=decoder_sequence_length,
-            add_end_value=False,
-        )
-
-        return {
-            "encoder_token_ids": encoder_token_ids,
-            "encoder_padding_mask": encoder_padding_mask,
-            "decoder_token_ids": decoder_token_ids,
-            "decoder_padding_mask": decoder_padding_mask,
-        }
-
-    def generate_postprocess(
-        self,
-        x,
-    ):
-        """Convert integer token output to strings for generation.
-
-        This method reverses `generate_preprocess()`, by first removing all
-        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["decoder_token_ids"],
-            x["decoder_padding_mask"],
-        )
-        ids_to_strip = (
-            self.tokenizer.start_token_id,
-            self.tokenizer.end_token_id,
+    backbone_cls = BartBackbone
+    tokenizer_cls = BartTokenizer
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # The decoder is packed a bit differently; the format is as follows:
+        # `[end_token_id, start_token_id, tokens..., end_token_id, padding...]`.
+        self.decoder_packer = StartEndPacker(
+            start_value=[
+                self.tokenizer.end_token_id,
+                self.tokenizer.start_token_id,
+            ],
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.decoder_sequence_length,
+            return_padding_mask=True,
         )
-        token_ids = strip_to_ragged(token_ids, padding_mask, ids_to_strip)
-        return self.tokenizer.detokenize(token_ids)

keras_hub/src/models/bart/bart_tokenizer.py

@@ -14,10 +14,16 @@
 
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.bart.bart_backbone import BartBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
-@keras_hub_export("keras_hub.models.BartTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.BartTokenizer",
+        "keras_hub.models.BartTokenizer",
+    ]
+)
 class BartTokenizer(BytePairTokenizer):
     """A BART tokenizer using Byte-Pair Encoding subword segmentation.
 
@@ -73,52 +79,19 @@ class BartTokenizer(BytePairTokenizer):
     ```
     """
 
+    backbone_cls = BartBackbone
+
     def __init__(
         self,
         vocabulary=None,
         merges=None,
         **kwargs,
     ):
-        self.start_token = "<s>"
-        self.pad_token = "<pad>"
-        self.end_token = "</s>"
-
+        self._add_special_token("<s>", "start_token")
+        self._add_special_token("</s>", "end_token")
+        self._add_special_token("<pad>", "pad_token")
         super().__init__(
             vocabulary=vocabulary,
             merges=merges,
-            unsplittable_tokens=[
-                self.start_token,
-                self.pad_token,
-                self.end_token,
-            ],
             **kwargs,
         )
-
-    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
-
-    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"]
-        return config

keras_hub/src/models/bert/__init__.py

@@ -13,11 +13,7 @@
 # limitations under the License.
 
 from keras_hub.src.models.bert.bert_backbone import BertBackbone
-from keras_hub.src.models.bert.bert_classifier import BertClassifier
 from keras_hub.src.models.bert.bert_presets import backbone_presets
-from keras_hub.src.models.bert.bert_presets import classifier_presets
-from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (BertBackbone, BertTokenizer))
-register_presets(classifier_presets, (BertClassifier, BertTokenizer))
+register_presets(backbone_presets, BertBackbone)

keras_hub/src/models/bert/bert_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.bert.bert_preprocessor import BertPreprocessor
+from keras_hub.src.models.bert.bert_backbone import BertBackbone
+from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
 
 
 @keras_hub_export("keras_hub.models.BertMaskedLMPreprocessor")
-class BertMaskedLMPreprocessor(BertPreprocessor):
+class BertMaskedLMPreprocessor(MaskedLMPreprocessor):
     """BERT preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -117,82 +113,5 @@ class BertMaskedLMPreprocessor(BertPreprocessor):
     ```
     """
 
-    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 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, segment_ids = (
-            x["token_ids"],
-            x["padding_mask"],
-            x["segment_ids"],
-        )
-        masker_outputs = self.masker(token_ids)
-        x = {
-            "token_ids": masker_outputs["token_ids"],
-            "padding_mask": padding_mask,
-            "segment_ids": segment_ids,
-            "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)
-
-    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
+    backbone_cls = BertBackbone
+    tokenizer_cls = BertTokenizer

keras_hub/src/models/bert/bert_presets.py

@@ -129,9 +129,6 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/bert/keras/bert_large_en/2",
     },
-}
-
-classifier_presets = {
     "bert_tiny_en_uncased_sst2": {
         "metadata": {
             "description": (
@@ -143,5 +140,5 @@ classifier_presets = {
             "model_card": "https://github.com/google-research/bert/blob/master/README.md",
         },
         "kaggle_handle": "kaggle://keras/bert/keras/bert_tiny_en_uncased_sst2/4",
-    }
+    },
 }

keras_hub/src/models/bert/{bert_classifier.py → bert_text_classifier.py}

@@ -17,12 +17,14 @@ import keras
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.bert.bert_backbone import BertBackbone
 from keras_hub.src.models.bert.bert_backbone import bert_kernel_initializer
-from keras_hub.src.models.bert.bert_preprocessor import BertPreprocessor
-from keras_hub.src.models.classifier import Classifier
+from keras_hub.src.models.bert.bert_text_classifier_preprocessor import (
+    BertTextClassifierPreprocessor,
+)
+from keras_hub.src.models.text_classifier import TextClassifier
 
 
-@keras_hub_export("keras_hub.models.BertClassifier")
-class BertClassifier(Classifier):
+@keras_hub_export("keras_hub.models.BertTextClassifier")
+class BertTextClassifier(TextClassifier):
     """An end-to-end BERT model for classification tasks.
 
     This model attaches a classification head to a
@@ -41,7 +43,7 @@ class BertClassifier(Classifier):
     Args:
         backbone: A `keras_hub.models.BertBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.BertPreprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.BertTextClassifierPreprocessor` 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
@@ -59,7 +61,7 @@ class BertClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier.
-    classifier = keras_hub.models.BertClassifier.from_preset(
+    classifier = keras_hub.models.BertTextClassifier.from_preset(
         "bert_base_en_uncased",
         num_classes=4,
     )
@@ -88,7 +90,7 @@ class BertClassifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.BertClassifier.from_preset(
+    classifier = keras_hub.models.BertTextClassifier.from_preset(
         "bert_base_en_uncased",
         num_classes=4,
         preprocessor=None,
@@ -106,7 +108,7 @@ class BertClassifier(Classifier):
     tokenizer = keras_hub.models.BertTokenizer(
         vocabulary=vocab,
     )
-    preprocessor = keras_hub.models.BertPreprocessor(
+    preprocessor = keras_hub.models.BertTextClassifierPreprocessor(
         tokenizer=tokenizer,
         sequence_length=128,
     )
@@ -118,7 +120,7 @@ class BertClassifier(Classifier):
         intermediate_dim=512,
         max_sequence_length=128,
     )
-    classifier = keras_hub.models.BertClassifier(
+    classifier = keras_hub.models.BertTextClassifier(
         backbone=backbone,
         preprocessor=preprocessor,
         num_classes=4,
@@ -128,7 +130,7 @@ class BertClassifier(Classifier):
     """
 
     backbone_cls = BertBackbone
-    preprocessor_cls = BertPreprocessor
+    preprocessor_cls = BertTextClassifierPreprocessor
 
     def __init__(
         self,