keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev20240915160609__py3-none-any.whl

keras_hub/src/utils/transformers/convert_mistral.py

@@ -0,0 +1,129 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+
+from keras_hub.src.models.mistral.mistral_backbone import MistralBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = MistralBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_query_heads": transformers_config["num_attention_heads"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "num_key_value_heads": transformers_config["num_key_value_heads"],
+        "rope_max_wavelength": transformers_config["rope_theta"],
+        "layer_norm_epsilon": transformers_config["rms_norm_eps"],
+        "sliding_window": transformers_config["sliding_window"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embeddings
+    loader.port_weight(
+        keras_variable=backbone.token_embedding.embeddings,
+        hf_weight_key="model.embed_tokens.weight",
+        hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float16),
+    )
+    loader.port_weight(
+        keras_variable=backbone.token_embedding.reverse_embeddings,
+        hf_weight_key="lm_head.weight",
+        hook_fn=lambda hf_tensor, _: np.transpose(
+            hf_tensor.astype(np.float16), axes=(1, 0)
+        ),
+    )
+
+    # Attention blocks
+    for index in range(backbone.num_layers):
+        decoder_layer = backbone.transformer_layers[index]
+
+        # Norm layers
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layernorm.scale,
+            hf_weight_key=f"model.layers.{index}.input_layernorm.weight",
+            hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float16),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layernorm.scale,
+            hf_weight_key=f"model.layers.{index}.post_attention_layernorm.weight",
+            hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float16),
+        )
+
+        # Attention layers
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._query_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.self_attn.q_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor.astype(np.float16)), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._key_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.self_attn.k_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor.astype(np.float16)), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._value_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.self_attn.v_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor.astype(np.float16)), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer._output_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.self_attn.o_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor.astype(np.float16)), keras_shape
+            ),
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_gate_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.mlp.gate_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(
+                hf_tensor.astype(np.float16), axes=(1, 0)
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.mlp.up_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(
+                hf_tensor.astype(np.float16), axes=(1, 0)
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_output_dense.kernel,
+            hf_weight_key=f"model.layers.{index}.mlp.down_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(
+                hf_tensor.astype(np.float16), axes=(1, 0)
+            ),
+        )
+
+    # Normalization
+    loader.port_weight(
+        keras_variable=backbone.layer_norm.scale,
+        hf_weight_key="model.norm.weight",
+        hook_fn=lambda hf_tensor, _: hf_tensor.astype(np.float16),
+    )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(get_file(preset, "tokenizer.model"), **kwargs)

keras_hub/src/utils/transformers/convert_pali_gemma.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.models.pali_gemma.pali_gemma_backbone import (
+    PaliGemmaBackbone,
+)
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+
+backbone_cls = PaliGemmaBackbone
 
 
 def convert_backbone_config(transformers_config):
@@ -275,29 +276,6 @@ def convert_weights(backbone, loader, transformers_config):
         hook_fn=lambda hf_tensor, keras_shape: hf_tensor[: keras_shape[0]],
     )
 
-    return backbone
-
-
-def load_pali_gemma_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader, transformers_config)
-    return backbone
-
-
-def load_pali_gemma_tokenizer(cls, preset):
-    """
-    Load the Gemma tokenizer.
-
-    Args:
-        cls (class): Tokenizer class.
-        preset (str): Preset configuration name.
 
-    Returns:
-        tokenizer: Initialized tokenizer.
-    """
-    return cls(get_file(preset, "tokenizer.model"))
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(get_file(preset, "tokenizer.model"), **kwargs)

keras_hub/src/utils/transformers/preset_loader.py

@@ -0,0 +1,77 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Convert huggingface models to KerasHub."""
+
+
+from keras_hub.src.utils.preset_utils import PresetLoader
+from keras_hub.src.utils.preset_utils import jax_memory_cleanup
+from keras_hub.src.utils.transformers import convert_albert
+from keras_hub.src.utils.transformers import convert_bart
+from keras_hub.src.utils.transformers import convert_bert
+from keras_hub.src.utils.transformers import convert_distilbert
+from keras_hub.src.utils.transformers import convert_gemma
+from keras_hub.src.utils.transformers import convert_gpt2
+from keras_hub.src.utils.transformers import convert_llama3
+from keras_hub.src.utils.transformers import convert_mistral
+from keras_hub.src.utils.transformers import convert_pali_gemma
+from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+
+
+class TransformersPresetLoader(PresetLoader):
+    def __init__(self, preset, config):
+        super().__init__(preset, config)
+        model_type = self.config["model_type"]
+        if model_type == "albert":
+            self.converter = convert_albert
+        elif model_type == "bart":
+            self.converter = convert_bart
+        elif model_type == "bert":
+            self.converter = convert_bert
+        elif model_type == "distilbert":
+            self.converter = convert_distilbert
+        elif model_type == "gemma" or model_type == "gemma2":
+            self.converter = convert_gemma
+        elif model_type == "gpt2":
+            self.converter = convert_gpt2
+        elif model_type == "llama":
+            # TODO: handle other llama versions.
+            self.converter = convert_llama3
+        elif model_type == "mistral":
+            self.converter = convert_mistral
+        elif model_type == "paligemma":
+            self.converter = convert_pali_gemma
+        else:
+            raise ValueError(
+                "KerasHub has no converter for huggingface/transformers models "
+                f"with model type `'{model_type}'`."
+            )
+
+    def check_backbone_class(self):
+        return self.converter.backbone_cls
+
+    def load_backbone(self, cls, load_weights, **kwargs):
+        keras_config = self.converter.convert_backbone_config(self.config)
+        backbone = cls(**{**keras_config, **kwargs})
+        if load_weights:
+            jax_memory_cleanup(backbone)
+            with SafetensorLoader(self.preset) as loader:
+                self.converter.convert_weights(backbone, loader, self.config)
+        return backbone
+
+    def load_tokenizer(self, cls, **kwargs):
+        return self.converter.convert_tokenizer(cls, self.preset, **kwargs)
+
+    def load_image_converter(self, cls, **kwargs):
+        # TODO: set image size for pali gemma checkpoints.
+        return None

keras_hub/src/utils/transformers/safetensor_utils.py

@@ -17,7 +17,7 @@ from keras_hub.src.utils.preset_utils import SAFETENSOR_CONFIG_FILE
 from keras_hub.src.utils.preset_utils import SAFETENSOR_FILE
 from keras_hub.src.utils.preset_utils import check_file_exists
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import load_config
+from keras_hub.src.utils.preset_utils import load_json
 
 try:
     import safetensors
@@ -38,7 +38,7 @@ class SafetensorLoader(contextlib.ExitStack):
 
         self.preset = preset
         if check_file_exists(preset, SAFETENSOR_CONFIG_FILE):
-            self.safetensor_config = load_config(preset, SAFETENSOR_CONFIG_FILE)
+            self.safetensor_config = load_json(preset, SAFETENSOR_CONFIG_FILE)
         else:
             self.safetensor_config = None
         self.safetensor_files = {}

keras_hub/src/version_utils.py

@@ -15,7 +15,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.15.0.dev20240823171555"
+__version__ = "0.16.0.dev20240915160609"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.15.0.dev20240823171555.dist-info → keras_hub_nightly-0.16.0.dev20240915160609.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: keras-hub-nightly
-Version: 0.15.0.dev20240823171555
+Version: 0.16.0.dev20240915160609
 Summary: 🚧🚧🚧 Work in progress. 🚧🚧🚧 More details soon!
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team
@@ -8,7 +8,6 @@ Author-email: keras-hub@google.com
 License: Apache License 2.0
 Classifier: Development Status :: 3 - Alpha
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11