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

keras_hub/src/models/deberta_v3/relative_embedding.py

@@ -0,0 +1,94 @@
+# 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 import ops
+
+
+class RelativeEmbedding(keras.layers.Layer):
+    """Relative embedding layer.
+
+    This is an implementation of relative embedding as described in the
+    paper ["DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing"](https://arxiv.org/abs/2111.09543).
+    This layer initializes an embedding matrix (of shape
+    `(2 * batch_size, hidden_dim)`) for relative position encoding. It then
+    applies layer normalization on the embedding matrix and returns the relative
+    embedding matrix.
+
+    Args:
+        hidden_dim: int. The size of the dense embedding.
+        bucket_size: int. The size of the relative position buckets.
+        layer_norm_epsilon: float. Epsilon value to initialize the layer
+            normalization layer.
+        kernel_initializer: string or `keras.initializers` initializer.
+            The kernel initializer for the dense embedding.
+            Defaults to `"glorot_uniform"`.
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        bucket_size,
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_dim = hidden_dim
+        self.bucket_size = bucket_size
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+
+        self.rel_embeddings = self.add_weight(
+            shape=(self.bucket_size * 2, self.hidden_dim),
+            initializer=self.kernel_initializer,
+            name="rel_embedding",
+        )
+        self.layer_norm = keras.layers.LayerNormalization(
+            epsilon=layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="rel_embeddings_layer_norm",
+        )
+
+    def call(self, inputs):
+        batch_size = ops.shape(inputs)[0]
+
+        rel_embeddings = ops.expand_dims(
+            ops.convert_to_tensor(self.rel_embeddings), axis=0
+        )
+        rel_embeddings = self.layer_norm(rel_embeddings)
+
+        # Repeat `rel_embeddings` along axis = 0 `batch_size` times. The
+        # resultant shape is `(batch_size, bucket_size * 2, hidden_dim)`.
+        rel_embeddings = ops.repeat(rel_embeddings, repeats=batch_size, axis=0)
+
+        return rel_embeddings
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "bucket_size": self.bucket_size,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return (input_shape[0],) + (self.bucket_size * 2, self.hidden_dim)

keras_hub/src/models/densenet/__init__.py

@@ -0,0 +1,13 @@
+# 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.

keras_hub/src/models/densenet/densenet_backbone.py

@@ -0,0 +1,210 @@
+# 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.models.backbone import Backbone
+
+BN_AXIS = 3
+BN_EPSILON = 1.001e-5
+
+
+@keras_hub_export("keras_hub.models.DenseNetBackbone")
+class DenseNetBackbone(Backbone):
+    """Instantiates the DenseNet architecture.
+
+    This class implements a DenseNet backbone as described in
+    [Densely Connected Convolutional Networks (CVPR 2017)](
+       https://arxiv.org/abs/1608.06993
+    ).
+
+    Args:
+        stackwise_num_repeats: list of ints, number of repeated convolutional
+            blocks per dense block.
+        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).
+        compression_ratio: float, compression rate at transition layers,
+            defaults to 0.5.
+        growth_rate: int, number of filters added by each dense block,
+            defaults to 32
+
+    Examples:
+    ```python
+    input_data = np.ones(shape=(8, 224, 224, 3))
+
+    # Pretrained backbone
+    model = keras_hub.models.DenseNetBackbone.from_preset("densenet121_imagenet")
+    model(input_data)
+
+    # Randomly initialized backbone with a custom config
+    model = keras_hub.models.DenseNetBackbone(
+        stackwise_num_repeats=[6, 12, 24, 16],
+        include_rescaling=False,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        stackwise_num_repeats,
+        include_rescaling=True,
+        image_shape=(224, 224, 3),
+        compression_ratio=0.5,
+        growth_rate=32,
+        **kwargs,
+    ):
+        # === Functional Model ===
+        image_input = keras.layers.Input(shape=image_shape)
+
+        x = image_input
+        if include_rescaling:
+            x = keras.layers.Rescaling(1 / 255.0)(x)
+
+        x = keras.layers.Conv2D(
+            64, 7, strides=2, use_bias=False, padding="same", name="conv1_conv"
+        )(x)
+        x = keras.layers.BatchNormalization(
+            axis=BN_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"
+        )(x)
+
+        for stack_index in range(len(stackwise_num_repeats) - 1):
+            index = stack_index + 2
+            x = apply_dense_block(
+                x,
+                stackwise_num_repeats[stack_index],
+                growth_rate,
+                name=f"conv{index}",
+            )
+            x = apply_transition_block(
+                x, compression_ratio, name=f"pool{index}"
+            )
+
+        x = apply_dense_block(
+            x,
+            stackwise_num_repeats[-1],
+            growth_rate,
+            name=f"conv{len(stackwise_num_repeats) + 1}",
+        )
+
+        x = keras.layers.BatchNormalization(
+            axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
+        )(x)
+        x = keras.layers.Activation("relu", name="relu")(x)
+
+        super().__init__(inputs=image_input, outputs=x, **kwargs)
+
+        # === Config ===
+        self.stackwise_num_repeats = stackwise_num_repeats
+        self.include_rescaling = include_rescaling
+        self.compression_ratio = compression_ratio
+        self.growth_rate = growth_rate
+        self.image_shape = image_shape
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "stackwise_num_repeats": self.stackwise_num_repeats,
+                "include_rescaling": self.include_rescaling,
+                "compression_ratio": self.compression_ratio,
+                "growth_rate": self.growth_rate,
+                "image_shape": self.image_shape,
+            }
+        )
+        return config
+
+
+def apply_dense_block(x, num_repeats, growth_rate, name=None):
+    """A dense block.
+
+    Args:
+      x: input tensor.
+      num_repeats: int, number of repeated convolutional blocks.
+      growth_rate: int, number of filters added by each dense block.
+      name: string, block label.
+    """
+    if name is 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}")
+    return x
+
+
+def apply_transition_block(x, compression_ratio, name=None):
+    """A transition block.
+
+    Args:
+      x: input tensor.
+      compression_ratio: float, compression rate at transition layers.
+      name: string, block label.
+    """
+    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"
+    )(x)
+    x = keras.layers.Activation("relu", name=f"{name}_relu")(x)
+    x = keras.layers.Conv2D(
+        int(x.shape[BN_AXIS] * compression_ratio),
+        1,
+        use_bias=False,
+        name=f"{name}_conv",
+    )(x)
+    x = keras.layers.AveragePooling2D(2, strides=2, name=f"{name}_pool")(x)
+    return x
+
+
+def apply_conv_block(x, growth_rate, name=None):
+    """A building block for a dense block.
+
+    Args:
+      x: input tensor.
+      growth_rate: int, number of filters added by each dense block.
+      name: string, block label.
+    """
+    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"
+    )(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"
+    )(x)
+    x = keras.layers.BatchNormalization(
+        axis=BN_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(
+        growth_rate,
+        3,
+        padding="same",
+        use_bias=False,
+        name=f"{name}_2_conv",
+    )(x)
+    x = keras.layers.Concatenate(axis=BN_AXIS, name=f"{name}_concat")(
+        [shortcut, x]
+    )
+    return x

keras_hub/src/models/densenet/densenet_image_classifier.py

@@ -0,0 +1,131 @@
+# 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.models.densenet.densenet_backbone import DenseNetBackbone
+from keras_hub.src.models.image_classifier import ImageClassifier
+
+
+@keras_hub_export("keras_hub.models.DenseNetImageClassifier")
+class DenseNetImageClassifier(ImageClassifier):
+    """DenseNet image classifier task model.
+
+    To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
+    where `x` is a tensor and `y` is a integer from `[0, num_classes)`.
+    All `ImageClassifier` tasks include a `from_preset()` constructor which can
+    be used to load a pre-trained config and weights.
+
+    Args:
+        backbone: A `keras_hub.models.DenseNetBackbone` instance.
+        num_classes: int. The number of classes to predict.
+        activation: `None`, str or callable. The activation function to use on
+            the `Dense` layer. Set `activation=None` to return the output
+            logits. Defaults to `"softmax"`.
+
+    Examples:
+
+    Call `predict()` to run inference.
+    ```python
+    # Load preset and train
+    images = np.ones((2, 224, 224, 3), dtype="float32")
+    classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
+        "densenet121_imagenet")
+    classifier.predict(images)
+    ```
+
+    Call `fit()` on a single batch.
+    ```python
+    # Load preset and train
+    images = np.ones((2, 224, 224, 3), dtype="float32")
+    labels = [0, 3]
+    classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
+        "densenet121_imagenet")
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+
+    Call `fit()` with custom loss, optimizer and backbone.
+    ```python
+    classifier = keras_hub.models.DenseNetImageClassifier.from_preset(
+        "densenet121_imagenet")
+    classifier.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+    )
+    classifier.backbone.trainable = False
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+
+    Custom backbone.
+    ```python
+    images = np.ones((2, 224, 224, 3), dtype="float32")
+    labels = [0, 3]
+    backbone = keras_hub.models.DenseNetBackbone(
+        stackwise_num_filters=[128, 256, 512, 1024],
+        stackwise_depth=[3, 9, 9, 3],
+        include_rescaling=False,
+        block_type="basic_block",
+        image_shape = (224, 224, 3),
+    )
+    classifier = keras_hub.models.DenseNetImageClassifier(
+        backbone=backbone,
+        num_classes=4,
+    )
+    classifier.fit(x=images, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = DenseNetBackbone
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        activation="softmax",
+        preprocessor=None,  # adding this dummy arg for saved model test
+        # TODO: once preprocessor flow is figured out, this needs to be updated
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.output_dense = keras.layers.Dense(
+            num_classes,
+            activation=activation,
+            name="predictions",
+        )
+
+        # === Functional Model ===
+        inputs = self.backbone.input
+        x = self.backbone(inputs)
+        outputs = self.output_dense(x)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.activation = activation
+
+    def get_config(self):
+        # Backbone serialized in `super`
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "activation": self.activation,
+            }
+        )
+        return config

keras_hub/src/models/distil_bert/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+from keras_hub.src.models.distil_bert.distil_bert_backbone import (
+    DistilBertBackbone,
+)
+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))

keras_hub/src/models/distil_bert/distil_bert_backbone.py

@@ -0,0 +1,187 @@
+# 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.modeling.token_and_position_embedding import (
+    TokenAndPositionEmbedding,
+)
+from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
+from keras_hub.src.models.backbone import Backbone
+
+
+def distilbert_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras_hub_export("keras_hub.models.DistilBertBackbone")
+class DistilBertBackbone(Backbone):
+    """A DistilBERT encoder network.
+
+    This network implements a bi-directional Transformer-based encoder as
+    described in ["DistilBERT, a distilled version of BERT: smaller, faster,
+    cheaper and lighter"](https://arxiv.org/abs/1910.01108). It includes the
+    embedding lookups and transformer layers, but not the masked language model
+    or classification task networks.
+
+    The default constructor gives a fully customizable, randomly initialized
+    DistilBERT encoder with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the
+    `from_preset()` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available
+    [here](https://github.com/huggingface/transformers).
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
+
+    Examples:
+    ```python
+    input_data = {
+        "token_ids": np.ones(shape=(1, 12), dtype="int32"),
+        "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
+    }
+
+    # Pretrained DistilBERT encoder.
+    model = keras_hub.models.DistilBertBackbone.from_preset(
+        "distil_bert_base_en_uncased"
+    )
+    model(input_data)
+
+    # Randomly initialized DistilBERT encoder with custom config.
+    model = keras_hub.models.DistilBertBackbone(
+        vocabulary_size=30552,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_sequence_length=128,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=512,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.embeddings = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            sequence_length=max_sequence_length,
+            embedding_dim=hidden_dim,
+            embeddings_initializer=distilbert_kernel_initializer(),
+            dtype=dtype,
+            name="token_and_position_embedding",
+        )
+        # Keep the token_embedding property for consistency across models.
+        self.token_embedding = self.embeddings.token_embedding
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
+            axis=-1,
+            epsilon=1e-12,
+            dtype=dtype,
+            name="embeddings_layer_norm",
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="embeddings_dropout",
+        )
+        self.transformer_layers = []
+        for i in range(num_layers):
+            layer = TransformerEncoder(
+                num_heads=num_heads,
+                intermediate_dim=intermediate_dim,
+                activation="gelu",
+                dropout=dropout,
+                layer_norm_epsilon=1e-12,
+                kernel_initializer=distilbert_kernel_initializer(),
+                dtype=dtype,
+                name=f"transformer_layer_{i}",
+            )
+            self.transformer_layers.append(layer)
+
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.embeddings(token_id_input)
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask_input)
+        super().__init__(
+            inputs={
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
+            },
+            outputs=x,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_sequence_length = max_sequence_length
+        self.cls_token_index = 0
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "dropout": self.dropout,
+                "max_sequence_length": self.max_sequence_length,
+            }
+        )
+        return config