Add large variations of electra in presets

keras-team · Mar 23, 2024 · 2889350 · 2889350
1 parent b268e26
commit 2889350
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/keras_nlp/models/electra/electra_presets.py b/keras_nlp/models/electra/electra_presets.py
@@ -66,4 +66,30 @@
         },
         "kaggle_handle": "kaggle://pranavprajapati16/electra/keras/electra_base_generator_en/1",
     },
+    "electra_large_generator_en": {
+        "metadata": {
+            "description": (
+                "ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators"
+                "This is large generator model with 24 layers."
+            ),
+            "params": 51065344,
+            "official_name": "ELECTRA",
+            "path": "electra",
+            "model_card": "https://github.com/google-research/electra",
+        },
+        "kaggle_handle": "kaggle://pranavprajapati16/electra/keras/electra_large_generator_en/1",
+    },
+    "electra_large_discriminator_en": {
+        "metadata": {
+            "description": (
+                "ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators"
+                "This is large generator model with 24 layers."
+            ),
+            "params": 335141888,
+            "official_name": "ELECTRA",
+            "path": "electra",
+            "model_card": "https://github.com/google-research/electra",
+        },
+        "kaggle_handle": "kaggle://pranavprajapati16/electra/keras/electra_large_discriminator_en/1",
+    },
 }
diff --git a/tools/checkpoint_conversion/convert_electra_checkpoints.py b/tools/checkpoint_conversion/convert_electra_checkpoints.py
@@ -37,6 +37,8 @@
     "electra_small_generator_en": "google/electra-small-generator",
     "electra_base_discriminator_en": "google/electra-base-discriminator",
     "electra_small_discriminator_en": "google/electra-small-discriminator",
+    "electra_large_discriminator_en": "google/electra-large-discriminator",
+    "electra_large_generator_en": "google/electra-large-generator",
 }
 
 EXTRACT_DIR = "./model"