keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl

keras_hub/src/utils/transformers/convert_bert.py

@@ -0,0 +1,173 @@
+# 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.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.utils.preset_utils import HF_TOKENIZER_CONFIG_FILE
+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
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embedding layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="bert.embeddings.word_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer(
+            "position_embedding"
+        ).position_embeddings,
+        hf_weight_key="bert.embeddings.position_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer("segment_embedding").embeddings,
+        hf_weight_key="bert.embeddings.token_type_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer("embeddings_layer_norm").beta,
+        hf_weight_key="bert.embeddings.LayerNorm.beta",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer("embeddings_layer_norm").gamma,
+        hf_weight_key="bert.embeddings.LayerNorm.gamma",
+    )
+
+    def transpose_and_reshape(x, shape):
+        return np.reshape(np.transpose(x), shape)
+
+    # Attention blocks
+    for i in range(backbone.num_layers):
+        block = backbone.get_layer(f"transformer_layer_{i}")
+        attn = block._self_attention_layer
+        hf_prefix = "bert.encoder.layer."
+        # Attention layers
+        loader.port_weight(
+            keras_variable=attn.query_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.query_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.key_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.key_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.value_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.value_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.output_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        # Attention layer norm.
+        loader.port_weight(
+            keras_variable=block._self_attention_layer_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.LayerNorm.beta",
+        )
+        loader.port_weight(
+            keras_variable=block._self_attention_layer_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.LayerNorm.gamma",
+        )
+        # MLP layers
+        loader.port_weight(
+            keras_variable=block._feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block._feedforward_intermediate_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.bias",
+        )
+        loader.port_weight(
+            keras_variable=block._feedforward_output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block._feedforward_output_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.bias",
+        )
+        # Output layer norm.
+        loader.port_weight(
+            keras_variable=block._feedforward_layer_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.output.LayerNorm.beta",
+        )
+        loader.port_weight(
+            keras_variable=block._feedforward_layer_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.output.LayerNorm.gamma",
+        )
+
+    loader.port_weight(
+        keras_variable=backbone.get_layer("pooled_dense").kernel,
+        hf_weight_key="bert.pooler.dense.weight",
+        hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer("pooled_dense").bias,
+        hf_weight_key="bert.pooler.dense.bias",
+    )
+
+
+def load_bert_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_bert_tokenizer(cls, preset):
+    transformers_config = load_config(preset, HF_TOKENIZER_CONFIG_FILE)
+    return cls(
+        get_file(preset, "vocab.txt"),
+        lowercase=transformers_config["do_lower_case"],
+    )

keras_hub/src/utils/transformers/convert_distilbert.py

@@ -0,0 +1,184 @@
+# 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.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.utils.preset_utils import HF_TOKENIZER_CONFIG_FILE
+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
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["n_layers"],
+        "num_heads": transformers_config["n_heads"],
+        "hidden_dim": transformers_config["dim"],
+        "intermediate_dim": transformers_config["hidden_dim"],
+        "dropout": transformers_config["dropout"],
+        "max_sequence_length": transformers_config["max_position_embeddings"],
+    }
+
+
+def convert_weights(backbone, loader):
+    # Embeddings
+    loader.port_weight(
+        keras_variable=backbone.get_layer(
+            "token_and_position_embedding"
+        ).token_embedding.embeddings,
+        hf_weight_key="distilbert.embeddings.word_embeddings.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.get_layer(
+            "token_and_position_embedding"
+        ).position_embedding.position_embeddings,
+        hf_weight_key="distilbert.embeddings.position_embeddings.weight",
+    )
+
+    # 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_layer_norm.gamma,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.sa_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer_norm.beta,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.sa_layer_norm.bias",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layer_norm.gamma,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.output_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layer_norm.beta,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.output_layer_norm.bias",
+        )
+
+        # Attention layers
+        # Query
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.query_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.q_lin.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.query_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.q_lin.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                hf_tensor, keras_shape
+            ),
+        )
+
+        # Key
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.key_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.k_lin.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.key_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.k_lin.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                hf_tensor, keras_shape
+            ),
+        )
+
+        # Value
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.value_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.v_lin.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.value_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.v_lin.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                hf_tensor, keras_shape
+            ),
+        )
+
+        # Output
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.output_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.out_lin.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer.output_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.attention.out_lin.bias",
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.ffn.lin1.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_intermediate_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.ffn.lin1.bias",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_output_dense.kernel,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.ffn.lin2.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_output_dense.bias,
+            hf_weight_key=f"distilbert.transformer.layer.{index}.ffn.lin2.bias",
+        )
+
+    # Normalization
+    loader.port_weight(
+        keras_variable=backbone.embeddings_layer_norm.gamma,
+        hf_weight_key="distilbert.embeddings.LayerNorm.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings_layer_norm.beta,
+        hf_weight_key="distilbert.embeddings.LayerNorm.bias",
+    )
+
+    return backbone
+
+
+def load_distilbert_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)
+    return backbone
+
+
+def load_distilbert_tokenizer(cls, preset):
+    transformers_config = load_config(preset, HF_TOKENIZER_CONFIG_FILE)
+    return cls(
+        get_file(preset, "vocab.txt"),
+        lowercase=transformers_config["do_lower_case"],
+    )

keras_hub/src/utils/transformers/convert_gemma.py

@@ -0,0 +1,187 @@
+# 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.utils.preset_utils import HF_CONFIG_FILE
+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
+
+
+def convert_backbone_config(transformers_config):
+    backbone_config = dict()
+    if transformers_config["model_type"] == "gemma":
+        # Build Gemma backbone configuration
+        backbone_config = {
+            "vocabulary_size": transformers_config["vocab_size"],
+            "num_layers": transformers_config["num_hidden_layers"],
+            "num_query_heads": transformers_config["num_attention_heads"],
+            "num_key_value_heads": transformers_config["num_key_value_heads"],
+            "hidden_dim": transformers_config["hidden_size"],
+            "intermediate_dim": transformers_config["intermediate_size"] * 2,
+            "head_dim": transformers_config["head_dim"],
+        }
+    elif transformers_config["model_type"] == "gemma2":
+        # Build Gemma 2 backbone configuration
+        backbone_config = {
+            "vocabulary_size": transformers_config["vocab_size"],
+            "num_layers": transformers_config["num_hidden_layers"],
+            "num_query_heads": transformers_config["num_attention_heads"],
+            "num_key_value_heads": transformers_config["num_key_value_heads"],
+            "hidden_dim": transformers_config["hidden_size"],
+            "intermediate_dim": transformers_config["intermediate_size"] * 2,
+            "head_dim": transformers_config["head_dim"],
+            "query_head_dim_normalize": (
+                transformers_config["head_dim"]
+                == transformers_config["query_pre_attn_scalar"]
+            ),
+            "use_post_ffw_norm": True,
+            "use_post_attention_norm": True,
+            "final_logit_soft_cap": transformers_config[
+                "final_logit_softcapping"
+            ],
+            "attention_logit_soft_cap": transformers_config[
+                "attn_logit_softcapping"
+            ],
+            "sliding_window_size": transformers_config["sliding_window"],
+            "use_sliding_window_attention": True,
+        }
+    return backbone_config
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embedding layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="model.embed_tokens.weight",
+    )
+
+    # Attention blocks
+    for i in range(backbone.num_layers):
+        decoder_layer = backbone.get_layer(f"decoder_block_{i}")
+        # Norm layers
+        loader.port_weight(
+            keras_variable=decoder_layer.pre_attention_norm.scale,
+            hf_weight_key=f"model.layers.{i}.input_layernorm.weight",
+        )
+
+        if decoder_layer.use_post_attention_norm:
+            loader.port_weight(
+                keras_variable=decoder_layer.post_attention_norm.scale,
+                hf_weight_key=f"model.layers.{i}.post_attention_layernorm.weight",
+            )
+
+        if transformers_config["model_type"] == "gemma":
+            loader.port_weight(
+                keras_variable=decoder_layer.pre_ffw_norm.scale,
+                hf_weight_key=f"model.layers.{i}.post_attention_layernorm.weight",
+            )
+        elif transformers_config["model_type"] == "gemma2":
+            loader.port_weight(
+                keras_variable=decoder_layer.pre_ffw_norm.scale,
+                hf_weight_key=f"model.layers.{i}.pre_feedforward_layernorm.weight",
+            )
+
+        if decoder_layer.use_post_ffw_norm:
+            loader.port_weight(
+                keras_variable=decoder_layer.post_ffw_norm.scale,
+                hf_weight_key=f"model.layers.{i}.post_feedforward_layernorm.weight",
+            )
+
+        # Attention layers
+        loader.port_weight(
+            keras_variable=decoder_layer.attention.query_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.q_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.transpose(
+                np.reshape(
+                    hf_tensor,
+                    (keras_shape[0], keras_shape[2], keras_shape[1]),
+                ),
+                axes=(0, 2, 1),
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer.attention.key_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.k_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.transpose(
+                np.reshape(
+                    hf_tensor,
+                    (keras_shape[0], keras_shape[2], keras_shape[1]),
+                ),
+                axes=(0, 2, 1),
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer.attention.value_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.v_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.transpose(
+                np.reshape(
+                    hf_tensor,
+                    (keras_shape[0], keras_shape[2], keras_shape[1]),
+                ),
+                axes=(0, 2, 1),
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer.attention.output_dense.kernel,
+            hf_weight_key=f"model.layers.{i}.self_attn.o_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.transpose(
+                np.reshape(
+                    hf_tensor,
+                    (keras_shape[2], keras_shape[0], keras_shape[1]),
+                ),
+                axes=(1, 2, 0),
+            ),
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=decoder_layer.gating_ffw.kernel,
+            hf_weight_key=f"model.layers.{i}.mlp.gate_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer.gating_ffw_2.kernel,
+            hf_weight_key=f"model.layers.{i}.mlp.up_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer.ffw_linear.kernel,
+            hf_weight_key=f"model.layers.{i}.mlp.down_proj.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+
+    # Final normalization layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("final_normalization").scale,
+        hf_weight_key="model.norm.weight",
+    )
+
+    return backbone
+
+
+def load_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_gemma_tokenizer(cls, preset):
+    return cls(get_file(preset, "tokenizer.model"))