keras-hub 0.20.0.dev1__py3-none-any.whl → 0.21.0.dev1__py3-none-any.whl

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -55,7 +55,7 @@ backbone_presets = {
             "params": 11765788416,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_12b_text/2",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_12b_text/3",
     },
     "gemma3_instruct_12b_text": {
         "metadata": {
@@ -66,7 +66,7 @@ backbone_presets = {
             "params": 11765788416,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_12b_text/2",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_12b_text/3",
     },
     "gemma3_27b_text": {
         "metadata": {
@@ -77,7 +77,7 @@ backbone_presets = {
             "params": 27009002240,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_27b_text/3",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_27b_text/4",
     },
     "gemma3_instruct_27b_text": {
         "metadata": {
@@ -88,7 +88,7 @@ backbone_presets = {
             "params": 27009002240,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_27b_text/2",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_27b_text/3",
     },
     "gemma3_4b": {
         "metadata": {
@@ -121,7 +121,7 @@ backbone_presets = {
             "params": 12187079280,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_12b/1",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_12b/2",
     },
     "gemma3_instruct_12b": {
         "metadata": {
@@ -132,7 +132,7 @@ backbone_presets = {
             "params": 12187079280,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_12b/1",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_12b/2",
     },
     "gemma3_27b": {
         "metadata": {
@@ -143,7 +143,7 @@ backbone_presets = {
             "params": 27432062576,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_27b/1",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_27b/2",
     },
     "gemma3_instruct_27b": {
         "metadata": {
@@ -154,6 +154,6 @@ backbone_presets = {
             "params": 27432062576,
             "path": "gemma3",
         },
-        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_27b/1",
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_27b/2",
     },
 }

keras_hub/src/models/gemma3/gemma3_vision_encoder.py

@@ -488,7 +488,7 @@ class Gemma3VisionEncoderBlock(keras.layers.Layer):
             # Fix the compatibility issue with Keras 3.1 where
             # `compute_output_spec` fails to propagate `inputs_shape`
             # correctly, causing it to be `None`.
-            inputs_shape = [None, None, None]
+            return [None, None, self.hidden_dim]
         return [
             None,
             (inputs_shape[2] // self.patch_size) ** 2,

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/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/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/src/models/llama/llama_presets.py

@@ -8,7 +8,7 @@ backbone_presets = {
             "params": 6738415616,
             "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en/2",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_7b_en/3",
     },
     "llama2_7b_en_int8": {
         "metadata": {
@@ -30,7 +30,7 @@ backbone_presets = {
             "params": 6738415616,
             "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en/2",
+        "kaggle_handle": "kaggle://keras/llama2/keras/llama2_instruct_7b_en/3",
     },
     "llama2_instruct_7b_en_int8": {
         "metadata": {
@@ -52,6 +52,6 @@ backbone_presets = {
             "params": 6738415616,
             "path": "llama",
         },
-        "kaggle_handle": "kaggle://keras/vicuna/keras/vicuna_1.5_7b_en/2",
+        "kaggle_handle": "kaggle://keras/vicuna/keras/vicuna_1.5_7b_en/3",
     },
 }

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/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/src/models/llama3/llama3_presets.py

@@ -8,7 +8,7 @@ backbone_presets = {
             "params": 8030261248,
             "path": "llama3",
         },
-        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en/4",
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_8b_en/5",
     },
     "llama3_8b_en_int8": {
         "metadata": {
@@ -30,7 +30,7 @@ backbone_presets = {
             "params": 8030261248,
             "path": "llama3",
         },
-        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_instruct_8b_en/4",
+        "kaggle_handle": "kaggle://keras/llama3/keras/llama3_instruct_8b_en/5",
     },
     "llama3_instruct_8b_en_int8": {
         "metadata": {
@@ -45,4 +45,86 @@ backbone_presets = {
             "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/src/models/mistral/mistral_presets.py

@@ -8,7 +8,7 @@ backbone_presets = {
             "params": 7241732096,
             "path": "mistral",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/7",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/8",
     },
     "mistral_instruct_7b_en": {
         "metadata": {
@@ -16,7 +16,7 @@ backbone_presets = {
             "params": 7241732096,
             "path": "mistral",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/7",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/8",
     },
     "mistral_0.2_instruct_7b_en": {
         "metadata": {
@@ -24,6 +24,6 @@ backbone_presets = {
             "params": 7241732096,
             "path": "mistral",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/2",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_0.2_instruct_7b_en/3",
     },
 }

keras_hub/src/models/mixtral/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.mixtral.mixtral_backbone import MixtralBackbone
+from keras_hub.src.models.mixtral.mixtral_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, MixtralBackbone)