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

keras_hub/src/layers/modeling/transformer_encoder.py

@@ -0,0 +1,262 @@
+# 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.utils.keras_utils import clone_initializer
+
+from keras_hub.src.layers.modeling.transformer_layer_utils import (  # isort:skip
+    merge_padding_and_attention_mask,
+)
+
+
+@keras_hub_export("keras_hub.layers.TransformerEncoder")
+class TransformerEncoder(keras.layers.Layer):
+    """Transformer encoder.
+
+    This class follows the architecture of the transformer encoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up an encoder.
+
+    This layer will correctly compute an attention mask from an implicit
+    Keras padding mask (for example, by passing `mask_zero=True` to a
+    `keras.layers.Embedding` layer). See the Masking and Padding
+    [guide](https://keras.io/guides/understanding_masking_and_padding/)
+    for more details.
+
+    Args:
+        intermediate_dim: int, the hidden size of feedforward network.
+        num_heads: int, the number of heads in the
+            `keras.layers.MultiHeadAttention` layer.
+        dropout: float. the dropout value, shared by
+            `keras.layers.MultiHeadAttention` and feedforward network.
+            Defaults to `0.`.
+        activation: string or `keras.activations`. the
+            activation function of feedforward network.
+            Defaults to `"relu"`.
+        layer_norm_epsilon: float. The epsilon value in layer
+            normalization components. Defaults to `1e-5`.
+        kernel_initializer: string or `keras.initializers` initializer.
+            The kernel initializer for the dense and multiheaded
+            attention layers. Defaults to `"glorot_uniform"`.
+        bias_initializer: string or `keras.initializers` initializer.
+            The bias initializer for the dense and multiheaded
+            attention layers. Defaults to `"zeros"`.
+        normalize_first: bool. If True, the inputs to the
+            attention layer and the intermediate dense layer  are normalized
+            (similar to GPT-2). If set to False, outputs of attention layer and
+            intermediate dense layer are normalized (similar to BERT).
+            Defaults to `False`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Example:
+
+    ```python
+    # Create a single transformer encoder layer.
+    encoder = keras_hub.layers.TransformerEncoder(
+        intermediate_dim=64, num_heads=8)
+
+    # Create a simple model containing the encoder.
+    input = keras.Input(shape=(10, 64))
+    output = encoder(input)
+    model = keras.Model(inputs=input, outputs=output)
+
+    # Call encoder on the inputs.
+    input_data = np.random.uniform(size=(2, 10, 64))
+    output = model(input_data)
+    ```
+
+    References:
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
+    """
+
+    def __init__(
+        self,
+        intermediate_dim,
+        num_heads,
+        dropout=0,
+        activation="relu",
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        normalize_first=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.intermediate_dim = intermediate_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.activation = keras.activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
+        self.normalize_first = normalize_first
+        self.supports_masking = True
+
+    def build(self, inputs_shape):
+        # Infer the dimension of our hidden feature size from the build shape.
+        hidden_dim = inputs_shape[-1]
+        # Attention head size is `hidden_dim` over the number of heads.
+        key_dim = int(hidden_dim // self.num_heads)
+        if key_dim == 0:
+            raise ValueError(
+                "Attention `key_dim` computed cannot be zero. "
+                f"The `hidden_dim` value of {hidden_dim} has to be equal to "
+                f"or greater than `num_heads` value of {self.num_heads}."
+            )
+
+        # Self attention layers.
+        self._self_attention_layer = keras.layers.MultiHeadAttention(
+            num_heads=self.num_heads,
+            key_dim=key_dim,
+            dropout=self.dropout,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
+            name="self_attention_layer",
+        )
+        if hasattr(self._self_attention_layer, "_build_from_signature"):
+            self._self_attention_layer._build_from_signature(
+                query=inputs_shape,
+                value=inputs_shape,
+            )
+        else:
+            self._self_attention_layer.build(
+                query_shape=inputs_shape,
+                value_shape=inputs_shape,
+            )
+        self._self_attention_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="self_attention_layer_norm",
+        )
+        self._self_attention_layer_norm.build(inputs_shape)
+        self._self_attention_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
+        )
+
+        # Feedforward layers.
+        self._feedforward_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="feedforward_layer_norm",
+        )
+        self._feedforward_layer_norm.build(inputs_shape)
+        self._feedforward_intermediate_dense = keras.layers.Dense(
+            self.intermediate_dim,
+            activation=self.activation,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
+            name="feedforward_intermediate_dense",
+        )
+        self._feedforward_intermediate_dense.build(inputs_shape)
+        self._feedforward_output_dense = keras.layers.Dense(
+            hidden_dim,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
+        )
+        intermediate_shape = list(inputs_shape)
+        intermediate_shape[-1] = self.intermediate_dim
+        self._feedforward_output_dense.build(tuple(intermediate_shape))
+        self._feedforward_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="feedforward_dropout",
+        )
+        self.built = True
+
+    def call(
+        self, inputs, padding_mask=None, attention_mask=None, training=None
+    ):
+        """Forward pass of the TransformerEncoder.
+
+        Args:
+            inputs: a Tensor. The input data to TransformerEncoder, should be
+                of shape [batch_size, sequence_length, hidden_dim].
+            padding_mask: a boolean Tensor. It indicates if the token should be
+                masked because the token is introduced due to padding.
+                `padding_mask` should have shape [batch_size, sequence_length].
+            attention_mask: a boolean Tensor. Customized mask used to mask out
+                certain tokens. `attention_mask` should have shape
+                [batch_size, sequence_length, sequence_length].
+            training: a boolean indicating whether the layer should behave in
+                training mode or in inference mode.
+
+        Returns:
+            A Tensor of the same shape as the `inputs`.
+        """
+        x = inputs  # Intermediate result.
+
+        # Compute self attention mask.
+        self_attention_mask = merge_padding_and_attention_mask(
+            inputs, padding_mask, attention_mask
+        )
+
+        # Self attention block.
+        residual = x
+        if self.normalize_first:
+            x = self._self_attention_layer_norm(x)
+        x = self._self_attention_layer(
+            query=x,
+            value=x,
+            attention_mask=self_attention_mask,
+            training=training,
+        )
+        x = self._self_attention_dropout(x, training=training)
+        x = x + residual
+        if not self.normalize_first:
+            x = self._self_attention_layer_norm(x)
+
+        # Feedforward block.
+        residual = x
+        if self.normalize_first:
+            x = self._feedforward_layer_norm(x)
+        x = self._feedforward_intermediate_dense(x)
+        x = self._feedforward_output_dense(x)
+        x = self._feedforward_dropout(x, training=training)
+        x = x + residual
+        if not self.normalize_first:
+            x = self._feedforward_layer_norm(x)
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "dropout": self.dropout,
+                "activation": keras.activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
+                "normalize_first": self.normalize_first,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, inputs_shape):
+        return inputs_shape

keras_hub/src/layers/modeling/transformer_layer_utils.py

@@ -0,0 +1,106 @@
+# 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 absl import logging
+from keras import ops
+
+
+def _check_masks_shapes(inputs, padding_mask, attention_mask):
+    mask = padding_mask
+    if hasattr(inputs, "_keras_mask") and mask is None:
+        mask = inputs._keras_mask
+    if mask is not None:
+        if len(mask.shape) != 2:
+            raise ValueError(
+                "`padding_mask` should have shape "
+                "(batch_size, target_length). "
+                f"Received shape `{mask.shape}`."
+            )
+    if attention_mask is not None:
+        if len(attention_mask.shape) != 3:
+            raise ValueError(
+                "`attention_mask` should have shape "
+                "(batch_size, target_length, source_length). "
+                f"Received shape `{mask.shape}`."
+            )
+
+
+def compute_causal_mask(batch_size, input_length, output_length, cache_index=0):
+    """Compute a causal attention mask for a transformer decoder.
+
+    Args:
+        batch_size: batch size for the mask.
+        input_length: the length of key/value tensors in the attention layer.
+        output_length: the length of query tensors in the attention layer.
+        cache_index: the current index for cached generation. If passed, the
+            query sequence will be considered to start at `cache_index` rather
+            than zero. For example, a causal mask with `output_length=1` and
+            `cache_index=5` would allow the query tensor to attend to the first
+            five positions of the key/value tensors.
+
+    Return:
+        A causal attention mask with shape
+        `(batch_size, output_length, input_length)` that can be passed to a
+        attention layer.
+    """
+    i = ops.arange(output_length, dtype="float32")
+    i = i + ops.cast(cache_index, "float32")
+    i = ops.expand_dims(i, axis=1)
+    j = ops.arange(input_length, dtype="float32")
+    mask = ops.expand_dims(i >= j, axis=0)
+
+    return ops.broadcast_to(mask, (batch_size, output_length, input_length))
+
+
+def merge_padding_and_attention_mask(
+    inputs,
+    padding_mask,
+    attention_mask,
+):
+    """Merge the padding mask with a customized attention mask.
+
+    Args:
+        inputs: the input sequence.
+        padding_mask: the 1D padding mask, of shape
+            [batch_size, sequence_length].
+        attention_mask: the 2D customized mask, of shape
+            [batch_size, sequence1_length, sequence2_length].
+
+    Return:
+        A merged 2D mask or None. If only `padding_mask` is provided, the
+        returned mask is padding_mask with one additional axis.
+    """
+    _check_masks_shapes(inputs, padding_mask, attention_mask)
+    mask = padding_mask
+    if hasattr(inputs, "_keras_mask"):
+        if mask is None:
+            # If no padding mask is explicitly provided, we look for padding
+            # mask from the input data.
+            mask = inputs._keras_mask
+        else:
+            logging.warning(
+                "You are explicitly setting `padding_mask` while the `inputs` "
+                "have built-in mask, so the built-in mask is ignored."
+            )
+    if mask is not None:
+        # Add an axis for broadcasting, the attention mask should be 2D
+        # (not including the batch axis).
+        mask = ops.cast(ops.expand_dims(mask, axis=1), "int32")
+    if attention_mask is not None:
+        attention_mask = ops.cast(attention_mask, "int32")
+        if mask is None:
+            return attention_mask
+        else:
+            return ops.minimum(mask, attention_mask)
+    return mask

keras_hub/src/layers/preprocessing/__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/layers/preprocessing/masked_lm_mask_generator.py

@@ -0,0 +1,220 @@
+# 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.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
+from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
+
+try:
+    import tensorflow as tf
+    import tensorflow_text as tf_text
+except ImportError:
+    tf = None
+    tf_text = None
+
+
+@keras_hub_export("keras_hub.layers.MaskedLMMaskGenerator")
+class MaskedLMMaskGenerator(PreprocessingLayer):
+    """Layer that applies language model masking.
+
+    This layer is useful for preparing inputs for masked language modeling
+    (MaskedLM) tasks. It follows the masking strategy described in the
+    [original BERT paper](https://arxiv.org/abs/1810.04805). Given tokenized
+    text, it randomly selects certain number of tokens for masking. Then for
+    each selected token, it has a chance (configurable) to be replaced by
+    "mask token" or random token, or stay unchanged.
+
+    Input data should be passed as tensors, `tf.RaggedTensor`s, or lists. For
+    batched input, inputs should be a list of lists or a rank two tensor. For
+    unbatched inputs, each element should be a list or a rank one tensor.
+
+    This layer can be used with `tf.data` to generate dynamic masks on the fly
+    during training.
+
+    Args:
+        vocabulary_size: int, the size of the vocabulary.
+        mask_selection_rate: float, the probability of a token is selected for
+            masking.
+        mask_token_id: int. The id of mask token.
+        mask_selection_length: int. Maximum number of tokens
+            selected for  masking in each sequence. If set, the output
+            `mask_positions`, `mask_ids` and `mask_weights` will be padded
+            to dense tensors of length `mask_selection_length`, otherwise
+            the output will be a RaggedTensor. Defaults to `None`.
+        unselectable_token_ids: A list of tokens id that should not be
+            considered eligible for masking. By default, we assume `0`
+            corresponds to a padding token and ignore it. Defaults to `[0]`.
+        mask_token_rate: float. `mask_token_rate` must be
+            between 0 and 1 which indicates how often the mask_token is
+            substituted for tokens selected for masking. Defaults to `0.8`.
+        random_token_rate: float. `random_token_rate` must be
+            between 0 and 1 which indicates how often a random token is
+            substituted for tokens selected for masking.
+            Note: mask_token_rate + random_token_rate <= 1,  and for
+            (1 - mask_token_rate - random_token_rate), the token will not be
+            changed. Defaults to `0.1`.
+
+    Returns:
+        A Dict with 4 keys:
+            token_ids: Tensor or RaggedTensor, has the same type and shape of
+                input. Sequence after getting masked.
+            mask_positions: Tensor, or RaggedTensor if `mask_selection_length`
+                is None. The positions of token_ids getting masked.
+            mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is
+                None. The original token ids at masked positions.
+            mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is
+                None. `mask_weights` has the same shape as `mask_positions` and
+                `mask_ids`. Each element in `mask_weights` should be 0 or 1,
+                1 means the corresponding position in `mask_positions` is an
+                actual mask, 0 means it is a pad.
+
+    Examples:
+
+    Basic usage.
+    ```python
+    masker = keras_hub.layers.MaskedLMMaskGenerator(
+        vocabulary_size=10,
+        mask_selection_rate=0.2,
+        mask_token_id=0,
+        mask_selection_length=5
+    )
+    # Dense input.
+    masker([1, 2, 3, 4, 5])
+
+    # Ragged input.
+    masker([[1, 2], [1, 2, 3, 4]])
+    ```
+
+    Masking a batch that contains special tokens.
+    ```python
+    pad_id, cls_id, sep_id, mask_id = 0, 1, 2, 3
+    batch = [
+        [cls_id,   4,    5,      6, sep_id,    7,    8, sep_id, pad_id, pad_id],
+        [cls_id,   4,    5, sep_id,      6,    7,    8,      9, sep_id, pad_id],
+    ]
+
+    masker = keras_hub.layers.MaskedLMMaskGenerator(
+        vocabulary_size = 10,
+        mask_selection_rate = 0.2,
+        mask_selection_length = 5,
+        mask_token_id = mask_id,
+        unselectable_token_ids = [
+            cls_id,
+            sep_id,
+            pad_id,
+        ]
+    )
+    masker(batch)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        mask_selection_rate,
+        mask_token_id,
+        mask_selection_length=None,
+        unselectable_token_ids=[0],
+        mask_token_rate=0.8,
+        random_token_rate=0.1,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.vocabulary_size = vocabulary_size
+        self.unselectable_token_ids = unselectable_token_ids
+        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
+
+        if mask_token_id >= vocabulary_size:
+            raise ValueError(
+                f"Mask token id should be in range [0, vocabulary_size - 1], "
+                f"but received mask_token_id={mask_token_id}."
+            )
+        self.mask_token_id = mask_token_id
+
+        max_selections = self.mask_selection_length
+        if max_selections is None:
+            # Set a large number to remove the `max_selections_per_batch` cap.
+            max_selections = 2**31 - 1
+        self._random_selector = tf_text.RandomItemSelector(
+            max_selections_per_batch=max_selections,
+            selection_rate=self.mask_selection_rate,
+            unselectable_ids=self.unselectable_token_ids,
+        )
+        self._mask_values_chooser = tf_text.MaskValuesChooser(
+            self.vocabulary_size,
+            self.mask_token_id,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+        )
+
+    def call(self, inputs):
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
+
+        (
+            token_ids,
+            mask_positions,
+            mask_ids,
+        ) = tf_text.mask_language_model(
+            inputs,
+            item_selector=self._random_selector,
+            mask_values_chooser=self._mask_values_chooser,
+        )
+
+        if rectangular:
+            # If we converted the input from dense to ragged, convert back.
+            token_ids = token_ids.to_tensor()
+
+        mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
+        # If `mask_selection_length` is set, convert to dense.
+        if self.mask_selection_length:
+            target_shape = tf.cast([-1, self.mask_selection_length], "int64")
+            mask_positions = mask_positions.to_tensor(shape=target_shape)
+            mask_ids = mask_ids.to_tensor(shape=target_shape)
+            mask_weights = mask_weights.to_tensor(shape=target_shape)
+
+        if unbatched:
+            # If inputs is 1D, we format the output to be 1D as well.
+            token_ids = tf.squeeze(token_ids, axis=0)
+            mask_positions = tf.squeeze(mask_positions, axis=0)
+            mask_ids = tf.squeeze(mask_ids, axis=0)
+            mask_weights = tf.squeeze(mask_weights, axis=0)
+
+        return {
+            "token_ids": token_ids,
+            "mask_positions": mask_positions,
+            "mask_ids": mask_ids,
+            "mask_weights": mask_weights,
+        }
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "unselectable_token_ids": self.unselectable_token_ids,
+                "mask_token_id": self.mask_token_id,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config