Covert a safetensor checkpoint from Hugging Face hub

keras_nlp/src/utils/preset_utils.py

@@ -59,6 +59,7 @@
 
 # Config file names.
 CONFIG_FILE = "config.json"
+HF_CONFIG_FILE = "config.json"
 TOKENIZER_CONFIG_FILE = "tokenizer.json"
 TASK_CONFIG_FILE = "task.json"
 PREPROCESSOR_CONFIG_FILE = "preprocessor.json"

keras_nlp/src/utils/transformers/convert.py

@@ -29,7 +29,11 @@ def load_transformers_backbone(cls, preset, load_weights):
         return load_gemma_backbone(cls, preset, load_weights)
     if cls.__name__ == "Llama3Backbone":
         return load_llama3_backbone(cls, preset, load_weights)
-    raise ValueError(f"No conversion huggingface/transformers to {cls}")
+    raise ValueError(
+        f"{cls} has not been ported from the Hugging Face format yet. "
+        "Please check Hugging Face Hub for the Keras model. "
+        "Models in Keras format should end with `-keras`. (e.g google/gemma-2b-keras)"
+    )
 
 
 def load_transformers_tokenizer(cls, preset):
@@ -39,4 +43,8 @@ def load_transformers_tokenizer(cls, preset):
         return load_gemma_tokenizer(cls, preset)
     if cls.__name__ == "Llama3Tokenizer":
         return load_llama3_tokenizer(cls, preset)
-    raise ValueError(f"No conversion huggingface/transformers to {cls}")
+    ValueError(
+        f"{cls} has not been ported from the Hugging Face format yet. "
+        "Please check Hugging Face Hub for the Keras model. "
+        "Models in Keras format should end with `-keras`. (e.g google/gemma-2b-keras)"
+    )

keras_nlp/src/utils/transformers/convert_gemma.py

@@ -15,6 +15,8 @@
 
 import numpy as np
 
+from keras_nlp.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_nlp.src.utils.preset_utils import SAFETENSOR_CONFIG_FILE
 from keras_nlp.src.utils.preset_utils import get_file
 from keras_nlp.src.utils.preset_utils import jax_memory_cleanup
 from keras_nlp.src.utils.preset_utils import load_config
@@ -33,7 +35,7 @@ def load_gemma_backbone(cls, preset, load_weights):
     Returns:
         backbone: Initialized Keras model backbone.
     """
-    transformers_config = load_config(preset, "config.json")
+    transformers_config = load_config(preset, HF_CONFIG_FILE)
 
     backbone = cls(
         vocabulary_size=transformers_config["vocab_size"],
@@ -50,7 +52,7 @@ def load_gemma_backbone(cls, preset, load_weights):
 
     jax_memory_cleanup(backbone)
     # Code to port the weights from safetensors into the keras nlp model
-    safetensor_config = load_config(preset, "model.safetensors.index.json")
+    safetensor_config = load_config(preset, SAFETENSOR_CONFIG_FILE)
     safetensor_files = {
         fname: get_file(preset, fname)
         for fname in set(safetensor_config["weight_map"].values())

keras_nlp/src/utils/transformers/convert_llama3.py

@@ -15,6 +15,8 @@
 
 import numpy as np
 
+from keras_nlp.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_nlp.src.utils.preset_utils import SAFETENSOR_CONFIG_FILE
 from keras_nlp.src.utils.preset_utils import get_file
 from keras_nlp.src.utils.preset_utils import jax_memory_cleanup
 from keras_nlp.src.utils.preset_utils import load_config
@@ -33,7 +35,7 @@ def load_llama3_backbone(cls, preset, load_weights):
     Returns:
         backbone: Initialized Keras model backbone.
     """
-    transformers_config = load_config(preset, "config.json")
+    transformers_config = load_config(preset, HF_CONFIG_FILE)
 
     backbone = cls(
         vocabulary_size=transformers_config["vocab_size"],
@@ -49,7 +51,7 @@ def load_llama3_backbone(cls, preset, load_weights):
 
     jax_memory_cleanup(backbone)
     # Code to port the weights from safetensors into the keras nlp model
-    safetensor_config = load_config(preset, "model.safetensors.index.json")
+    safetensor_config = load_config(preset, SAFETENSOR_CONFIG_FILE)
     safetensor_files = {
         fname: get_file(preset, fname)
         for fname in set(safetensor_config["weight_map"].values())