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

keras_hub/src/utils/tensor_utils.py

@@ -0,0 +1,206 @@
+# 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 keras
+from keras import ops
+
+try:
+    import tensorflow as tf
+    import tensorflow_text as tf_text
+except ImportError:
+    tf = None
+    tf_text = None
+
+
+def _decode_strings_to_utf8(inputs):
+    """Recursively decodes to list of strings with 'utf-8' encoding."""
+    if isinstance(inputs, bytes):
+        # Handles the case when the input is a scalar string.
+        return inputs.decode("utf-8", errors="ignore")
+    else:
+        # Recursively iterate when input is a list.
+        return [_decode_strings_to_utf8(x) for x in inputs]
+
+
+def tensor_to_list(inputs):
+    """Converts a tensor to nested lists.
+
+    Args:
+        inputs: Input tensor, or dict/list/tuple of input tensors.
+    """
+    if not isinstance(inputs, (tf.RaggedTensor, tf.Tensor)):
+        inputs = tf.convert_to_tensor(inputs)
+    if isinstance(inputs, tf.RaggedTensor):
+        list_outputs = inputs.to_list()
+    elif isinstance(inputs, tf.Tensor):
+        list_outputs = inputs.numpy()
+        if inputs.shape.rank != 0:
+            list_outputs = list_outputs.tolist()
+    if inputs.dtype == tf.string:
+        list_outputs = _decode_strings_to_utf8(list_outputs)
+    return list_outputs
+
+
+def convert_to_backend_tensor_or_python_list(x):
+    """
+    Convert a tensor to the backend friendly representation of the data.
+
+    This wraps `ops.convert_to_tensor` to account for the fact that torch and
+    jax both lack native types for ragged and string data.
+
+    If we encounter one of these types in torch or jax, we will instead covert
+    the tensor to simple pythonic types (lists of strings).
+    """
+    if isinstance(x, tf.RaggedTensor) or getattr(x, "dtype", None) == tf.string:
+        return tensor_to_list(x)
+    dtype = getattr(x, "dtype", "float32")
+    dtype = keras.backend.standardize_dtype(dtype)
+    return ops.convert_to_tensor(x, dtype=dtype)
+
+
+def convert_to_ragged_batch(inputs):
+    """Convert pythonic or numpy-like input to a 2-D `tf.RaggedTensor`.
+
+    This is useful for text preprocessing layers which deal with already
+    tokenized or split text.
+
+    Args:
+        inputs: A pythonic or numpy-like input to covert. This input should
+            represent a possibly batched list of token sequences.
+
+    Returns:
+        An `(inputs, unbatched, rectangular)` tuple, where `inputs` is a
+        2-D `tf.RaggedTensor`, `unbatched` is `True` if the inputs were
+        origianlly rank 1, and `rectangular` is `True` if the inputs rows are
+        all of equal lengths.
+    """
+    # `tf.keras.layers.Layer` does a weird conversion in __call__, where a list
+    # of lists of ints will become a list of list of scalar tensors. We could
+    # clean this up if we no longer need to care about that case.
+    if isinstance(inputs, (list, tuple)):
+        if isinstance(inputs[0], (list, tuple)):
+            rectangular = len(set([len(row) for row in inputs])) == 1
+            rows = [
+                tf.convert_to_tensor(row, dtype_hint="int32") for row in inputs
+            ]
+            inputs = tf.ragged.stack(rows).with_row_splits_dtype("int64")
+        else:
+            inputs = tf.convert_to_tensor(inputs)
+            rectangular = True
+    elif isinstance(inputs, tf.Tensor):
+        rectangular = True
+    elif isinstance(inputs, tf.RaggedTensor):
+        rectangular = False
+    elif hasattr(inputs, "__array__"):
+        inputs = tf.convert_to_tensor(ops.convert_to_numpy(inputs))
+        rectangular = True
+    else:
+        raise ValueError(
+            f"Unknown tensor type. Tensor input can be passed as "
+            "tensors, numpy arrays, or python lists. Received: "
+            f"`type(inputs)={type(inputs)}`"
+        )
+    if inputs.shape.rank < 1 or inputs.shape.rank > 2:
+        raise ValueError(
+            f"Tokenized tensor input should be rank 1 (unbatched) or "
+            f"rank 2 (batched). Received: `inputs.shape={input.shape}`"
+        )
+    unbatched = inputs.shape.rank == 1
+    rectangular = rectangular or unbatched
+    if unbatched:
+        inputs = tf.expand_dims(inputs, 0)
+    if isinstance(inputs, tf.Tensor):
+        inputs = tf.RaggedTensor.from_tensor(inputs)
+    return inputs, unbatched, rectangular
+
+
+def truncate_at_token(inputs, token, mask):
+    """Truncate at first instance of `token`, ignoring `mask`."""
+    matches = (inputs == token) & (~mask)
+    end_indices = tf.cast(tf.math.argmax(matches, -1), "int32")
+    end_indices = tf.where(end_indices == 0, tf.shape(inputs)[-1], end_indices)
+    return tf.RaggedTensor.from_tensor(inputs, end_indices)
+
+
+def strip_to_ragged(token_ids, mask, ids_to_strip):
+    """Remove masked and special tokens from a sequence before detokenizing."""
+    token_ids = ops.convert_to_numpy(token_ids)
+    token_ids = token_ids.astype("int32")
+    mask = ops.convert_to_numpy(mask)
+    mask = mask.astype("bool")
+    for id in ids_to_strip:
+        mask = mask & (token_ids != id)
+    return tf.ragged.boolean_mask(token_ids, mask)
+
+
+def assert_tf_libs_installed(symbol_name):
+    if tf_text is None or tf is None:
+        raise ImportError(
+            f"{symbol_name} requires `tensorflow` and `tensorflow-text` for "
+            "text processing. Run `pip install tensorflow-text` to install "
+            "both packages or visit https://www.tensorflow.org/install\n\n"
+            "If `tensorflow-text` is already installed, try importing it "
+            "in a clean python session. Your installation may have errors.\n\n"
+            "KerasHub uses `tf.data` and `tensorflow-text` to preprocess text "
+            "on all Keras backends. If you are running on Jax or Torch, this "
+            "installation does not need GPU support."
+        )
+
+
+def assert_tf_backend(symbol_name):
+    if keras.config.backend() != "tensorflow":
+        raise RuntimeError(
+            f"{symbol_name} requires the `tensorflow` backend. "
+            "Please set `KERAS_BACKEND=tensorflow` when running your program."
+        )
+
+
+def is_tensor_type(x):
+    return hasattr(x, "__array__")
+
+
+def is_float_dtype(dtype):
+    return "float" in keras.backend.standardize_dtype(dtype)
+
+
+def is_int_dtype(dtype):
+    return "int" in keras.backend.standardize_dtype(dtype)
+
+
+def is_string_dtype(dtype):
+    return "string" in keras.backend.standardize_dtype(dtype)
+
+
+def any_equal(inputs, values, padding_mask):
+    """Return a mask that is True anywhere `inputs` has a value in `values`.
+
+    Final mask has `padding_mask` applied.
+
+    Args:
+        inputs: Input tensor.
+        values: List or iterable of tensors shaped like `inputs` or broadcastable
+            by bit operators.
+        padding_mask: Tensor with shape compatible with inputs that will condition
+            output.
+
+    Returns:
+        A tensor with `inputs` shape where each position is True if it contains
+            a value from any `values`. Padding mask will be applied before
+            returning."""
+    output = ops.equal(inputs, values[0])
+    for value in values[1:]:
+        value_equality = ops.equal(inputs, value)
+        output = ops.logical_or(output, value_equality)
+
+    return ops.logical_and(output, padding_mask)

keras_hub/src/utils/timm/__init__.py

@@ -0,0 +1,13 @@
+# 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.

keras_hub/src/utils/timm/convert.py

@@ -0,0 +1,37 @@
+# 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 timm models to KerasHub."""
+
+from keras_hub.src.utils.timm.convert_resnet import load_resnet_backbone
+
+
+def load_timm_backbone(cls, preset, load_weights, **kwargs):
+    """Load a timm model config and weights as a KerasHub backbone.
+
+    Args:
+        cls (class): Keras model class.
+        preset (str): Preset configuration name.
+        load_weights (bool): Whether to load the weights.
+
+    Returns:
+        backbone: Initialized Keras model backbone.
+    """
+    if cls is None:
+        raise ValueError("Backbone class is None")
+    if cls.__name__ == "ResNetBackbone":
+        return load_resnet_backbone(cls, preset, load_weights, **kwargs)
+    raise ValueError(
+        f"{cls} has not been ported from the Hugging Face format yet. "
+        "Please check Hugging Face Hub for the Keras model. "
+    )

keras_hub/src/utils/timm/convert_resnet.py

@@ -0,0 +1,171 @@
+# 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 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(timm_config):
+    timm_architecture = timm_config["architecture"]
+
+    if "resnetv2_" in timm_architecture:
+        use_pre_activation = True
+    else:
+        use_pre_activation = False
+
+    if timm_architecture == "resnet18":
+        stackwise_num_blocks = [2, 2, 2, 2]
+        block_type = "basic_block"
+    elif timm_architecture == "resnet26":
+        stackwise_num_blocks = [2, 2, 2, 2]
+        block_type = "bottleneck_block"
+    elif timm_architecture == "resnet34":
+        stackwise_num_blocks = [3, 4, 6, 3]
+        block_type = "basic_block"
+    elif timm_architecture in ("resnet50", "resnetv2_50"):
+        stackwise_num_blocks = [3, 4, 6, 3]
+        block_type = "bottleneck_block"
+    elif timm_architecture in ("resnet101", "resnetv2_101"):
+        stackwise_num_blocks = [3, 4, 23, 3]
+        block_type = "bottleneck_block"
+    elif timm_architecture in ("resnet152", "resnetv2_152"):
+        stackwise_num_blocks = [3, 8, 36, 3]
+        block_type = "bottleneck_block"
+    else:
+        raise ValueError(
+            f"Currently, the architecture {timm_architecture} is not supported."
+        )
+
+    return dict(
+        stackwise_num_filters=[64, 128, 256, 512],
+        stackwise_num_blocks=stackwise_num_blocks,
+        stackwise_num_strides=[1, 2, 2, 2],
+        block_type=block_type,
+        use_pre_activation=use_pre_activation,
+    )
+
+
+def convert_weights(backbone, loader, timm_config):
+    def port_conv2d(keras_layer_name, hf_weight_prefix):
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).kernel,
+            hf_weight_key=f"{hf_weight_prefix}.weight",
+            hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+        )
+
+    def port_batch_normalization(keras_layer_name, hf_weight_prefix):
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).gamma,
+            hf_weight_key=f"{hf_weight_prefix}.weight",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).beta,
+            hf_weight_key=f"{hf_weight_prefix}.bias",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).moving_mean,
+            hf_weight_key=f"{hf_weight_prefix}.running_mean",
+        )
+        loader.port_weight(
+            backbone.get_layer(keras_layer_name).moving_variance,
+            hf_weight_key=f"{hf_weight_prefix}.running_var",
+        )
+
+    version = "v1" if not backbone.use_pre_activation else "v2"
+    block_type = backbone.block_type
+
+    # Stem
+    if version == "v1":
+        port_conv2d("conv1_conv", "conv1")
+        port_batch_normalization("conv1_bn", "bn1")
+    else:
+        port_conv2d("conv1_conv", "stem.conv")
+
+    # Stages
+    num_stacks = len(backbone.stackwise_num_filters)
+    for stack_index in range(num_stacks):
+        for block_idx in range(backbone.stackwise_num_blocks[stack_index]):
+            if version == "v1":
+                keras_name = f"v1_stack{stack_index}_block{block_idx}"
+                hf_name = f"layer{stack_index+1}.{block_idx}"
+            else:
+                keras_name = f"v2_stack{stack_index}_block{block_idx}"
+                hf_name = f"stages.{stack_index}.blocks.{block_idx}"
+
+            if version == "v1":
+                if block_idx == 0 and (
+                    block_type == "bottleneck_block" or stack_index > 0
+                ):
+                    port_conv2d(
+                        f"{keras_name}_0_conv", f"{hf_name}.downsample.0"
+                    )
+                    port_batch_normalization(
+                        f"{keras_name}_0_bn", f"{hf_name}.downsample.1"
+                    )
+                port_conv2d(f"{keras_name}_1_conv", f"{hf_name}.conv1")
+                port_batch_normalization(f"{keras_name}_1_bn", f"{hf_name}.bn1")
+                port_conv2d(f"{keras_name}_2_conv", f"{hf_name}.conv2")
+                port_batch_normalization(f"{keras_name}_2_bn", f"{hf_name}.bn2")
+                if block_type == "bottleneck_block":
+                    port_conv2d(f"{keras_name}_3_conv", f"{hf_name}.conv3")
+                    port_batch_normalization(
+                        f"{keras_name}_3_bn", f"{hf_name}.bn3"
+                    )
+            else:
+                if block_idx == 0 and (
+                    block_type == "bottleneck_block" or stack_index > 0
+                ):
+                    port_conv2d(
+                        f"{keras_name}_0_conv", f"{hf_name}.downsample.conv"
+                    )
+                port_batch_normalization(
+                    f"{keras_name}_pre_activation_bn", f"{hf_name}.norm1"
+                )
+                port_conv2d(f"{keras_name}_1_conv", f"{hf_name}.conv1")
+                port_batch_normalization(
+                    f"{keras_name}_1_bn", f"{hf_name}.norm2"
+                )
+                port_conv2d(f"{keras_name}_2_conv", f"{hf_name}.conv2")
+                if block_type == "bottleneck_block":
+                    port_batch_normalization(
+                        f"{keras_name}_2_bn", f"{hf_name}.norm3"
+                    )
+                    port_conv2d(f"{keras_name}_3_conv", f"{hf_name}.conv3")
+
+    # Post
+    if version == "v2":
+        port_batch_normalization("post_bn", "norm")
+
+    # Rebuild normalization layer with pretrained mean & std
+    mean = timm_config["pretrained_cfg"]["mean"]
+    std = timm_config["pretrained_cfg"]["std"]
+    normalization_layer = backbone.get_layer("normalization")
+    normalization_layer.input_mean = mean
+    normalization_layer.input_variance = [s**2 for s in std]
+    normalization_layer.build(normalization_layer._build_input_shape)
+
+
+def load_resnet_backbone(cls, preset, load_weights, **kwargs):
+    timm_config = load_config(preset, HF_CONFIG_FILE)
+    keras_config = convert_backbone_config(timm_config)
+    backbone = cls(**keras_config, **kwargs)
+    if load_weights:
+        jax_memory_cleanup(backbone)
+        # Use prefix="" to avoid using `get_prefixed_key`.
+        with SafetensorLoader(preset, prefix="") as loader:
+            convert_weights(backbone, loader, timm_config)
+    return backbone

keras_hub/src/utils/transformers/__init__.py

@@ -0,0 +1,13 @@
+# 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.

keras_hub/src/utils/transformers/convert.py

@@ -0,0 +1,101 @@
+# 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.transformers.convert_bert import load_bert_backbone
+from keras_hub.src.utils.transformers.convert_bert import load_bert_tokenizer
+from keras_hub.src.utils.transformers.convert_distilbert import (
+    load_distilbert_backbone,
+)
+from keras_hub.src.utils.transformers.convert_distilbert import (
+    load_distilbert_tokenizer,
+)
+from keras_hub.src.utils.transformers.convert_gemma import load_gemma_backbone
+from keras_hub.src.utils.transformers.convert_gemma import load_gemma_tokenizer
+from keras_hub.src.utils.transformers.convert_gpt2 import load_gpt2_backbone
+from keras_hub.src.utils.transformers.convert_gpt2 import load_gpt2_tokenizer
+from keras_hub.src.utils.transformers.convert_llama3 import load_llama3_backbone
+from keras_hub.src.utils.transformers.convert_llama3 import (
+    load_llama3_tokenizer,
+)
+from keras_hub.src.utils.transformers.convert_pali_gemma import (
+    load_pali_gemma_backbone,
+)
+from keras_hub.src.utils.transformers.convert_pali_gemma import (
+    load_pali_gemma_tokenizer,
+)
+
+
+def load_transformers_backbone(cls, preset, load_weights):
+    """
+    Load a Transformer model config and weights as a KerasHub backbone.
+
+    Args:
+        cls (class): Keras model class.
+        preset (str): Preset configuration name.
+        load_weights (bool): Whether to load the weights.
+
+    Returns:
+        backbone: Initialized Keras model backbone.
+    """
+    if cls is None:
+        raise ValueError("Backbone class is None")
+    if cls.__name__ == "BertBackbone":
+        return load_bert_backbone(cls, preset, load_weights)
+    if cls.__name__ == "GemmaBackbone":
+        return load_gemma_backbone(cls, preset, load_weights)
+    if cls.__name__ == "Llama3Backbone":
+        return load_llama3_backbone(cls, preset, load_weights)
+    if cls.__name__ == "PaliGemmaBackbone":
+        return load_pali_gemma_backbone(cls, preset, load_weights)
+    if cls.__name__ == "GPT2Backbone":
+        return load_gpt2_backbone(cls, preset, load_weights)
+    if cls.__name__ == "DistilBertBackbone":
+        return load_distilbert_backbone(cls, preset, load_weights)
+    raise ValueError(
+        f"{cls} has not been ported from the Hugging Face format yet. "
+        "Please check Hugging Face Hub for the Keras model. "
+    )
+
+
+def load_transformers_tokenizer(cls, preset):
+    """
+    Load a Transformer tokenizer assets as a KerasHub tokenizer.
+
+    Args:
+        cls (class): Tokenizer class.
+        preset (str): Preset configuration name.
+
+    Returns:
+        tokenizer: Initialized tokenizer.
+    """
+    if cls is None:
+        raise ValueError("Tokenizer class is None")
+    if cls.__name__ == "BertTokenizer":
+        return load_bert_tokenizer(cls, preset)
+    if cls.__name__ == "GemmaTokenizer":
+        return load_gemma_tokenizer(cls, preset)
+    if cls.__name__ == "Llama3Tokenizer":
+        return load_llama3_tokenizer(cls, preset)
+    if cls.__name__ == "PaliGemmaTokenizer":
+        return load_pali_gemma_tokenizer(cls, preset)
+    if cls.__name__ == "GPT2Tokenizer":
+        return load_gpt2_tokenizer(cls, preset)
+    if cls.__name__ == "DistilBertTokenizer":
+        return load_distilbert_tokenizer(cls, preset)
+    raise ValueError(
+        f"{cls} has not been ported from the Hugging Face format yet. "
+        "Please check Hugging Face Hub for the Keras model. "
+    )