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

keras_hub/src/layers/modeling/alibi_bias.py

@@ -0,0 +1,143 @@
+# 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 math
+
+import keras
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.layers.AlibiBias")
+class AlibiBias(keras.layers.Layer):
+    """A layer that adds the alibi bias to attention scores.
+
+    This layer adds the alibi bias to the attention scores. Alibi bias is a
+    linear, non-learned bias. Defined and formalized in
+    [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://arxiv.org/abs/2108.12409).
+
+    This layer takes as input the attention scores. and returns the attention
+    scores after adding the alibi bias to it. The output will have the same
+    shape as the input.
+
+    Args:
+        alibi_bias_max: int. This value will be used to compute the slope of
+            each head. The heads' slopes are a geometric sequence that starts at
+            `2**(-alibi_bias_max/num_heads)` and uses that same value as its
+            ratio. Defaults to 8.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Call arguments:
+        attention_scores: The result of multipying the query and the key of the
+            multi-head attention layer of the transformer to add alibi bias to
+            it. With shape `(batch_size, num_heads, query_length, key_length)`.
+
+    Example:
+    ```python
+    query_length = 10
+    key_length = 10
+    num_heads = 4
+    batch_size = 2
+    hidden_dim = 8
+
+    # Create new alibi layer.
+    alibi_layer = keras_hub.layers.AlibiBias()
+
+    query = np.zeros((batch_size, num_heads, query_length, hidden_dim))
+    key = np.zeros((batch_size, num_heads, hidden_dim, key_length))
+
+    attention_scores = keras.ops.matmul(query, key)
+
+    # Add alibi bias to attention scores.
+    attention_scores = alibi_layer(attention_scores)
+    ```
+
+    References:
+     - [Press et al., 2021](https://arxiv.org/abs/2108.12409)
+    """
+
+    def __init__(
+        self,
+        alibi_bias_max=8,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.alibi_bias_max = alibi_bias_max
+
+    def call(self, attention_scores):
+        shape = ops.shape(attention_scores)
+        if len(shape) != 4:
+            raise ValueError(
+                "Expected `attention_scores` shape to be "
+                "`(batch_size, num_heads, query_length, key_Length)`."
+                f" Recived shape={shape}"
+            )
+
+        key_length = shape[-1]
+        num_heads = shape[-3]
+
+        alibi_bias = self._get_alibi_bias(num_heads, key_length)
+
+        return ops.add(attention_scores, alibi_bias)
+
+    def _get_alibi_bias(self, num_heads, key_length):
+        slopes = ops.convert_to_tensor(
+            self._get_slopes(num_heads), dtype=self.compute_dtype
+        )
+        slopes = ops.expand_dims(slopes, 1)
+
+        seq_range = ops.expand_dims(
+            ops.arange(1 - key_length, 1, dtype="int32"), 0
+        )
+        seq_range = ops.cast(seq_range, dtype=self.compute_dtype)
+
+        alibi_bias = ops.multiply(slopes, seq_range)
+        alibi_bias = ops.expand_dims(alibi_bias, 1)
+
+        # return shape is `(1, num_heads, 1, key_length)`
+        return ops.expand_dims(alibi_bias, 0)
+
+    def _get_slopes(self, num_heads):
+        # this function is adopted from Alibi original implementation.
+        # https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742
+        def get_slopes_power_of_2(n):
+            start = 2 ** (
+                -(2 ** -(math.log2(n) - math.log2(self.alibi_bias_max)))
+            )
+            ratio = start
+            return [start * ratio**i for i in range(n)]
+
+        if math.log2(num_heads).is_integer():
+            return get_slopes_power_of_2(num_heads)
+        else:
+            closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+            return (
+                get_slopes_power_of_2(closest_power_of_2)
+                + self._get_slopes(2 * closest_power_of_2)[0::2][
+                    : num_heads - closest_power_of_2
+                ]
+            )
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "alibi_bias_max": self.alibi_bias_max,
+            }
+        )
+        return config

keras_hub/src/layers/modeling/cached_multi_head_attention.py

@@ -0,0 +1,137 @@
+# 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
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.layers.CachedMultiHeadAttention")
+class CachedMultiHeadAttention(keras.layers.MultiHeadAttention):
+    """MultiHeadAttention layer with cache support.
+
+    This layer is suitable for use in autoregressive decoding. It can be used
+    to cache decoder self-attention and cross-attention. The forward pass
+    can happen in one of three modes:
+
+    - No cache, same as regular multi-head attention.
+    - Static cache (`cache_update_index` is None). In this case, the
+        cached key/value projections will be used and the input values will
+        be ignored.
+    - Updated cache (`cache_update_index` is not None). In this case, new
+        key/value projections are computed using the input, and spliced into
+        the cache at the specified index.
+
+    Note that caching is useful only during inference and should not be used
+    during training.
+
+    We use the notation `B`, `T`, `S` below, where `B` is the batch dimension,
+    `T` is the target sequence length, and `S` in the source sequence length.
+    Note that during generative decoding, `T` is usually 1 (you are
+    generating a target sequence of length one to predict the next token).
+
+    Call arguments:
+        query: Query `Tensor` of shape `(B, T, dim)`.
+        value: Value `Tensor` of shape `(B, S*, dim)`. if `cache` is None`, `S*`
+            must equal `S` and match the shape of `attention_mask`. If cache` is
+            not `None`, `S*` can be any length less than `S`, and the computed
+            value will be spliced into `cache` at `cache_update_index`.
+        key: Optional key `Tensor` of shape `(B, S*, dim)`. If `cache` is
+            `None`, `S*` must equal `S` and match the shape of
+            `attention_mask`. If `cache` is not `None`, `S*` can be any length
+            less than `S`, and the computed value will be spliced into `cache`
+            at `cache_update_index`.
+        attention_mask: a boolean mask of shape `(B, T, S)`. `attention_mask`
+            prevents attention to certain positions. The boolean mask specifies
+            which query elements can attend to which key elements, 1 indicates
+            attention and 0 indicates no attention. Broadcasting can happen for
+            the missing batch dimensions and the head dimension.
+        cache: a dense float Tensor. The key/value cache, of shape
+            `[B, 2, S, num_heads, key_dims]`, where `S` must agree with the
+            `attention_mask` shape. This argument is intended for use during
+            generation to avoid recomputing intermediate state.
+        cache_update_index: a int or int Tensor, the index at which to update
+            `cache` (usually the index of the current token being processed
+            when running generation). If `cache_update_index=None` while `cache`
+            is set, the cache will not be updated.
+        training: a boolean indicating whether the layer should behave in
+            training mode or in inference mode.
+
+    Returns:
+        An `(attention_output, cache)` tuple. `attention_output` is the result
+        of the computation, of shape `(B, T, dim)`, where `T` is for target
+        sequence shapes and `dim` is the query input last dimension if
+        `output_shape` is `None`. Otherwise, the multi-head outputs are
+        projected to the shape specified by `output_shape`. `cache` is the
+        updated cache.
+    """
+
+    def call(
+        self,
+        query,
+        value,
+        key=None,
+        attention_mask=None,
+        cache=None,
+        cache_update_index=None,
+        training=None,
+    ):
+        if key is None:
+            key = value
+
+        query = self._query_dense(query)
+
+        # If cache is not `None`, we will use the cache to compute the final key
+        # and value tensors. If `cache_update_index` is not None, we will first
+        # update the cache before use. To do this, we first call the
+        # `_key_dense` and `_value_dense` layers, and copy the outputs into the
+        # cache at the specified index. `cache = None` handles the training
+        # case, where we don't use the cache at all.
+        if cache is not None:
+            key_cache = cache[:, 0, ...]
+            value_cache = cache[:, 1, ...]
+            if cache_update_index is None:
+                key = key_cache
+                value = value_cache
+            else:
+                key_update = self._key_dense(key)
+                value_update = self._value_dense(value)
+                start = [0, cache_update_index, 0, 0]
+                key = ops.slice_update(key_cache, start, key_update)
+                value = ops.slice_update(value_cache, start, value_update)
+                cache = ops.stack((key, value), axis=1)
+        else:
+            if cache_update_index is not None:
+                raise ValueError(
+                    "`cache_update_index` should not be set if `cache` is "
+                    f"`None`. Received: cache={cache}, "
+                    f"cache_update_index={cache_update_index}"
+                )
+            key = self._key_dense(key)
+            value = self._value_dense(value)
+
+        attention_output, attention_scores = self._compute_attention(
+            query=query,
+            key=key,
+            value=value,
+            attention_mask=attention_mask,
+            training=training,
+        )
+
+        attention_output = self._output_dense(attention_output)
+
+        if cache is not None:
+            return attention_output, cache
+        return attention_output

keras_hub/src/layers/modeling/f_net_encoder.py

@@ -0,0 +1,200 @@
+# 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
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.utils.keras_utils import clone_initializer
+
+
+@keras_hub_export("keras_hub.layers.FNetEncoder")
+class FNetEncoder(keras.layers.Layer):
+    """FNet encoder.
+
+    This class follows the architecture of FNet encoder layer in the
+    [FNet paper](https://arxiv.org/abs/2105.03824). Users can instantiate
+    multiple instances of this class to stack up the encoder.
+
+    Note on masking: In the official FNet code, padding tokens are added to the
+    the input. However, the padding masks are deleted, i.e., mixing of
+    all tokens is done. This is because certain frequencies will be zeroed
+    out if we apply padding masks in every encoder layer. Hence, we don't
+    take padding mask as input in the call() function.
+
+    Args:
+        intermediate_dim: int. The hidden size of feedforward network.
+        dropout: float. The dropout value, applied in the
+            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: `str` or `keras.initializers` initializer.
+            The kernel initializer for the dense layers.
+            Defaults to `"glorot_uniform"`.
+        bias_initializer: "string" or `keras.initializers` initializer.
+            The bias initializer for the dense layers.
+            Defaults to `"zeros"`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Example:
+
+    ```python
+    # Create a single FNet encoder layer.
+    encoder = keras_hub.layers.FNetEncoder(
+        intermediate_dim=64)
+
+    # 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=(1, 10, 64))
+    output = model(input_data)
+    ```
+
+    References:
+     - [Lee-Thorp et al., 2021](https://arxiv.org/abs/2105.03824)
+    """
+
+    def __init__(
+        self,
+        intermediate_dim,
+        dropout=0,
+        activation="relu",
+        layer_norm_epsilon=1e-5,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.intermediate_dim = intermediate_dim
+        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)
+
+    def build(self, inputs_shape):
+        # Create layers based on input shape.
+        feature_size = inputs_shape[-1]
+
+        # Layer Norm layers.
+        self._mixing_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="mixing_layer_norm",
+        )
+        self._mixing_layer_norm.build(inputs_shape)
+        self._output_layer_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="output_layer_norm",
+        )
+        self._output_layer_norm.build(inputs_shape)
+
+        # Feedforward layers.
+        self._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="intermediate_dense",
+        )
+        self._intermediate_dense.build(inputs_shape)
+        self._output_dense = keras.layers.Dense(
+            feature_size,
+            kernel_initializer=clone_initializer(self.kernel_initializer),
+            bias_initializer=clone_initializer(self.bias_initializer),
+            dtype=self.dtype_policy,
+            name="output_dense",
+        )
+        self._output_dense.build(
+            self._intermediate_dense.compute_output_shape(inputs_shape)
+        )
+        self._output_dropout = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="output_dropout",
+        )
+        self.built = True
+
+    def call(self, inputs, training=None):
+        """Forward pass of the FNetEncoder.
+
+        Args:
+            inputs: a Tensor. The input data to TransformerEncoder, should be
+                of shape [batch_size, sequence_length, feature_dim].
+            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`.
+        """
+
+        def fourier_transform(input):
+            # Apply FFT on the input and take the real part.
+            input_dtype = input.dtype
+            # FFT transforms do not support float16.
+            input = ops.cast(input, "float32")
+            real_in, imaginary_in = (input, ops.zeros_like(input))
+            real_out, _ = ops.fft2((real_in, imaginary_in))
+            return ops.cast(real_out, input_dtype)
+
+        def add_and_norm(input1, input2, norm_layer):
+            return norm_layer(input1 + input2)
+
+        def feed_forward(input):
+            x = self._intermediate_dense(input)
+            x = self._output_dense(x)
+            return self._output_dropout(x, training=training)
+
+        mixing_output = fourier_transform(inputs)
+
+        mixing_output = add_and_norm(
+            inputs, mixing_output, self._mixing_layer_norm
+        )
+
+        feed_forward_output = feed_forward(mixing_output)
+
+        x = add_and_norm(
+            mixing_output, feed_forward_output, self._output_layer_norm
+        )
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "intermediate_dim": self.intermediate_dim,
+                "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
+                ),
+            }
+        )
+        return config
+
+    def compute_output_shape(self, inputs_shape):
+        return inputs_shape