Register Embedding Gemma3 presets (keras-team#2598)

* Register presets

* Fix format

Author: laxmareddyp

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -292,4 +292,15 @@
         },
         "kaggle_handle": "kaggle://keras/translategemma/keras/translategemma_27b_it/1",
     },
+    # Embedding Gemma presets
+    "embedding_gemma3_300m": {
+        "metadata": {
+            "description": (
+                "Embedding-focused Gemma 3 model (300M parameters, 24 layers)."
+            ),
+            "params": 307581696,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/embeddinggemma/keras/embedding_gemma3_300m/2",
+    },
 }

tools/checkpoint_conversion/convert_gemma3_embedding_checkpoints.py

@@ -20,6 +20,7 @@
 
 import json
 import os
+import tempfile
 
 os.environ["KERAS_BACKEND"] = "torch"
 os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
@@ -257,69 +258,77 @@ def main(_):
     print(f"Converting: {hf_model_id} -> {preset}")
     print(f"{'=' * 60}\n")
 
-    # Download config to get parameters
-    print("Downloading config...")
-    config_path = hf_hub_download(hf_model_id, "config.json")
-    with open(config_path, "r") as f:
-        transformers_config = json.load(f)
-
-    # Convert config
-    keras_config = convert_backbone_config(transformers_config)
-
-    # Create KerasHub model
-    print("\nCreating KerasHub Gemma3Backbone...")
-    # Use float32 to match HF model default often
-    keras_model = keras_hub.models.Gemma3Backbone(
-        **keras_config, dtype="float32"
-    )
+    # Use a temporary directory for downloading HuggingFace artifacts
+    # This keeps the output preset directory clean
+    with tempfile.TemporaryDirectory() as temp_dir:
+        print(f"Using temporary directory for downloads: {temp_dir}")
+
+        # Download config to get parameters
+        print("Downloading config...")
+        config_path = hf_hub_download(hf_model_id, "config.json")
+        with open(config_path, "r") as f:
+            transformers_config = json.load(f)
+
+        # Convert config
+        keras_config = convert_backbone_config(transformers_config)
+
+        # Create KerasHub model
+        print("\nCreating KerasHub Gemma3Backbone...")
+        # Use float32 to match HF model default often
+        keras_model = keras_hub.models.Gemma3Backbone(
+            **keras_config, dtype="float32"
+        )
 
-    # Build model by calling it once
-    dummy_input = {
-        "token_ids": np.ones((1, 16), dtype="int32"),
-        "padding_mask": np.ones((1, 16), dtype="int32"),
-    }
-    _ = keras_model(dummy_input)
+        # Build model by calling it once
+        dummy_input = {
+            "token_ids": np.ones((1, 16), dtype="int32"),
+            "padding_mask": np.ones((1, 16), dtype="int32"),
+        }
+        _ = keras_model(dummy_input)
 
-    print(f"KerasHub model parameters: {keras_model.count_params():,}")
+        print(f"KerasHub model parameters: {keras_model.count_params():,}")
 
-    # Transfer weights
-    print("\nAttaching simple loader...")
-    print(f"Downloading model.safetensors into {preset}...")
-    os.makedirs(preset, exist_ok=True)
+        # Transfer weights
+        print("\nAttaching simple loader...")
+        print("Downloading model.safetensors into temporary directory...")
 
-    hf_hub_download(hf_model_id, "model.safetensors", local_dir=preset)
+        hf_hub_download(hf_model_id, "model.safetensors", local_dir=temp_dir)
 
-    # We also need to download the dense layer weights for our new logic
-    print("Downloading dense layer weights...")
-    hf_hub_download(hf_model_id, "2_Dense/model.safetensors", local_dir=preset)
-    hf_hub_download(hf_model_id, "3_Dense/model.safetensors", local_dir=preset)
+        # We also need to download the dense layer weights for our new logic
+        print("Downloading dense layer weights...")
+        hf_hub_download(
+            hf_model_id, "2_Dense/model.safetensors", local_dir=temp_dir
+        )
+        hf_hub_download(
+            hf_model_id, "3_Dense/model.safetensors", local_dir=temp_dir
+        )
 
-    print("Converting weights...")
+        print("Converting weights...")
 
-    with SafetensorLoader(preset) as loader:
-        convert_weights(keras_model, loader, transformers_config)
+        with SafetensorLoader(temp_dir) as loader:
+            convert_weights(keras_model, loader, transformers_config)
 
-    # Convert tokenizer
-    print("\nConverting tokenizer...")
-    # Helper to clean up tokenizer download if needed
-    hf_hub_download(hf_model_id, "tokenizer.model", local_dir=preset)
-    keras_tokenizer = convert_tokenizer(
-        keras_hub.models.Gemma3Tokenizer, preset
-    )
+        # Convert tokenizer
+        print("\nConverting tokenizer...")
+        # Helper to clean up tokenizer download if needed
+        hf_hub_download(hf_model_id, "tokenizer.model", local_dir=temp_dir)
+        keras_tokenizer = convert_tokenizer(
+            keras_hub.models.Gemma3Tokenizer, temp_dir
+        )
 
-    # Validate
-    passed = validate_output(keras_model, keras_tokenizer, hf_model_id)
+        # Validate
+        passed = validate_output(keras_model, keras_tokenizer, hf_model_id)
 
-    if not passed:
-        print("\n⚠️  Verification failed. Check weight mapping.")
-        return
+        if not passed:
+            print("\n⚠️  Verification failed. Check weight mapping.")
+            return
 
-    # Save preset
-    print(f"\nSaving to preset: ./{preset}")
-    keras_model.save_to_preset(preset)
-    keras_tokenizer.save_to_preset(preset)
+        # Save preset
+        print(f"\nSaving to preset: ./{preset}")
+        keras_model.save_to_preset(preset)
+        keras_tokenizer.save_to_preset(preset)
 
-    print(f"\n✅ Successfully converted and saved to: ./{preset}")
+        print(f"\n✅ Successfully converted and saved to: ./{preset}")
 
 
 if __name__ == "__main__":