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

keras_hub/src/models/gemma/gemma_decoder_block.py

@@ -0,0 +1,241 @@
+# 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.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_hub.src.models.gemma.gemma_attention import CachedGemmaAttention
+from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
+
+
+class GemmaDecoderBlock(keras.layers.Layer):
+    def __init__(
+        self,
+        hidden_dim,
+        intermediate_dim,
+        head_dim,
+        num_query_heads,
+        num_key_value_heads,
+        query_head_dim_normalize=True,
+        use_post_ffw_norm=False,
+        use_post_attention_norm=False,
+        logit_soft_cap=None,
+        use_sliding_window_attention=False,
+        sliding_window_size=4096,
+        layer_norm_epsilon=1e-6,
+        dropout=0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.intermediate_dim = intermediate_dim
+        self.hidden_dim = hidden_dim
+        self.num_query_heads = num_query_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.dropout = dropout
+        self.query_head_dim_normalize = query_head_dim_normalize
+        self.use_post_ffw_norm = use_post_ffw_norm
+        self.use_post_attention_norm = use_post_attention_norm
+        self.logit_soft_cap = logit_soft_cap
+        self.use_sliding_window_attention = use_sliding_window_attention
+        self.sliding_window_size = sliding_window_size
+
+        self.pre_attention_norm = RMSNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="pre_attention_norm",
+        )
+
+        if use_post_attention_norm:
+            self.post_attention_norm = RMSNormalization(
+                epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name="pre_attention_norm",
+            )
+
+        self.attention = CachedGemmaAttention(
+            head_dim=head_dim,
+            num_query_heads=num_query_heads,
+            num_key_value_heads=num_key_value_heads,
+            logit_soft_cap=logit_soft_cap,
+            use_sliding_window_attention=use_sliding_window_attention,
+            sliding_window_size=sliding_window_size,
+            query_head_dim_normalize=True,
+            dropout=dropout,
+            dtype=self.dtype_policy,
+            name="attention",
+        )
+
+        if self.dropout > 0:
+            self.attention_dropout = keras.layers.Dropout(rate=dropout)
+            self.feedforward_dropout = keras.layers.Dropout(rate=dropout)
+
+        self.pre_ffw_norm = RMSNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="pre_ffw_norm",
+        )
+
+        if use_post_ffw_norm:
+            self.post_ffw_norm = RMSNormalization(
+                epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name="post_ffw_norm",
+            )
+
+        self.gating_ffw = keras.layers.EinsumDense(
+            equation="btd,df->btf",
+            output_shape=(None, self.intermediate_dim // 2),
+            dtype=self.dtype_policy,
+            name="ffw_gating",
+        )
+
+        self.gating_ffw_2 = keras.layers.EinsumDense(
+            equation="btd,df->btf",
+            output_shape=(None, self.intermediate_dim // 2),
+            dtype=self.dtype_policy,
+            name="ffw_gating_2",
+        )
+
+        self.ffw_linear = keras.layers.EinsumDense(
+            equation="btf,fd->btd",
+            output_shape=(None, self.hidden_dim),
+            dtype=self.dtype_policy,
+            name="ffw_linear",
+        )
+
+    def build(self, input_shape):
+        self.pre_attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+
+        if self.use_post_attention_norm:
+            shape = self.attention.compute_output_shape(input_shape)
+            self.post_attention_norm.build(shape)
+
+        shape = input_shape
+        self.pre_ffw_norm.build(shape)
+        self.gating_ffw.build(shape)
+        self.gating_ffw_2.build(shape)
+
+        shape = self.gating_ffw.compute_output_shape(shape)
+        self.ffw_linear.build(shape)
+
+        if self.use_post_ffw_norm:
+            shape = self.ffw_linear.compute_output_shape(shape)
+            self.post_ffw_norm.build(shape)
+
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        # Isometric
+        return input_shape
+
+    def _compute_attention_mask(
+        self, x, padding_mask, cache, cache_update_index
+    ):
+        decoder_mask = merge_padding_and_attention_mask(
+            inputs=x, padding_mask=padding_mask, attention_mask=None
+        )
+        batch_size = ops.shape(x)[0]
+        input_length = output_length = ops.shape(x)[1]
+        if cache is not None:
+            input_length = ops.shape(cache)[2]
+
+        causal_mask = compute_causal_mask(
+            batch_size=batch_size,
+            input_length=input_length,
+            output_length=output_length,
+            cache_index=cache_update_index,
+        )
+
+        return (
+            ops.minimum(decoder_mask, causal_mask)
+            if decoder_mask is not None
+            else causal_mask
+        )
+
+    def call(
+        self,
+        x,
+        padding_mask=None,
+        cache=None,
+        cache_update_index=0,
+    ):
+        normalized_x = self.pre_attention_norm(x)
+        attention_mask = self._compute_attention_mask(
+            normalized_x, padding_mask, cache, cache_update_index
+        )
+        if cache is not None:
+            attention, new_cache = self.attention(
+                normalized_x,
+                attention_mask=attention_mask,
+                cache=cache,
+                cache_update_index=cache_update_index,
+            )
+        else:
+            attention = self.attention(
+                normalized_x,
+                attention_mask=attention_mask,
+            )
+
+        if self.use_post_attention_norm:
+            attention = self.post_attention_norm(attention)
+
+        if self.dropout:
+            attention = self.attention_dropout(attention)
+
+        attention_x = x + attention
+        normalized_x = self.pre_ffw_norm(attention_x)
+
+        x1 = self.gating_ffw(normalized_x)
+        x2 = self.gating_ffw_2(normalized_x)
+        x = keras.activations.gelu(x1, approximate=True) * x2
+        x = self.ffw_linear(x)
+
+        if self.use_post_ffw_norm:
+            x = self.post_ffw_norm(x)
+
+        x = x + attention_x
+
+        if cache is not None:
+            return x, new_cache
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "head_dim": self.head_dim,
+                "num_query_heads": self.num_query_heads,
+                "num_key_value_heads": self.num_key_value_heads,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "dropout": self.dropout,
+                "use_post_ffw_norm": self.use_post_ffw_norm,
+                "use_post_attention_norm": self.use_post_attention_norm,
+                "logit_soft_cap": self.logit_soft_cap,
+                "use_sliding_window_attention": self.use_sliding_window_attention,
+                "sliding_window_size": self.sliding_window_size,
+                "query_head_dim_normalize": self.query_head_dim_normalize,
+            }
+        )
+        return config

keras_hub/src/models/gemma/gemma_preprocessor.py

@@ -0,0 +1,191 @@
+# 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.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+
+
+@keras_hub_export("keras_hub.models.GemmaPreprocessor")
+class GemmaPreprocessor(Preprocessor):
+    """Gemma preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do 2 things:
+
+    - Tokenize the inputs using the `tokenizer`.
+    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
+        be passed directly to a `keras_hub.models.GemmaBackbone`.
+
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    The call method of this layer accepts three arguments, `x`, `y`, and
+    `sample_weight`. `x` can be a python string or tensor representing a single
+    segment, a list of python strings representing a batch of single segments,
+    or a list of tensors representing multiple segments to be packed together.
+    `y` and `sample_weight` are both optional, can have any format, and will be
+    passed through unaltered.
+
+    `GemmaPreprocessor` expects the input to have only one segment, as Gemma is
+    mainly used for generation tasks. For tasks having multi-segment inputs
+    please combine inputs into a single string input before passing to the
+    preprocessor layer.
+
+    Args:
+        tokenizer: A `keras_hub.models.GemmaTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence.
+        add_end_token: If `True`, the preprocessor will append the tokenizer
+            end token to each input sequence.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through unaltered.
+        sequence_length: Pass to override the configured `sequence_length` of
+            the layer.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    preprocessor = keras_hub.models.GemmaPreprocessor.from_preset(
+        "gemma_2b_en"
+    )
+
+    # Tokenize and pack a single sentence.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize a batch of sentences.
+    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
+
+    # Custom vocabulary.
+    bytes_io = io.BytesIO()
+    ds = tf.data.Dataset.from_tensor_slices(["The quick brown fox jumped."])
+    sentencepiece.SentencePieceTrainer.train(
+        sentence_iterator=ds.as_numpy_iterator(),
+        model_writer=bytes_io,
+        vocab_size=8,
+        model_type="WORD",
+        pad_id=0,
+        bos_id=1,
+        eos_id=2,
+        unk_id=3,
+        pad_piece="<pad>",
+        bos_piece="<bos>",
+        eos_piece="<eos>",
+        unk_piece="<unk>",
+    )
+    tokenizer = keras_hub.models.GemmaTokenizer(
+        proto=bytes_io.getvalue(),
+    )
+    preprocessor = keras_hub.models.GemmaPreprocessor(tokenizer=tokenizer)
+    preprocessor("The quick brown fox jumped.")
+    ```
+
+    Apply preprocessing to a `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.GemmaPreprocessor.from_preset(
+        "gemma_2b_en"
+    )
+
+    text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    label = tf.constant([1, 1])
+
+    # Map labeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices((text, label))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map unlabeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(text)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    tokenizer_cls = GemmaTokenizer
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=1024,
+        add_start_token=True,
+        add_end_token=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.tokenizer = tokenizer
+        self.sequence_length = sequence_length
+        self.add_start_token = add_start_token
+        self.add_end_token = add_end_token
+
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.packer = StartEndPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+        )
+        self.built = True
+
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        x = convert_inputs_to_list_of_tensor_segments(x)
+        if len(x) != 1:
+            raise ValueError(
+                "GemmaPreprocessor requires each input to contain only "
+                f"one segment, but received {len(x)}. If you are using Gemma "
+                "for a multi-segment classification task, please combine your "
+                "input into a single string."
+            )
+        sequence_length = sequence_length or self.sequence_length
+        token_ids, padding_mask = self.packer(
+            self.tokenizer(x[0]),
+            sequence_length=sequence_length,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "add_start_token": self.add_start_token,
+                "add_end_token": self.add_end_token,
+            }
+        )
+        return config

keras_hub/src/models/gemma/gemma_presets.py

@@ -0,0 +1,248 @@
+# 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.
+"""Gemma model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "gemma_2b_en": {
+        "metadata": {
+            "description": "2 billion parameter, 18-layer, base Gemma model.",
+            "params": 2506172416,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_2b_en/2",
+    },
+    "gemma_instruct_2b_en": {
+        "metadata": {
+            "description": (
+                "2 billion parameter, 18-layer, instruction tuned Gemma model."
+            ),
+            "params": 2506172416,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_2b_en/2",
+    },
+    "gemma_1.1_instruct_2b_en": {
+        "metadata": {
+            "description": (
+                "2 billion parameter, 18-layer, instruction tuned Gemma model. "
+                "The 1.1 update improves model quality."
+            ),
+            "params": 2506172416,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_2b_en/3",
+    },
+    "code_gemma_1.1_2b_en": {
+        "metadata": {
+            "description": (
+                "2 billion parameter, 18-layer, CodeGemma model. This model "
+                "has been trained on a fill-in-the-middle (FIM) task for code "
+                "completion. The 1.1 update improves model quality."
+            ),
+            "params": 2506172416,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_2b_en/1",
+    },
+    "code_gemma_2b_en": {
+        "metadata": {
+            "description": (
+                "2 billion parameter, 18-layer, CodeGemma model. This model "
+                "has been trained on a fill-in-the-middle (FIM) task for code "
+                "completion."
+            ),
+            "params": 2506172416,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_2b_en/1",
+    },
+    "gemma_7b_en": {
+        "metadata": {
+            "description": "7 billion parameter, 28-layer, base Gemma model.",
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_7b_en/2",
+    },
+    "gemma_instruct_7b_en": {
+        "metadata": {
+            "description": (
+                "7 billion parameter, 28-layer, instruction tuned Gemma model."
+            ),
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_7b_en/2",
+    },
+    "gemma_1.1_instruct_7b_en": {
+        "metadata": {
+            "description": (
+                "7 billion parameter, 28-layer, instruction tuned Gemma model. "
+                "The 1.1 update improves model quality."
+            ),
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma/keras/gemma_1.1_instruct_7b_en/3",
+    },
+    "code_gemma_7b_en": {
+        "metadata": {
+            "description": (
+                "7 billion parameter, 28-layer, CodeGemma model. This model "
+                "has been trained on a fill-in-the-middle (FIM) task for code "
+                "completion."
+            ),
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_7b_en/1",
+    },
+    "code_gemma_instruct_7b_en": {
+        "metadata": {
+            "description": (
+                "7 billion parameter, 28-layer, instruction tuned CodeGemma "
+                "model. This model has been trained for chat use cases related "
+                "to code."
+            ),
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_instruct_7b_en/1",
+    },
+    "code_gemma_1.1_instruct_7b_en": {
+        "metadata": {
+            "description": (
+                "7 billion parameter, 28-layer, instruction tuned CodeGemma "
+                "model. This model has been trained for chat use cases related "
+                "to code. The 1.1 update improves model quality."
+            ),
+            "params": 8537680896,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/codegemma/keras/code_gemma_1.1_instruct_7b_en/1",
+    },
+    "gemma2_2b_en": {
+        "metadata": {
+            "description": "2 billion parameter, 26-layer, base Gemma model.",
+            "params": 2614341888,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_2b_en/1",
+    },
+    "gemma2_instruct_2b_en": {
+        "metadata": {
+            "description": "2 billion parameter, 26-layer, instruction tuned Gemma model.",
+            "params": 2614341888,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_2b_en/1",
+    },
+    "gemma2_9b_en": {
+        "metadata": {
+            "description": "9 billion parameter, 42-layer, base Gemma model.",
+            "params": 9241705984,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_9b_en/2",
+    },
+    "gemma2_instruct_9b_en": {
+        "metadata": {
+            "description": "9 billion parameter, 42-layer, instruction tuned Gemma model.",
+            "params": 9241705984,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_9b_en/2",
+    },
+    "gemma2_27b_en": {
+        "metadata": {
+            "description": "27 billion parameter, 42-layer, base Gemma model.",
+            "params": 27227128320,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_27b_en/1",
+    },
+    "gemma2_instruct_27b_en": {
+        "metadata": {
+            "description": "27 billion parameter, 42-layer, instruction tuned Gemma model.",
+            "params": 27227128320,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://keras/gemma2/keras/gemma2_instruct_27b_en/1",
+    },
+    "shieldgemma_2b_en": {
+        "metadata": {
+            "description": "2 billion parameter, 26-layer, ShieldGemma model.",
+            "params": 2614341888,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_2b_en/1",
+    },
+    "shieldgemma_9b_en": {
+        "metadata": {
+            "description": "9 billion parameter, 42-layer, ShieldGemma model.",
+            "params": 9241705984,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_9b_en/1",
+    },
+    "shieldgemma_27b_en": {
+        "metadata": {
+            "description": "27 billion parameter, 42-layer, ShieldGemma model.",
+            "params": 27227128320,
+            "official_name": "Gemma",
+            "path": "gemma",
+            "model_card": "https://www.kaggle.com/models/google/gemma",
+        },
+        "kaggle_handle": "kaggle://google/shieldgemma/keras/shieldgemma_27b_en/1",
+    },
+}