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

keras_hub/src/models/causal_lm_preprocessor.py

@@ -0,0 +1,195 @@
+# 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.start_end_packer import StartEndPacker
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+from keras_hub.src.utils.tensor_utils import strip_to_ragged
+
+
+@keras_hub_export("keras_hub.models.CausalLMPreprocessor")
+class CausalLMPreprocessor(Preprocessor):
+    """Base class for causal language modeling preprocessing layers.
+
+    `CausalLMPreprocessor` tasks wrap a `keras_hub.tokenizer.Tokenizer` to
+    create a preprocessing layer for causal language modeling tasks. It is
+    intended to be paired with a `keras.models.CausalLM` task.
+
+    All `CausalLMPreprocessor` take inputs a single input. This can be a single
+    string or a batch of strings. See examples below. These inputs
+    will be tokenized and padded/truncated to a fixed sequence length.
+
+    This layer will always output a `(x, y, sample_weight)` tuple, where `x`
+    is a dictionary with the tokenized inputs, `y` contains the tokens from `x`
+    offset by 1, and `sample_weight` marks where `y` contains padded
+    values. The exact contents of `x` will vary depending on the model being
+    used.
+
+    a `CausalLMPreprocessor` contains two extra methods, `generate_preprocess`
+    and `generate_postprocess` for use with generation. See examples below.
+
+    All `CausalLMPreprocessor` 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.CausalLMPreprocessor.from_preset(
+        "bert_base_en_uncased",
+        sequence_length=256, # Optional.
+    )
+
+    # Tokenize, mask and pack a single sentence.
+    x = "The quick brown fox jumped."
+    x, y, sample_weight = preprocessor(x)
+
+    # Tokenize and pad/truncate a batch of labeled sentences.
+    x = ["The quick brown fox jumped.", "Call me Ishmael."]
+    x, y, sample_weight = preprocessor(x)
+
+    # With a `tf.data.Dataset`.
+    ds = tf.data.Dataset.from_tensor_slices(x)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Generate preprocess and postprocess.
+    x = preprocessor.generate_preprocess(x)  # Tokenized numeric inputs.
+    x = preprocessor.generate_postprocess(x)  # Detokenized string outputs.
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=1024,
+        add_start_token=True,
+        add_end_token=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tokenizer = tokenizer
+        self.packer = None
+        self.sequence_length = sequence_length
+        self.add_start_token = add_start_token
+        self.add_end_token = add_end_token
+
+    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 = StartEndPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+        )
+        self.built = True
+
+    @preprocessing_function
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        sequence_length = sequence_length or self.sequence_length
+        x = self.tokenizer(x)
+        # Pad with one extra token to account for the truncation below.
+        token_ids, padding_mask = self.packer(
+            x,
+            sequence_length=sequence_length + 1,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        # The last token does not have a next token, so we truncate it out.
+        x = {
+            "token_ids": token_ids[..., :-1],
+            "padding_mask": padding_mask[..., :-1],
+        }
+        # Target `y` will be the next token.
+        y, sample_weight = token_ids[..., 1:], padding_mask[..., 1:]
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    @preprocessing_function
+    def generate_preprocess(
+        self,
+        x,
+        sequence_length=None,
+    ):
+        """Convert strings to integer token input for generation.
+
+        Similar to calling the layer for training, this method takes in strings
+        or tensor strings, 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 sequence (as generation is expected to continue at the end of the
+        inputted prompt).
+        """
+        if not self.built:
+            self.build(None)
+
+        x = self.tokenizer(x)
+        token_ids, padding_mask = self.packer(
+            x, sequence_length=sequence_length, add_end_value=False
+        )
+        return {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+
+    @preprocessing_function
+    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["token_ids"], x["padding_mask"]
+        ids_to_strip = self.tokenizer.special_token_ids
+        token_ids = strip_to_ragged(token_ids, padding_mask, ids_to_strip)
+        return self.tokenizer.detokenize(token_ids)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "add_start_token": self.add_start_token,
+                "add_end_token": self.add_end_token,
+            }
+        )
+        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/csp_darknet/csp_darknet_backbone.py

@@ -15,11 +15,11 @@ import keras
 from keras import layers
 
 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
 
 
 @keras_hub_export("keras_hub.models.CSPDarkNetBackbone")
-class CSPDarkNetBackbone(Backbone):
+class CSPDarkNetBackbone(FeaturePyramidBackbone):
     """This class represents Keras Backbone of CSPDarkNet model.
 
     This class implements a CSPDarkNet backbone as described in
@@ -65,12 +65,15 @@ class CSPDarkNetBackbone(Backbone):
         self,
         stackwise_num_filters,
         stackwise_depth,
-        include_rescaling,
+        include_rescaling=True,
         block_type="basic_block",
-        image_shape=(224, 224, 3),
+        image_shape=(None, None, 3),
         **kwargs,
     ):
         # === Functional Model ===
+        channel_axis = (
+            -1 if keras.config.image_data_format() == "channels_last" else 1
+        )
         apply_ConvBlock = (
             apply_darknet_conv_block_depthwise
             if block_type == "depthwise_block"
@@ -83,15 +86,22 @@ class CSPDarkNetBackbone(Backbone):
         if include_rescaling:
             x = layers.Rescaling(scale=1 / 255.0)(x)
 
-        x = apply_focus(name="stem_focus")(x)
+        x = apply_focus(channel_axis, name="stem_focus")(x)
         x = apply_darknet_conv_block(
-            base_channels, kernel_size=3, strides=1, name="stem_conv"
+            base_channels,
+            channel_axis,
+            kernel_size=3,
+            strides=1,
+            name="stem_conv",
         )(x)
+
+        pyramid_outputs = {}
         for index, (channels, depth) in enumerate(
             zip(stackwise_num_filters, stackwise_depth)
         ):
             x = apply_ConvBlock(
                 channels,
+                channel_axis,
                 kernel_size=3,
                 strides=2,
                 name=f"dark{index + 2}_conv",
@@ -100,17 +110,20 @@ class CSPDarkNetBackbone(Backbone):
             if index == len(stackwise_depth) - 1:
                 x = apply_spatial_pyramid_pooling_bottleneck(
                     channels,
+                    channel_axis,
                     hidden_filters=channels // 2,
                     name=f"dark{index + 2}_spp",
                 )(x)
 
             x = apply_cross_stage_partial(
                 channels,
+                channel_axis,
                 num_bottlenecks=depth,
                 block_type="basic_block",
                 residual=(index != len(stackwise_depth) - 1),
                 name=f"dark{index + 2}_csp",
             )(x)
+            pyramid_outputs[f"P{index + 2}"] = x
 
         super().__init__(inputs=image_input, outputs=x, **kwargs)
 
@@ -120,6 +133,7 @@ class CSPDarkNetBackbone(Backbone):
         self.include_rescaling = include_rescaling
         self.block_type = block_type
         self.image_shape = image_shape
+        self.pyramid_outputs = pyramid_outputs
 
     def get_config(self):
         config = super().get_config()
@@ -135,7 +149,7 @@ class CSPDarkNetBackbone(Backbone):
         return config
 
 
-def apply_focus(name=None):
+def apply_focus(channel_axis, name=None):
     """A block used in CSPDarknet to focus information into channels of the
     image.
 
@@ -151,7 +165,7 @@ def apply_focus(name=None):
     """
 
     def apply(x):
-        return layers.Concatenate(name=name)(
+        return layers.Concatenate(axis=channel_axis, name=name)(
             [
                 x[..., ::2, ::2, :],
                 x[..., 1::2, ::2, :],
@@ -164,7 +178,13 @@ def apply_focus(name=None):
 
 
 def apply_darknet_conv_block(
-    filters, kernel_size, strides, use_bias=False, activation="silu", name=None
+    filters,
+    channel_axis,
+    kernel_size,
+    strides,
+    use_bias=False,
+    activation="silu",
+    name=None,
 ):
     """
     The basic conv block used in Darknet. Applies Conv2D followed by a
@@ -193,11 +213,12 @@ def apply_darknet_conv_block(
             kernel_size,
             strides,
             padding="same",
+            data_format=keras.config.image_data_format(),
             use_bias=use_bias,
             name=name + "_conv",
         )(inputs)
 
-        x = layers.BatchNormalization(name=name + "_bn")(x)
+        x = layers.BatchNormalization(axis=channel_axis, name=name + "_bn")(x)
 
         if activation == "silu":
             x = layers.Lambda(lambda x: keras.activations.silu(x))(x)
@@ -212,7 +233,7 @@ def apply_darknet_conv_block(
 
 
 def apply_darknet_conv_block_depthwise(
-    filters, kernel_size, strides, activation="silu", name=None
+    filters, channel_axis, kernel_size, strides, activation="silu", name=None
 ):
     """
     The depthwise conv block used in CSPDarknet.
@@ -236,9 +257,13 @@ def apply_darknet_conv_block_depthwise(
 
     def apply(inputs):
         x = layers.DepthwiseConv2D(
-            kernel_size, strides, padding="same", use_bias=False
+            kernel_size,
+            strides,
+            padding="same",
+            data_format=keras.config.image_data_format(),
+            use_bias=False,
         )(inputs)
-        x = layers.BatchNormalization()(x)
+        x = layers.BatchNormalization(axis=channel_axis)(x)
 
         if activation == "silu":
             x = layers.Lambda(lambda x: keras.activations.swish(x))(x)
@@ -248,7 +273,11 @@ def apply_darknet_conv_block_depthwise(
             x = layers.LeakyReLU(0.1)(x)
 
         x = apply_darknet_conv_block(
-            filters, kernel_size=1, strides=1, activation=activation
+            filters,
+            channel_axis,
+            kernel_size=1,
+            strides=1,
+            activation=activation,
         )(x)
 
         return x
@@ -258,6 +287,7 @@ def apply_darknet_conv_block_depthwise(
 
 def apply_spatial_pyramid_pooling_bottleneck(
     filters,
+    channel_axis,
     hidden_filters=None,
     kernel_sizes=(5, 9, 13),
     activation="silu",
@@ -291,6 +321,7 @@ def apply_spatial_pyramid_pooling_bottleneck(
     def apply(x):
         x = apply_darknet_conv_block(
             hidden_filters,
+            channel_axis,
             kernel_size=1,
             strides=1,
             activation=activation,
@@ -304,13 +335,15 @@ def apply_spatial_pyramid_pooling_bottleneck(
                     kernel_size,
                     strides=1,
                     padding="same",
+                    data_format=keras.config.image_data_format(),
                     name=f"{name}_maxpool_{kernel_size}",
                 )(x[0])
             )
 
-        x = layers.Concatenate(name=f"{name}_concat")(x)
+        x = layers.Concatenate(axis=channel_axis, name=f"{name}_concat")(x)
         x = apply_darknet_conv_block(
             filters,
+            channel_axis,
             kernel_size=1,
             strides=1,
             activation=activation,
@@ -324,6 +357,7 @@ def apply_spatial_pyramid_pooling_bottleneck(
 
 def apply_cross_stage_partial(
     filters,
+    channel_axis,
     num_bottlenecks,
     residual=True,
     block_type="basic_block",
@@ -361,6 +395,7 @@ def apply_cross_stage_partial(
 
         x1 = apply_darknet_conv_block(
             hidden_channels,
+            channel_axis,
             kernel_size=1,
             strides=1,
             activation=activation,
@@ -369,6 +404,7 @@ def apply_cross_stage_partial(
 
         x2 = apply_darknet_conv_block(
             hidden_channels,
+            channel_axis,
             kernel_size=1,
             strides=1,
             activation=activation,
@@ -379,6 +415,7 @@ def apply_cross_stage_partial(
             residual_x = x1
             x1 = apply_darknet_conv_block(
                 hidden_channels,
+                channel_axis,
                 kernel_size=1,
                 strides=1,
                 activation=activation,
@@ -386,6 +423,7 @@ def apply_cross_stage_partial(
             )(x1)
             x1 = ConvBlock(
                 hidden_channels,
+                channel_axis,
                 kernel_size=3,
                 strides=1,
                 activation=activation,
@@ -399,6 +437,7 @@ def apply_cross_stage_partial(
         x = layers.Concatenate(name=f"{name}_concat")([x1, x2])
         x = apply_darknet_conv_block(
             filters,
+            channel_axis,
             kernel_size=1,
             strides=1,
             activation=activation,

keras_hub/src/models/deberta_v3/__init__.py

@@ -16,9 +16,6 @@ from keras_hub.src.models.deberta_v3.deberta_v3_backbone import (
     DebertaV3Backbone,
 )
 from keras_hub.src.models.deberta_v3.deberta_v3_presets import backbone_presets
-from keras_hub.src.models.deberta_v3.deberta_v3_tokenizer import (
-    DebertaV3Tokenizer,
-)
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (DebertaV3Backbone, DebertaV3Tokenizer))
+register_presets(backbone_presets, DebertaV3Backbone)

keras_hub/src/models/deberta_v3/deberta_v3_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.deberta_v3.deberta_v3_backbone import (
+    DebertaV3Backbone,
 )
-from keras_hub.src.models.deberta_v3.deberta_v3_preprocessor import (
-    DebertaV3Preprocessor,
+from keras_hub.src.models.deberta_v3.deberta_v3_tokenizer import (
+    DebertaV3Tokenizer,
 )
+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.DebertaV3MaskedLMPreprocessor")
-class DebertaV3MaskedLMPreprocessor(DebertaV3Preprocessor):
+class DebertaV3MaskedLMPreprocessor(MaskedLMPreprocessor):
     """DeBERTa preprocessing for the masked language modeling task.
 
     This preprocessing layer will prepare inputs for a masked language modeling
@@ -115,77 +116,13 @@ class DebertaV3MaskedLMPreprocessor(DebertaV3Preprocessor):
     ```
     """
 
-    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
+    backbone_cls = DebertaV3Backbone
+    tokenizer_cls = DebertaV3Tokenizer
 
+    @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/deberta_v3/{deberta_v3_classifier.py → deberta_v3_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.deberta_v3.deberta_v3_backbone import (
     DebertaV3Backbone,
 )
 from keras_hub.src.models.deberta_v3.deberta_v3_backbone import (
     deberta_kernel_initializer,
 )
-from keras_hub.src.models.deberta_v3.deberta_v3_preprocessor import (
-    DebertaV3Preprocessor,
+from keras_hub.src.models.deberta_v3.deberta_v3_text_classifier_preprocessor import (
+    DebertaV3TextClassifierPreprocessor,
 )
+from keras_hub.src.models.text_classifier import TextClassifier
 
 
-@keras_hub_export("keras_hub.models.DebertaV3Classifier")
-class DebertaV3Classifier(Classifier):
+@keras_hub_export("keras_hub.models.DebertaV3TextClassifier")
+class DebertaV3TextClassifier(TextClassifier):
     """An end-to-end DeBERTa model for classification tasks.
 
     This model attaches a classification head to a
@@ -53,7 +53,7 @@ class DebertaV3Classifier(Classifier):
     Args:
         backbone: A `keras_hub.models.DebertaV3` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.DebertaV3Preprocessor` or `None`. If
+        preprocessor: A `keras_hub.models.DebertaV3TextClassifierPreprocessor` 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
@@ -72,7 +72,7 @@ class DebertaV3Classifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier.
-    classifier = keras_hub.models.DebertaV3Classifier.from_preset(
+    classifier = keras_hub.models.DebertaV3TextClassifier.from_preset(
         "deberta_v3_base_en",
         num_classes=4,
     )
@@ -100,7 +100,7 @@ class DebertaV3Classifier(Classifier):
     labels = [0, 3]
 
     # Pretrained classifier without preprocessing.
-    classifier = keras_hub.models.DebertaV3Classifier.from_preset(
+    classifier = keras_hub.models.DebertaV3TextClassifier.from_preset(
         "deberta_v3_base_en",
         num_classes=4,
         preprocessor=None,
@@ -132,7 +132,7 @@ class DebertaV3Classifier(Classifier):
     tokenizer = keras_hub.models.DebertaV3Tokenizer(
         proto=bytes_io.getvalue(),
     )
-    preprocessor = keras_hub.models.DebertaV3Preprocessor(
+    preprocessor = keras_hub.models.DebertaV3TextClassifierPreprocessor(
         tokenizer=tokenizer,
         sequence_length=128,
     )
@@ -144,7 +144,7 @@ class DebertaV3Classifier(Classifier):
         intermediate_dim=512,
         max_sequence_length=128,
     )
-    classifier = keras_hub.models.DebertaV3Classifier(
+    classifier = keras_hub.models.DebertaV3TextClassifier(
         backbone=backbone,
         preprocessor=preprocessor,
         num_classes=4,
@@ -154,7 +154,7 @@ class DebertaV3Classifier(Classifier):
     """
 
     backbone_cls = DebertaV3Backbone
-    preprocessor_cls = DebertaV3Preprocessor
+    preprocessor_cls = DebertaV3TextClassifierPreprocessor
 
     def __init__(
         self,