Updates presets for Gemma 1.1 (keras-team#1556)

SamanehSaadat · Apr 5, 2024 · 1c0c02a · 1c0c02a
1 parent b2379d9
commit 1c0c02a
Showing 1 changed file with 29 additions and 7 deletions.
diff --git a/keras_nlp/models/gemma/gemma_presets.py b/keras_nlp/models/gemma/gemma_presets.py
@@ -17,9 +17,7 @@
 backbone_presets = {
     "gemma_2b_en": {
         "metadata": {
-            "description": (
-                "18-layer Gemma model (Gemma with 2B parameters). "
-            ),
+            "description": "2 billion parameter, 18-layer, base Gemma model.",
             "params": 2506172416,
             "official_name": "Gemma",
             "path": "gemma",
@@ -30,7 +28,7 @@
     "gemma_instruct_2b_en": {
         "metadata": {
             "description": (
-                "18-layer Gemma model (Gemma with 2B parameters). "
+                "2 billion parameter, 18-layer, instruction tuned Gemma model."
             ),
             "params": 2506172416,
             "official_name": "Gemma",
@@ -39,11 +37,22 @@
         },
         "kaggle_handle": "kaggle://keras/gemma/keras/gemma_instruct_2b_en/2",
     },
-    "gemma_7b_en": {
+    "gemma_1.1_instruct_2b_en": {
         "metadata": {
             "description": (
-                "28-layer Gemma model (Gemma with 7B parameters). "
+                "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/1",
+    },
+    "gemma_7b_en": {
+        "metadata": {
+            "description": "7 billion parameter, 28-layer, base Gemma model.",
             "params": 8537680896,
             "official_name": "Gemma",
             "path": "gemma",
@@ -54,7 +63,7 @@
     "gemma_instruct_7b_en": {
         "metadata": {
             "description": (
-                "28-layer Gemma model (Gemma with 7B parameters). "
+                "7 billion parameter, 28-layer, instruction tuned Gemma model."
             ),
             "params": 8537680896,
             "official_name": "Gemma",
@@ -63,4 +72,17 @@
         },
         "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/1",
+    },
 }