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

keras_hub/src/layers/modeling/rotary_embedding.py

@@ -0,0 +1,169 @@
+# 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.RotaryEmbedding")
+class RotaryEmbedding(keras.layers.Layer):
+    """Rotary positional encoding layer.
+
+    This layer encodes absolute positional information with a rotation
+    matrix. It calculates the rotary encoding with a mix of sine and
+    cosine functions with geometrically increasing wavelengths.
+    Defined and formulated in [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4).
+    The input must be a tensor with shape a sequence dimension and a feature
+    dimension. Typically, this will either an input with shape
+    `(batch_size, sequence_length, feature_length)` or
+    `(batch_size, sequence_length, num_heads, feature_length)`.
+    This layer will return a new tensor with the rotary embedding applied to
+    the input tensor.
+
+    Args:
+        max_wavelength: int. The maximum angular wavelength of the sine/cosine
+            curves.
+        scaling_factor: float. The scaling factor used to scale positions of
+            the tokens.
+        sequence_axis: int. Sequence axis in the input tensor.
+        feature_axis: int. Feature axis in the input tensor.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Call arguments:
+        inputs: The tensor inputs to apply the embedding to. This can have
+            any shape, but must contain both a sequence and feature axis. The
+            rotary embedding will be applied to `inputs` and returned.
+        start_index: An integer or integer tensor. The starting position to
+            compute the rotary embedding from. This is useful during cached
+            decoding, where each position is predicted separately in a loop.
+
+    Examples:
+
+    ```python
+    batch_size = 16
+    feature_length = 18
+    sequence_length = 256
+    num_heads = 8
+
+    # No multi-head dimension.
+    tensor = np.ones((batch_size, sequence_length, feature_length))
+    rot_emb_layer = RotaryEmbedding()
+    tensor_rot = rot_emb_layer(tensor)
+
+    # With multi-head dimension.
+    tensor = np.ones((batch_size, sequence_length, num_heads, feature_length))
+    tensor_rot = rot_emb_layer(tensor)
+    ```
+
+    References:
+     - [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4)
+    """
+
+    def __init__(
+        self,
+        max_wavelength=10000,
+        scaling_factor=1.0,
+        sequence_axis=1,
+        feature_axis=-1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.max_wavelength = max_wavelength
+        self.sequence_axis = sequence_axis
+        self.feature_axis = feature_axis
+        self.scaling_factor = scaling_factor
+        self.built = True
+
+    def call(self, inputs, start_index=0, positions=None):
+        inputs = ops.moveaxis(
+            inputs, (self.feature_axis, self.sequence_axis), (-1, 1)
+        )
+        cos_emb, sin_emb = self._compute_cos_sin_embedding(
+            inputs, start_index, positions
+        )
+        output = self._apply_rotary_pos_emb(inputs, cos_emb, sin_emb)
+        return ops.moveaxis(
+            output, (-1, 1), (self.feature_axis, self.sequence_axis)
+        )
+
+    def _apply_rotary_pos_emb(self, tensor, cos_emb, sin_emb):
+        x1, x2 = ops.split(tensor, 2, axis=-1)
+        # Avoid `ops.concatenate` for now, to avoid a obscure bug with XLA
+        # compilation on jax. We should be able to remove this once the
+        # following PR is in all jax releases we care about:
+        # https://github.com/openxla/xla/pull/7875
+        half_rot_tensor = ops.stack((-x2, x1), axis=-2)
+        half_rot_tensor = ops.reshape(half_rot_tensor, ops.shape(tensor))
+        return (tensor * cos_emb) + (half_rot_tensor * sin_emb)
+
+    def _compute_positions(self, inputs, start_index=0):
+        seq_len = ops.shape(inputs)[1]
+        positions = ops.arange(seq_len, dtype="float32")
+        return positions + ops.cast(start_index, dtype="float32")
+
+    def _compute_cos_sin_embedding(self, inputs, start_index=0, positions=None):
+        feature_axis = len(inputs.shape) - 1
+        sequence_axis = 1
+
+        rotary_dim = ops.shape(inputs)[feature_axis]
+        inverse_freq = self._get_inverse_freq(rotary_dim)
+
+        if positions is None:
+            positions = self._compute_positions(inputs, start_index)
+        else:
+            positions = ops.cast(positions, "float32")
+
+        positions = positions / ops.cast(self.scaling_factor, "float32")
+        freq = ops.einsum("i,j->ij", positions, inverse_freq)
+        embedding = ops.stack((freq, freq), axis=-2)
+        embedding = ops.reshape(
+            embedding, (*ops.shape(freq)[:-1], ops.shape(freq)[-1] * 2)
+        )
+
+        # Reshape the embedding to be broadcastable with input shape.
+        if feature_axis < sequence_axis:
+            embedding = ops.transpose(embedding)
+        for axis in range(len(inputs.shape)):
+            if axis != sequence_axis and axis != feature_axis:
+                embedding = ops.expand_dims(embedding, axis)
+
+        cos_emb = ops.cast(ops.cos(embedding), self.compute_dtype)
+        sin_emb = ops.cast(ops.sin(embedding), self.compute_dtype)
+        return cos_emb, sin_emb
+
+    def _get_inverse_freq(self, rotary_dim):
+        freq_range = ops.divide(
+            ops.arange(0, rotary_dim, 2, dtype="float32"),
+            ops.cast(rotary_dim, "float32"),
+        )
+        inverse_freq = 1.0 / (self.max_wavelength**freq_range)
+        return inverse_freq
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_wavelength": self.max_wavelength,
+                "scaling_factor": self.scaling_factor,
+                "sequence_axis": self.sequence_axis,
+                "feature_axis": self.feature_axis,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

keras_hub/src/layers/modeling/sine_position_encoding.py

@@ -0,0 +1,108 @@
+# 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.SinePositionEncoding")
+class SinePositionEncoding(keras.layers.Layer):
+    """Sinusoidal positional encoding layer.
+
+    This layer calculates the position encoding as a mix of sine and cosine
+    functions with geometrically increasing wavelengths. Defined and formulized
+    in [Attention is All You Need](https://arxiv.org/abs/1706.03762).
+
+    Takes as input an embedded token tensor. The input must have shape
+    [batch_size, sequence_length, feature_size]. This layer will return a
+    positional encoding the same size as the embedded token tensor, which
+    can be added directly to the embedded token tensor.
+
+    Args:
+        max_wavelength: The maximum angular wavelength of the sine/cosine
+            curves, as described in Attention is All You Need. Defaults to
+            `10000`.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Call arguments:
+        inputs: The tensor inputs to compute an embedding for, with shape
+            `(batch_size, sequence_length, hidden_dim)`.
+        start_index: An integer or integer tensor. The starting position to
+            compute the encoding from. This is useful during cached decoding,
+            where each position is predicted separately in a loop.
+
+    Example:
+    ```python
+    # create a simple embedding layer with sinusoidal positional encoding
+    seq_len = 100
+    vocab_size = 1000
+    embedding_dim = 32
+    inputs = keras.Input((seq_len,), dtype="float32")
+    embedding = keras.layers.Embedding(
+        input_dim=vocab_size, output_dim=embedding_dim
+    )(inputs)
+    positional_encoding = keras_hub.layers.SinePositionEncoding()(embedding)
+    outputs = embedding + positional_encoding
+    ```
+
+    References:
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
+    """
+
+    def __init__(
+        self,
+        max_wavelength=10000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.max_wavelength = max_wavelength
+        self.built = True
+
+    def call(self, inputs, start_index=0):
+        shape = ops.shape(inputs)
+        seq_length = shape[-2]
+        hidden_size = shape[-1]
+        positions = ops.arange(seq_length)
+        positions = ops.cast(positions + start_index, self.compute_dtype)
+        min_freq = ops.cast(1 / self.max_wavelength, dtype=self.compute_dtype)
+        timescales = ops.power(
+            min_freq,
+            ops.cast(2 * (ops.arange(hidden_size) // 2), self.compute_dtype)
+            / ops.cast(hidden_size, self.compute_dtype),
+        )
+        angles = ops.expand_dims(positions, 1) * ops.expand_dims(timescales, 0)
+        # even indices are sine, odd are cosine
+        cos_mask = ops.cast(ops.arange(hidden_size) % 2, self.compute_dtype)
+        sin_mask = 1 - cos_mask
+        # embedding shape is [seq_length, hidden_size]
+        positional_encodings = (
+            ops.sin(angles) * sin_mask + ops.cos(angles) * cos_mask
+        )
+
+        return ops.broadcast_to(positional_encodings, shape)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_wavelength": self.max_wavelength,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

keras_hub/src/layers/modeling/token_and_position_embedding.py

@@ -0,0 +1,150 @@
+# 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.position_embedding import PositionEmbedding
+from keras_hub.src.layers.modeling.reversible_embedding import (
+    ReversibleEmbedding,
+)
+from keras_hub.src.utils.keras_utils import clone_initializer
+
+
+@keras_hub_export("keras_hub.layers.TokenAndPositionEmbedding")
+class TokenAndPositionEmbedding(keras.layers.Layer):
+    """A layer which sums a token and position embedding.
+
+    Token and position embeddings are ways of representing words and their order
+    in a sentence. This layer creates a `keras.layers.Embedding` token embedding
+    and a `keras_hub.layers.PositionEmbedding` position embedding and sums their
+    output when called. This layer assumes that the last dimension in the input
+    corresponds to the sequence dimension.
+
+    Args:
+        vocabulary_size: The size of the vocabulary.
+        sequence_length: The maximum length of input sequence
+        embedding_dim: The output dimension of the embedding layer
+        tie_weights: Boolean, whether or not the matrix for embedding and
+            the matrix for the `reverse` projection should share the same
+            weights.
+        embeddings_initializer: The initializer to use for the Embedding
+            Layers
+        mask_zero: Boolean, whether or not the input value 0 is a special
+            "padding" value that should be masked out.
+            This is useful when using recurrent layers which may take variable
+            length input. If this is True, then all subsequent layers in the
+            model need to support masking or an exception will be raised.
+            If mask_zero` is set to True, as a consequence, index 0 cannot be
+            used in the vocabulary
+            (input_dim should equal size of vocabulary + 1).
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `trainable`, `dtype` etc.
+
+    Example:
+    ```python
+    inputs = np.ones(shape=(1, 50), dtype="int32")
+    embedding_layer = keras_hub.layers.TokenAndPositionEmbedding(
+        vocabulary_size=10_000,
+        sequence_length=50,
+        embedding_dim=128,
+    )
+    outputs = embedding_layer(inputs)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        sequence_length,
+        embedding_dim,
+        tie_weights=True,
+        embeddings_initializer="uniform",
+        mask_zero=False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if vocabulary_size is None:
+            raise ValueError(
+                "`vocabulary_size` must be an Integer, received `None`."
+            )
+        if sequence_length is None:
+            raise ValueError(
+                "`sequence_length` must be an Integer, received `None`."
+            )
+        if embedding_dim is None:
+            raise ValueError(
+                "`embedding_dim` must be an Integer, received `None`."
+            )
+        self.vocabulary_size = int(vocabulary_size)
+        self.sequence_length = int(sequence_length)
+        self.embedding_dim = int(embedding_dim)
+        self.embeddings_initializer = keras.initializers.get(
+            embeddings_initializer
+        )
+        self.token_embedding = ReversibleEmbedding(
+            vocabulary_size,
+            embedding_dim,
+            tie_weights=tie_weights,
+            embeddings_initializer=clone_initializer(
+                self.embeddings_initializer
+            ),
+            mask_zero=mask_zero,
+            dtype=self.dtype_policy,
+            name="token_embedding",
+        )
+        self.position_embedding = PositionEmbedding(
+            sequence_length=sequence_length,
+            initializer=clone_initializer(self.embeddings_initializer),
+            dtype=self.dtype_policy,
+            name="position_embedding",
+        )
+        self.supports_masking = self.token_embedding.supports_masking
+
+    def build(self, input_shape):
+        input_shape = tuple(input_shape)
+        self.token_embedding.build(input_shape)
+        self.position_embedding.build(input_shape + (self.embedding_dim,))
+        self.built = True
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "sequence_length": self.sequence_length,
+                "embedding_dim": self.embedding_dim,
+                "embeddings_initializer": keras.initializers.serialize(
+                    self.embeddings_initializer
+                ),
+                "tie_weights": self.token_embedding.tie_weights,
+                "mask_zero": self.token_embedding.mask_zero,
+            }
+        )
+        return config
+
+    def call(self, inputs, start_index=0):
+        embedded_tokens = self.token_embedding(inputs)
+        embedded_positions = self.position_embedding(
+            embedded_tokens,
+            start_index=start_index,
+        )
+        outputs = embedded_tokens + embedded_positions
+        return outputs
+
+    def compute_mask(self, inputs, mask=None):
+        return self.token_embedding.compute_mask(inputs, mask=mask)
+
+    def compute_output_shape(self, input_shape):
+        return tuple(input_shape) + (self.embedding_dim,)