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

keras_hub/src/models/backbone.py

@@ -0,0 +1,311 @@
+# 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 os
+
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.utils.keras_utils import assert_quantization_support
+from keras_hub.src.utils.preset_utils import CONFIG_FILE
+from keras_hub.src.utils.preset_utils import MODEL_WEIGHTS_FILE
+from keras_hub.src.utils.preset_utils import check_config_class
+from keras_hub.src.utils.preset_utils import check_format
+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 list_presets
+from keras_hub.src.utils.preset_utils import list_subclasses
+from keras_hub.src.utils.preset_utils import load_serialized_object
+from keras_hub.src.utils.preset_utils import save_metadata
+from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.python_utils import classproperty
+from keras_hub.src.utils.timm.convert import load_timm_backbone
+from keras_hub.src.utils.transformers.convert import load_transformers_backbone
+
+
+@keras_hub_export("keras_hub.models.Backbone")
+class Backbone(keras.Model):
+    """Base class for all `Backbone` models.
+
+    A `Backbone` is the basic architecture for a given NLP model. Unlike a
+    `keras_hub.models.Task`, a `Backbone` is not tailored to any specific loss
+    function and training setup. A `Backbone` generally outputs the last hidden
+    states of an architecture before any output predictions.
+
+    A `Backbone` can be used in one of two ways:
+
+    1. Through a `Task` class, which will wrap and extend a `Backbone` so it
+       can be used with high level Keras functions like `fit()`, `predict()` or
+       `evaluate()`. `Task` classes are built with a particular training
+       objective in mind (e.g. classification or language modeling).
+    2. Directly, by extending underlying functional model with additional
+       outputs and training setup. This is the most flexible approach, and can
+       allow for any outputs, loss, or custom training loop.
+
+    All backbones include a `from_preset()` constructor which can be used to
+    load a pre-trained config and weights.
+
+    Example:
+    ```python
+    # Load a BERT backbone with pre-trained weights.
+    backbone = keras_hub.models.Backbone.from_preset(
+        "bert_base_en",
+    )
+    # Load a GPT2 backbone with pre-trained weights at bfloat16 precision.
+    backbone = keras_hub.models.Backbone.from_preset(
+        "gpt2_base_en",
+        dtype="bfloat16",
+        trainable=False,
+    )
+    ```
+    """
+
+    def __init__(self, *args, dtype=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._functional_layer_ids = set(
+            id(layer) for layer in self._flatten_layers()
+        )
+        self._initialized = True
+        if dtype is not None:
+            try:
+                self.dtype_policy = keras.dtype_policies.get(dtype)
+            # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+            except AttributeError:
+                if isinstance(dtype, keras.DTypePolicy):
+                    dtype = dtype.name
+                self.dtype_policy = keras.DTypePolicy(dtype)
+
+    def __setattr__(self, name, value):
+        # Work around setattr issues for Keras 2 and Keras 3 torch backend.
+        # Since all our state is covered by functional model we can route
+        # around custom setattr calls.
+        is_property = isinstance(getattr(type(self), name, None), property)
+        is_unitialized = not hasattr(self, "_initialized")
+        simple_setattr = keras.config.backend() == "torch"
+        if simple_setattr and (is_property or is_unitialized):
+            return object.__setattr__(self, name, value)
+        return super().__setattr__(name, value)
+
+    @property
+    def token_embedding(self):
+        """A `keras.layers.Embedding` instance for embedding token ids.
+
+        This layer embeds integer token ids to the hidden dim of the model.
+        """
+        return getattr(self, "_token_embedding", None)
+
+    @token_embedding.setter
+    def token_embedding(self, value):
+        self._token_embedding = value
+
+    def quantize(self, mode, **kwargs):
+        assert_quantization_support()
+        return super().quantize(mode, **kwargs)
+
+    def get_config(self):
+        # Don't chain to super here. `get_config()` for functional models is
+        # a nested layer config and cannot be passed to Backbone constructors.
+        config = {
+            "name": self.name,
+            "trainable": self.trainable,
+        }
+
+        # Add quantization support by utilizing `DTypePolicyMap`
+        try:
+            if isinstance(
+                self.dtype_policy, keras.dtype_policies.DTypePolicyMap
+            ):
+                config.update({"dtype": self.dtype_policy})
+            else:
+                policy_map = keras.dtype_policies.DTypePolicyMap()
+                for layer in self._flatten_layers():
+                    if layer.quantization_mode is not None:
+                        policy_map[layer.path] = layer.dtype_policy
+                if len(policy_map) > 0:
+                    config.update({"dtype": policy_map})
+        # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+        except AttributeError:
+            pass
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        # The default `from_config()` for functional models will return a
+        # vanilla `keras.Model`. We override it to get a subclass instance back.
+        return cls(**config)
+
+    @classproperty
+    def presets(cls):
+        """List built-in presets for a `Task` subclass."""
+        presets = list_presets(cls)
+        for subclass in list_subclasses(cls):
+            presets.update(subclass.presets)
+        return presets
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        """Instantiate a `keras_hub.models.Backbone` from a model preset.
+
+        A preset is a directory of configs, weights and other file assets used
+        to save and load a pre-trained model. The `preset` can be passed as a
+        one of:
+
+        1. a built in preset identifier like `'bert_base_en'`
+        2. a Kaggle Models handle like `'kaggle://user/bert/keras/bert_base_en'`
+        3. a Hugging Face handle like `'hf://user/bert_base_en'`
+        4. a path to a local preset directory like `'./bert_base_en'`
+
+        This constructor can be called in one of two ways. Either from the base
+        class like `keras_hub.models.Backbone.from_preset()`, or from
+        a model class like `keras_hub.models.GemmaBackbone.from_preset()`.
+        If calling from the base class, the subclass of the returning object
+        will be inferred from the config in the preset directory.
+
+        For any `Backbone` subclass, you can run `cls.presets.keys()` to list
+        all built-in presets available on the class.
+
+        Args:
+            preset: string. A built in preset identifier, a Kaggle Models
+                handle, a Hugging Face handle, or a path to a local directory.
+            load_weights: bool. If `True`, the weights will be loaded into the
+                model architecture. If `False`, the weights will be randomly
+                initialized.
+
+        Examples:
+        ```python
+        # Load a Gemma backbone with pre-trained weights.
+        model = keras_hub.models.Backbone.from_preset(
+            "gemma_2b_en",
+        )
+
+        # Load a Bert backbone with a pre-trained config and random weights.
+        model = keras_hub.models.Backbone.from_preset(
+            "bert_base_en",
+            load_weights=False,
+        )
+        ```
+        """
+        format = check_format(preset)
+
+        if format == "transformers":
+            return load_transformers_backbone(cls, preset, load_weights)
+        elif format == "timm":
+            return load_timm_backbone(cls, preset, load_weights, **kwargs)
+
+        preset_cls = check_config_class(preset)
+        if not issubclass(preset_cls, cls):
+            raise ValueError(
+                f"Preset has type `{preset_cls.__name__}` which is not a "
+                f"a subclass of calling class `{cls.__name__}`. Call "
+                f"`from_preset` directly on `{preset_cls.__name__}` instead."
+            )
+
+        backbone = load_serialized_object(preset, CONFIG_FILE, **kwargs)
+        if load_weights:
+            jax_memory_cleanup(backbone)
+            backbone.load_weights(get_file(preset, MODEL_WEIGHTS_FILE))
+
+        return backbone
+
+    def save_to_preset(self, preset_dir):
+        """Save backbone to a preset directory.
+
+        Args:
+            preset_dir: The path to the local model preset directory.
+        """
+        save_serialized_object(self, preset_dir, config_file=CONFIG_FILE)
+        self.save_weights(os.path.join(preset_dir, MODEL_WEIGHTS_FILE))
+        save_metadata(self, preset_dir)
+
+    def enable_lora(self, rank):
+        """Enable Lora on the backbone.
+
+        Calling this method will freeze all weights on the backbone,
+        while enabling Lora on the query & value `EinsumDense` layers
+        of the attention layers.
+        """
+        target_names = ["query_dense", "value_dense", "query", "value"]
+        self.trainable = True
+        self._lora_enabled_layers = []
+        self._lora_rank = rank
+        for layer in self._flatten_layers(include_self=False):
+            layer.trainable = False
+        all_layers = self._flatten_layers(include_self=False)
+        all_layers = [lyr for lyr in all_layers if lyr.weights]
+        for i, layer in enumerate(all_layers):
+            for name in target_names:
+                if layer.name == name:
+                    if hasattr(layer, "enable_lora"):
+                        layer.trainable = True
+                        layer.enable_lora(rank)
+                        self._lora_enabled_layers.append(i)
+
+    def save_lora_weights(self, filepath):
+        if not getattr(self, "_lora_enabled_layers", []):
+            raise ValueError(
+                "There are no lora-enabled layers in this model. "
+                "Make sure to call `.enable_lora(rank)` first."
+            )
+        if not str(filepath).endswith(".lora.h5"):
+            raise ValueError(
+                "The filename must end in `.lora.h5`. "
+                f"Received: filepath={filepath}"
+            )
+
+        store = keras.src.saving.saving_lib.H5IOStore(filepath, mode="w")
+        lora_store = store.make("lora")
+        lora_store["rank"] = self._lora_rank
+        # We cannot identify layers by name since names are non-unique,
+        # so we identify them by index in the topologically sorted list
+        # of layers that have weights.
+        all_layers = self._flatten_layers(include_self=False)
+        all_layers = [lyr for lyr in all_layers if lyr.weights]
+        for layer_index in self._lora_enabled_layers:
+            # We only lora the einsumdense layers,
+            # so the factored weights are always named `kernel`
+            layer = all_layers[layer_index]
+            inner_store = store.make(f"lora/{layer_index}")
+            inner_store["lora_kernel_a"] = layer.lora_kernel_a
+            inner_store["lora_kernel_b"] = layer.lora_kernel_b
+        store.close()
+
+    def load_lora_weights(self, filepath):
+        store = keras.src.saving.saving_lib.H5IOStore(filepath, mode="r")
+        lora_store = store.get("lora")
+        rank = int(lora_store["rank"][()])
+
+        if not getattr(self, "_lora_enabled_layers", []):
+            self.enable_lora(rank)
+        else:
+            if self._lora_rank != rank:
+                raise ValueError(
+                    f"The Lora rank expected by file '{filepath}' "
+                    f"is rank={rank}, but the model was called with "
+                    f"`.enable_lora(rank={self._lora_rank})`. "
+                    "Both ranks must match."
+                )
+        all_layers = self._flatten_layers(include_self=False)
+        all_layers = [lyr for lyr in all_layers if lyr.weights]
+        for layer_index in self._lora_enabled_layers:
+            layer = all_layers[layer_index]
+            lora_kernel_a = store.get(f"lora/{layer_index}")["lora_kernel_a"]
+            lora_kernel_b = store.get(f"lora/{layer_index}")["lora_kernel_b"]
+            layer.lora_kernel_a.assign(lora_kernel_a)
+            layer.lora_kernel_b.assign(lora_kernel_b)
+        store.close()

keras_hub/src/models/bart/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+from keras_hub.src.models.bart.bart_backbone import BartBackbone
+from keras_hub.src.models.bart.bart_presets import backbone_presets
+from keras_hub.src.models.bart.bart_tokenizer import BartTokenizer
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, (BartBackbone, BartTokenizer))

keras_hub/src/models/bart/bart_backbone.py

@@ -0,0 +1,261 @@
+# 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_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
+from keras_hub.src.layers.modeling.reversible_embedding import (
+    ReversibleEmbedding,
+)
+from keras_hub.src.layers.modeling.transformer_decoder import TransformerDecoder
+from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
+from keras_hub.src.models.backbone import Backbone
+
+
+def bart_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras_hub_export("keras_hub.models.BartBackbone")
+class BartBackbone(Backbone):
+    """BART encoder-decoder network.
+
+    This class implements a Transformer-based encoder-decoder model as
+    described in
+    ["BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension"](https://arxiv.org/abs/1910.13461).
+
+    The default constructor gives a fully customizable, randomly initialized BART
+    model with any number of layers, heads, and embedding dimensions. To load
+    preset architectures and weights, use the `from_preset` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available
+    [here](https://github.com/facebookresearch/fairseq/).
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer encoder layers and
+            transformer decoder layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
+
+    Examples:
+    ```python
+    input_data = {
+        "encoder_token_ids": np.ones(shape=(1, 12), dtype="int32"),
+        "encoder_padding_mask": np.array(
+            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]
+        ),
+        "decoder_token_ids": np.ones(shape=(1, 12), dtype="int32"),
+        "decoder_padding_mask": np.array(
+            [[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]
+        ),
+    }
+
+    # Pretrained BART encoder.
+    model = keras_hub.models.BartBackbone.from_preset("bart_base_en")
+    model(input_data)
+
+    # Randomly initialized BART encoder-decoder model with a custom config
+    model = keras_hub.models.BartBackbone(
+        vocabulary_size=50265,
+        num_layers=6,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
+    output = model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        dropout=0.1,
+        max_sequence_length=1024,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            embeddings_initializer=bart_kernel_initializer(),
+            dtype=dtype,
+            name="token_embedding",
+        )
+        self.encoder_position_embedding = PositionEmbedding(
+            initializer=bart_kernel_initializer(),
+            sequence_length=max_sequence_length,
+            dtype=dtype,
+            name="encoder_position_embedding",
+        )
+        self.encoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="encoder_embeddings_add",
+        )
+        self.encoder_embeddings_layer_norm = keras.layers.LayerNormalization(
+            axis=-1,
+            epsilon=1e-5,
+            dtype=dtype,
+            name="encoder_embeddings_layer_norm",
+        )
+        self.encoder_embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="encoder_embeddings_dropout",
+        )
+        self.encoder_transformer_layers = []
+        for i in range(num_layers):
+            layer = TransformerEncoder(
+                num_heads=num_heads,
+                intermediate_dim=intermediate_dim,
+                activation=keras.activations.gelu,
+                dropout=dropout,
+                layer_norm_epsilon=1e-5,
+                kernel_initializer=bart_kernel_initializer(),
+                dtype=dtype,
+                name=f"transformer_encoder_layer_{i}",
+            )
+            self.encoder_transformer_layers.append(layer)
+        self.decoder_position_embedding = PositionEmbedding(
+            initializer=bart_kernel_initializer(),
+            sequence_length=max_sequence_length,
+            dtype=dtype,
+            name="decoder_position_embedding",
+        )
+        self.decoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="decoder_embeddings_add",
+        )
+        self.decoder_embeddings_layer_norm = keras.layers.LayerNormalization(
+            axis=-1,
+            epsilon=1e-5,
+            dtype=dtype,
+            name="decoder_embeddings_layer_norm",
+        )
+        self.decoder_embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="decoder_embeddings_dropout",
+        )
+        self.decoder_transformer_layers = []
+        for i in range(num_layers):
+            layer = TransformerDecoder(
+                intermediate_dim=intermediate_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                activation=keras.activations.gelu,
+                layer_norm_epsilon=1e-5,
+                kernel_initializer=bart_kernel_initializer(),
+                dtype=dtype,
+                name=f"transformer_decoder_layer_{i}",
+            )
+            self.decoder_transformer_layers.append(layer)
+
+        # === Functional Model ===
+        encoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_token_ids"
+        )
+        encoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_padding_mask"
+        )
+        decoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        decoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_padding_mask"
+        )
+        # Encoder.
+        tokens = self.token_embedding(encoder_token_id_input)
+        positions = self.encoder_position_embedding(tokens)
+        x = self.encoder_embeddings_add((tokens, positions))
+        x = self.encoder_embeddings_layer_norm(x)
+        x = self.encoder_embeddings_dropout(x)
+        for transformer_layer in self.encoder_transformer_layers:
+            x = transformer_layer(x, padding_mask=encoder_padding_mask_input)
+        encoder_output = x
+        # Decoder.
+        tokens = self.token_embedding(decoder_token_id_input)
+        positions = self.decoder_position_embedding(tokens)
+        x = self.decoder_embeddings_add((tokens, positions))
+        x = self.decoder_embeddings_layer_norm(x)
+        x = self.decoder_embeddings_dropout(x)
+        for transformer_layer in self.decoder_transformer_layers:
+            x = transformer_layer(
+                decoder_sequence=x,
+                encoder_sequence=encoder_output,
+                decoder_padding_mask=decoder_padding_mask_input,
+                encoder_padding_mask=encoder_padding_mask_input,
+            )
+        decoder_output = x
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "encoder_token_ids": encoder_token_id_input,
+                "encoder_padding_mask": encoder_padding_mask_input,
+                "decoder_token_ids": decoder_token_id_input,
+                "decoder_padding_mask": decoder_padding_mask_input,
+            },
+            outputs={
+                "encoder_sequence_output": encoder_output,
+                "decoder_sequence_output": decoder_output,
+            },
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_sequence_length = max_sequence_length
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "dropout": self.dropout,
+                "max_sequence_length": self.max_sequence_length,
+            }
+        )
+
+        return config