keras-hub-nightly 0.21.0.dev202504160404__tar.gz → 0.21.0.dev202504180401__tar.gz

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.21.0.dev202504160404
+Version: 0.21.0.dev202504180401
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/falcon/falcon_backbone.py

@@ -29,7 +29,7 @@ class FalconBackbone(Backbone):
         layer_norm_epsilon: float. Epsilon for the layer normalization layers in
             the transformer decoder.
         attention_dropout_rate: float. Dropout probability for the attention.
-        feedforward_dropout_rate: flaot. Dropout probability for the
+        feedforward_dropout_rate: float. Dropout probability for the
             feedforward.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/gemma3/gemma3_causal_lm_preprocessor.py

@@ -512,6 +512,7 @@ class Gemma3CausalLMPreprocessor(CausalLMPreprocessor):
 
         # Extract text part of the input.
         prompts, responses = x["prompts"], x["responses"]
+        tf.debugging.assert_shapes([(prompts, ("N",)), (responses, ("N",))])
 
         # Find out if the input is batched/not batched. Uprank if not batched.
         # In other preprocessors, we don't have to do this, but here, all

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/llama/llama_attention.py

@@ -3,7 +3,9 @@ import math
 import keras
 from keras import ops
 
-from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
+from keras_hub.src.models.llama.llama_rotary_embedding import (
+    LlamaRotaryEmbedding,
+)
 from keras_hub.src.utils.keras_utils import clone_initializer
 from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
@@ -16,7 +18,11 @@ class LlamaAttention(keras.layers.Layer):
         num_query_heads,
         num_key_value_heads,
         rope_max_wavelength=10000,
-        rope_scaling_factor=1.0,
+        rope_position_scaling_factor=1.0,
+        rope_frequency_adjustment_factor=None,
+        rope_low_freq_factor=None,
+        rope_high_freq_factor=None,
+        rope_pretraining_sequence_length=None,
         kernel_initializer="glorot_uniform",
         dropout=0,
         **kwargs,
@@ -28,13 +34,16 @@ class LlamaAttention(keras.layers.Layer):
 
         self.num_key_value_groups = num_query_heads // num_key_value_heads
         self.rope_max_wavelength = rope_max_wavelength
+        self.rope_position_scaling_factor = rope_position_scaling_factor
+        self.rope_frequency_adjustment_factor = rope_frequency_adjustment_factor
+        self.rope_low_freq_factor = rope_low_freq_factor
+        self.rope_high_freq_factor = rope_high_freq_factor
+        self.rope_pretraining_sequence_length = rope_pretraining_sequence_length
 
         self.kernel_initializer = keras.initializers.get(
             clone_initializer(kernel_initializer)
         )
 
-        self.rope_scaling_factor = rope_scaling_factor
-
     def build(self, inputs_shape):
         # Einsum variables:
         # b = batch size
@@ -103,9 +112,13 @@ class LlamaAttention(keras.layers.Layer):
         )
         self._output_dense.build((None, None, self.num_query_heads, head_dim))
 
-        self.rotary_embedding_layer = RotaryEmbedding(
+        self.rotary_embedding_layer = LlamaRotaryEmbedding(
             max_wavelength=self.rope_max_wavelength,
-            scaling_factor=self.rope_scaling_factor,
+            position_scaling_factor=self.rope_position_scaling_factor,
+            frequency_adjustment_factor=self.rope_frequency_adjustment_factor,
+            low_freq_factor=self.rope_low_freq_factor,
+            high_freq_factor=self.rope_high_freq_factor,
+            pretraining_sequence_length=self.rope_pretraining_sequence_length,
             dtype=self.dtype_policy,
         )
 
@@ -224,6 +237,11 @@ class LlamaAttention(keras.layers.Layer):
                 "num_key_value_heads": self.num_key_value_heads,
                 "rope_max_wavelength": self.rope_max_wavelength,
                 "rope_scaling_factor": self.rope_scaling_factor,
+                "rope_low_freq_factor": self.rope_low_freq_factor,
+                "rope_high_freq_factor": self.rope_high_freq_factor,
+                "rope_pretraining_sequence_length": (
+                    self.rope_pretraining_sequence_length
+                ),
                 "kernel_initializer": keras.initializers.serialize(
                     self.kernel_initializer
                 ),

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/llama/llama_backbone.py

@@ -30,22 +30,30 @@ class LlamaBackbone(Backbone):
     constructor.
 
     Args:
-        vocabulary_size (int): The size of the token vocabulary.
-        num_layers (int): The number of transformer layers.
-        num_query_heads (int): The number of query attention heads for
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_query_heads : int.  The number of query attention heads for
             each transformer.
-        hidden_dim (int): The size of the transformer encoding and pooling
+        hidden_dim : int.  The size of the transformer encoding and pooling
             layers.
-        intermediate_dim (int): The output dimension of the first Dense layer in
+        intermediate_dim : int. The output dimension of the first Dense layer in
             a three-layer feedforward network for each transformer.
-        num_key_value_heads (int): The number of key and value attention heads
+        num_key_value_heads : int. The number of key and value attention heads
             for each transformer.
-        rope_max_wavelength (int, optional): The maximum angular wavelength of
+        rope_max_wavelength : int. The maximum angular wavelength of
             the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
-        rope_scaling_factor (float, optional): The scaling factor for
-            calculation of roatary embedding. Defaults to `1.0`.
-        layer_norm_epsilon (float, optional): Epsilon for the layer
-            normalization layers in the transformer decoder. Defaults to `1e-6`.
+        rope_position_scaling_factor: float. The scaling factor for
+            calculation of rotary embedding. Defaults to `1.0`
+        rope_frequency_adjustment_factor: float. The scaling factor
+            used to scale the inverse frequencies.  Defaults to `None`.
+        rope_low_freq_factor: float. The low frequency scaling
+            factor. Defaults to `None`.
+        rope_high_freq_factor: float. Used for Llama3.1+. The high
+            frequency scaling factor. Defaults to `None`.
+        rope_pretraining_sequence_length: int. Used for Llama3.1+.
+            Defaults to `None`.
+        layer_norm_epsilon : float. Epsilon for the layer normalization layers
+            in the transformer decoder. Defaults to `1e-6`.
         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
@@ -87,7 +95,11 @@ class LlamaBackbone(Backbone):
         intermediate_dim,
         num_key_value_heads,
         rope_max_wavelength=10000,
-        rope_scaling_factor=1.0,
+        rope_position_scaling_factor=1.0,
+        rope_frequency_adjustment_factor=None,
+        rope_low_freq_factor=None,
+        rope_high_freq_factor=None,
+        rope_pretraining_sequence_length=None,
         layer_norm_epsilon=1e-6,
         dropout=0,
         dtype=None,
@@ -110,7 +122,15 @@ class LlamaBackbone(Backbone):
                 num_query_heads=num_query_heads,
                 num_key_value_heads=num_key_value_heads,
                 rope_max_wavelength=rope_max_wavelength,
-                rope_scaling_factor=rope_scaling_factor,
+                rope_position_scaling_factor=rope_position_scaling_factor,
+                rope_frequency_adjustment_factor=(
+                    rope_frequency_adjustment_factor
+                ),
+                rope_low_freq_factor=rope_low_freq_factor,
+                rope_high_freq_factor=rope_high_freq_factor,
+                rope_pretraining_sequence_length=(
+                    rope_pretraining_sequence_length
+                ),
                 layer_norm_epsilon=layer_norm_epsilon,
                 activation=ops.silu,
                 kernel_initializer=_llama_kernel_initializer(stddev=0.02),
@@ -152,9 +172,13 @@ class LlamaBackbone(Backbone):
         self.num_query_heads = num_query_heads
         self.hidden_dim = hidden_dim
         self.intermediate_dim = intermediate_dim
-        self.rope_max_wavelength = rope_max_wavelength
         self.num_key_value_heads = num_key_value_heads
-        self.rope_scaling_factor = rope_scaling_factor
+        self.rope_max_wavelength = rope_max_wavelength
+        self.rope_position_scaling_factor = rope_position_scaling_factor
+        self.rope_frequency_adjustment_factor = rope_frequency_adjustment_factor
+        self.rope_low_freq_factor = rope_low_freq_factor
+        self.rope_high_freq_factor = rope_high_freq_factor
+        self.rope_pretraining_sequence_length = rope_pretraining_sequence_length
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
         self.tie_word_embeddings = tie_word_embeddings
@@ -169,7 +193,17 @@ class LlamaBackbone(Backbone):
                 "hidden_dim": self.hidden_dim,
                 "intermediate_dim": self.intermediate_dim,
                 "rope_max_wavelength": self.rope_max_wavelength,
-                "rope_scaling_factor": self.rope_scaling_factor,
+                "rope_position_scaling_factor": (
+                    self.rope_position_scaling_factor
+                ),
+                "rope_frequency_adjustment_factor": (
+                    self.rope_frequency_adjustment_factor
+                ),
+                "rope_low_freq_factor": self.rope_low_freq_factor,
+                "rope_high_freq_factor": self.rope_high_freq_factor,
+                "rope_pretraining_sequence_length": (
+                    self.rope_pretraining_sequence_length
+                ),
                 "num_key_value_heads": self.num_key_value_heads,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/llama/llama_decoder.py

@@ -21,7 +21,11 @@ class LlamaTransformerDecoder(keras.layers.Layer):
         num_query_heads,
         num_key_value_heads,
         rope_max_wavelength=10000,
-        rope_scaling_factor=1.0,
+        rope_position_scaling_factor=1.0,
+        rope_frequency_adjustment_factor=None,
+        rope_low_freq_factor=None,
+        rope_high_freq_factor=None,
+        rope_pretraining_sequence_length=None,
         activation="silu",
         layer_norm_epsilon=1e-5,
         kernel_initializer="glorot_uniform",
@@ -34,7 +38,11 @@ class LlamaTransformerDecoder(keras.layers.Layer):
         self.num_key_value_heads = num_key_value_heads
 
         self.rope_max_wavelength = rope_max_wavelength
-        self.rope_scaling_factor = rope_scaling_factor
+        self.rope_position_scaling_factor = rope_position_scaling_factor
+        self.rope_frequency_adjustment_factor = rope_frequency_adjustment_factor
+        self.rope_low_freq_factor = rope_low_freq_factor
+        self.rope_high_freq_factor = rope_high_freq_factor
+        self.rope_pretraining_sequence_length = rope_pretraining_sequence_length
 
         self.dropout = dropout
 
@@ -53,7 +61,11 @@ class LlamaTransformerDecoder(keras.layers.Layer):
             num_query_heads=self.num_query_heads,
             num_key_value_heads=self.num_key_value_heads,
             rope_max_wavelength=self.rope_max_wavelength,
-            rope_scaling_factor=self.rope_scaling_factor,
+            rope_position_scaling_factor=self.rope_position_scaling_factor,
+            rope_frequency_adjustment_factor=self.rope_frequency_adjustment_factor,
+            rope_low_freq_factor=self.rope_low_freq_factor,
+            rope_high_freq_factor=self.rope_high_freq_factor,
+            rope_pretraining_sequence_length=self.rope_pretraining_sequence_length,
             kernel_initializer=clone_initializer(self.kernel_initializer),
             dropout=self.dropout,
             dtype=self.dtype_policy,
@@ -221,6 +233,11 @@ class LlamaTransformerDecoder(keras.layers.Layer):
                 "num_query_heads": self.num_query_heads,
                 "rope_max_wavelength": self.rope_max_wavelength,
                 "rope_scaling_factor": self.rope_scaling_factor,
+                "rope_low_freq_factor": self.rope_low_freq_factor,
+                "rope_high_freq_factor": self.rope_high_freq_factor,
+                "rope_pretraining_sequence_length": (
+                    self.rope_pretraining_sequence_length
+                ),
                 "num_key_value_heads": self.num_key_value_heads,
                 "activation": keras.activations.serialize(self.activation),
                 "layer_norm_epsilon": self.layer_norm_epsilon,

keras_hub_nightly-0.21.0.dev202504180401/keras_hub/src/models/llama/llama_rotary_embedding.py

@@ -0,0 +1,180 @@
+import math
+
+from keras import ops
+
+from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
+
+
+class LlamaRotaryEmbedding(RotaryEmbedding):
+    """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.
+    It is extended from `RotaryEmbedding` layer in `keras_hub.layers`.
+    It has additional smoothening and interpolation for some frequency ranges.
+
+    Args:
+        max_wavelength: int. The maximum angular wavelength of the sine/cosine
+            curves. Defaults to `10000`.
+        position_scaling_factor: float. The scaling factor used to scale
+            positions of the tokens. Defaults to `1.0`.
+        frequency_adjustment_factor: float. The scaling factor used to scale the
+            inverse frequencies. Defaults to `None`.
+        low_freq_factor: float. The low frequency scaling factor.
+            Defaults to `None`.
+        high_freq_factor: float. The high frequency scaling factor.
+            Defaults to `None`.
+        pretraining_sequence_length: int. Used for Llama3.1+, the original
+            context length at time of pretraining. Defaults to `None`.
+        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,
+        position_scaling_factor=1.0,
+        sequence_axis=1,
+        feature_axis=-1,
+        frequency_adjustment_factor=None,
+        low_freq_factor=None,
+        high_freq_factor=None,
+        pretraining_sequence_length=None,
+        **kwargs,
+    ):
+        super().__init__(
+            max_wavelength=max_wavelength,
+            scaling_factor=position_scaling_factor,
+            sequence_axis=sequence_axis,
+            feature_axis=feature_axis,
+            **kwargs,
+        )
+        self.max_wavelength = max_wavelength
+        self.sequence_axis = sequence_axis
+        self.feature_axis = feature_axis
+        self.position_scaling_factor = position_scaling_factor
+        self.frequency_adjustment_factor = frequency_adjustment_factor
+        self.low_freq_factor = low_freq_factor
+        self.high_freq_factor = high_freq_factor
+        self.pretraining_sequence_length = pretraining_sequence_length
+
+        grouped_args = [
+            low_freq_factor,
+            high_freq_factor,
+            frequency_adjustment_factor,
+            pretraining_sequence_length,
+        ]
+        args_none = [x is None for x in grouped_args]
+        if any(args_none) and not all(args_none):
+            raise ValueError(
+                "Either all of `low_freq_factor`,`high_freq_factor`, "
+                "`frequency_adjustment_factor` and "
+                "`pretraining_sequence_length` should be set, or all of should"
+                " be set `None`."
+            )
+        self.built = True
+
+    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)
+
+        # From llama3.1+ we have additional smoothening and interpolation.
+        # low_freq_factor, high_freq_factor, pretraining_sequence_length,
+        # frequency_adjustment_factor are all set at once so it is fine.
+        if self.low_freq_factor is not None:
+            low_freq_wavelen = (
+                self.pretraining_sequence_length / self.low_freq_factor
+            )
+            high_freq_wavelen = (
+                self.pretraining_sequence_length / self.high_freq_factor
+            )
+            wavelen = 2 * math.pi / inverse_freq
+
+            # wavelen < high_freq_wavelen: do nothing
+            # wavelen > low_freq_wavelen: divide by factor
+            inverse_freq = ops.where(
+                ops.greater(wavelen, low_freq_wavelen),
+                (inverse_freq / self.frequency_adjustment_factor),
+                inverse_freq,
+            )
+
+            # otherwise: interpolate between the two, using a smooth factor
+            smooth_factor = (
+                (self.pretraining_sequence_length / wavelen)
+                - self.low_freq_factor
+            ) / (self.high_freq_factor - self.low_freq_factor)
+            smoothed_inv_freq = (1 - smooth_factor) * (
+                inverse_freq / self.frequency_adjustment_factor
+            ) + (smooth_factor * inverse_freq)
+            is_medium_freq = ops.logical_and(
+                ops.greater_equal(wavelen, high_freq_wavelen),
+                ops.less_equal(wavelen, low_freq_wavelen),
+            )
+
+            inverse_freq = ops.where(
+                is_medium_freq, smoothed_inv_freq, inverse_freq
+            )
+
+        return inverse_freq
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_wavelength": self.max_wavelength,
+                "sequence_axis": self.sequence_axis,
+                "feature_axis": self.feature_axis,
+                "position_scaling_factor": self.position_scaling_factor,
+                "frequency_adjustment_factor": self.frequency_adjustment_factor,
+                "low_freq_factor": self.low_freq_factor,
+                "high_freq_factor": self.high_freq_factor,
+                "original_max_embeddings": self.pretraining_sequence_length,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/models/llama3/llama3_backbone.py

@@ -32,8 +32,16 @@ class Llama3Backbone(LlamaBackbone):
             fo each transformer.
         rope_max_wavelength (int, optional): The maximum angular wavelength of
             the sine/cosine curves, for rotary embeddings. Defaults to `10000`.
-        rope_scaling_factor (float, optional): The scaling factor for
-            calculation of roatary embedding. Defaults to `1.0`.
+        rope_position_scaling_factor (float, optional): The scaling factor for
+            calculation of roatary embedding. Defaults to `1.0`
+        rope_requency_adjustment_factor (float, optional): The scaling factor
+            used to scale the inverse frequencies.
+        rope_low_freq_factor (float, optional): The low frequency factor.
+            Defaults to None.
+        rope_high_freq_factor: (float, optional) Used for Llama3.1+. The high
+            frequency factor. Defaults to None.
+        rope_pretraining_sequence_length: (int, optional) Sequence length during
+            original pretraining. Defaults to None.
         layer_norm_epsilon (float, optional): Epsilon for the layer
             normalization layers in the transformer decoder. Defaults to `1e-6`.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use

keras_hub_nightly-0.21.0.dev202504180401/keras_hub/src/models/llama3/llama3_presets.py

@@ -0,0 +1,130 @@
+"""Llama 3 model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "llama3_8b_en": {
+        "metadata": {
+            "description": "8 billion parameter, 32-layer, base LLaMA 3 model.",
+            "params": 8030261248,
+            "path": "llama3",
+        },
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en/4",
+    },
+    "llama3_8b_en_int8": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, base LLaMA 3 model with "
+                "activation and weights quantized to int8."
+            ),
+            "params": 8031894016,
+            "path": "llama3",
+        },
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en_int8/2",
+    },
+    "llama3_instruct_8b_en": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, instruction tuned LLaMA 3 "
+                "model."
+            ),
+            "params": 8030261248,
+            "path": "llama3",
+        },
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_instruct_8b_en/4",
+    },
+    "llama3_instruct_8b_en_int8": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, instruction tuned LLaMA 3 "
+                "model with activation and weights quantized to int8."
+            ),
+            "params": 8031894016,
+            "path": "llama3",
+        },
+        "kaggle_handle": (
+            "kaggle://keras/llama3/keras/llama3_instruct_8b_en_int8/2"
+        ),
+    },
+    "llama3.1_8b": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, based LLaMA 3.1 model. "
+            ),
+            "params": 8030261248,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.1_8b/1"),
+    },
+    "llama3.1_instruct_8b": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, instruction tuned LLaMA 3.1. "
+            ),
+            "params": 8030261248,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/lama3.1_instruct_8b/1"),
+    },
+    "llama3.1_guard_8b": {
+        "metadata": {
+            "description": (
+                "8 billion parameter, 32-layer, LLaMA 3.1 fine-tuned for "
+                "consent safety classification. "
+            ),
+            "params": 8030261248,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.1_guard_8b/1"),
+    },
+    "llama3.2_1b": {
+        "metadata": {
+            "description": (
+                "1 billion parameter, 16-layer, based LLaMA 3.2 model. "
+            ),
+            "params": 1498482688,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.2_1b/1"),
+    },
+    "llama3.2_instruct_1b": {
+        "metadata": {
+            "description": (
+                "1 billion parameter, 16-layer, instruction tuned LLaMA 3.2. "
+            ),
+            "params": 1498482688,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.2_instruct_1b/1"),
+    },
+    "llama3.2_3b": {
+        "metadata": {
+            "description": (
+                "3 billion parameter, 26-layer, based LLaMA 3.2 model. "
+            ),
+            "params": 3606752256,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.2_3b/1"),
+    },
+    "llama3.2_instruct_3b": {
+        "metadata": {
+            "description": (
+                "3 billion parameter, 28-layer, instruction tuned LLaMA 3.2. "
+            ),
+            "params": 3606752256,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.2_instruct_3b/1"),
+    },
+    "llama3.2_guard_1b": {
+        "metadata": {
+            "description": (
+                "1 billion parameter, 16-layer, based LLaMA 3.2 model "
+                "fine-tuned for consent safety classification. "
+            ),
+            "params": 1498482688,
+            "path": "llama3",
+        },
+        "kaggle_handle": ("kaggle://keras/llama3/keras/llama3.2_guard_1b/1"),
+    },
+}

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/utils/transformers/convert_llama3.py

@@ -7,7 +7,7 @@ backbone_cls = Llama3Backbone
 
 
 def convert_backbone_config(transformers_config):
-    return {
+    backbone_config = {
         "vocabulary_size": transformers_config["vocab_size"],
         "num_layers": transformers_config["num_hidden_layers"],
         "num_query_heads": transformers_config["num_attention_heads"],
@@ -15,8 +15,28 @@ def convert_backbone_config(transformers_config):
         "intermediate_dim": transformers_config["intermediate_size"],
         "num_key_value_heads": transformers_config["num_key_value_heads"],
         "tie_word_embeddings": transformers_config["tie_word_embeddings"],
+        "rope_max_wavelength": transformers_config["rope_theta"],
     }
 
+    if transformers_config.get("rope_scaling", None) is not None:
+        if transformers_config["rope_scaling"]["rope_type"] != "llama3":
+            raise ValueError("The config should be a valid llama3 config.")
+        backbone_config["rope_frequency_adjustment_factor"] = (
+            transformers_config["rope_scaling"]["factor"]
+        )
+        backbone_config["rope_low_freq_factor"] = transformers_config[
+            "rope_scaling"
+        ]["low_freq_factor"]
+        backbone_config["rope_high_freq_factor"] = transformers_config[
+            "rope_scaling"
+        ]["high_freq_factor"]
+        backbone_config["rope_pretraining_sequence_length"] = (
+            transformers_config["rope_scaling"][
+                "original_max_position_embeddings"
+            ]
+        )
+    return backbone_config
+
 
 def convert_weights(backbone, loader, transformers_config):
     loader.port_weight(

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.21.0.dev202504160404"
+__version__ = "0.21.0.dev202504180401"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.21.0.dev202504160404
+Version: 0.21.0.dev202504180401
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/keras_hub_nightly.egg-info/SOURCES.txt

@@ -232,6 +232,7 @@ keras_hub/src/models/llama/llama_causal_lm_preprocessor.py
 keras_hub/src/models/llama/llama_decoder.py
 keras_hub/src/models/llama/llama_layernorm.py
 keras_hub/src/models/llama/llama_presets.py
+keras_hub/src/models/llama/llama_rotary_embedding.py
 keras_hub/src/models/llama/llama_tokenizer.py
 keras_hub/src/models/llama3/__init__.py
 keras_hub/src/models/llama3/llama3_backbone.py

{keras_hub_nightly-0.21.0.dev202504160404 → keras_hub_nightly-0.21.0.dev202504180401}/setup.py

@@ -23,7 +23,7 @@ def get_version(rel_path):
 
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
-VERSION = "0.21.0.dev202504160404"  # get_version("keras_hub/src/version_utils.py")
+VERSION = "0.21.0.dev202504180401"  # get_version("keras_hub/src/version_utils.py")
 
 setup(
     name="keras-hub-nightly",  # "keras-hub",