Shield gemma shortnames (#1731)

keras-team · Aug 2, 2024 · c1afb07 · c1afb07
1 parent 765c402
commit c1afb07
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/keras_nlp/src/models/gemma/gemma_presets.py b/keras_nlp/src/models/gemma/gemma_presets.py
@@ -215,4 +215,34 @@
         },
         "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",
+    },
 }