Add presets for gemma 2 2b (#1721)

keras-team · Jul 31, 2024 · b5375c9 · b5375c9
1 parent 39802c7
commit b5375c9
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/keras_nlp/src/models/gemma/gemma_presets.py b/keras_nlp/src/models/gemma/gemma_presets.py
@@ -155,6 +155,26 @@
         },
         "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.",