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

keras_hub/src/models/xlm_roberta/xlm_roberta_backbone.py

@@ -0,0 +1,81 @@
+# 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.api_export import keras_hub_export
+from keras_hub.src.models.roberta import roberta_backbone
+
+
+@keras_hub_export("keras_hub.models.XLMRobertaBackbone")
+class XLMRobertaBackbone(roberta_backbone.RobertaBackbone):
+    """An XLM-RoBERTa encoder network.
+
+    This class implements a bi-directional Transformer-based encoder as
+    described in ["Unsupervised Cross-lingual Representation Learning at Scale"](https://arxiv.org/abs/1911.02116).
+    It includes the embedding lookups and transformer layers, but it does not
+    include the masked language modeling head used during pretraining.
+
+    The default constructor gives a fully customizable, randomly initialized
+    RoBERTa encoder 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 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 layer.
+        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 this encoder can
+            consume. The sequence length of the input must be less than
+            `max_sequence_length` default value. 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 = {
+        "token_ids": np.ones(shape=(1, 12), dtype="int32"),
+        "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
+    }
+
+    # Pretrained XLM-R encoder.
+    model = keras_hub.models.XLMRobertaBackbone.from_preset(
+        "xlm_roberta_base_multi",
+    )
+    model(input_data)
+
+    # Randomly initialized XLM-R model with custom config.
+    model = keras_hub.models.XLMRobertaBackbone(
+        vocabulary_size=250002,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_sequence_length=128
+    )
+    model(input_data)
+    ```
+    """

keras_hub/src/models/xlm_roberta/xlm_roberta_classifier.py

@@ -0,0 +1,225 @@
+# 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.models.classifier import Classifier
+from keras_hub.src.models.roberta.roberta_backbone import (
+    roberta_kernel_initializer,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
+    XLMRobertaPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.XLMRobertaClassifier")
+class XLMRobertaClassifier(Classifier):
+    """An end-to-end XLM-RoBERTa model for classification tasks.
+
+    This model attaches a classification head to a
+    `keras_hub.model.XLMRobertaBackbone` instance, mapping from the backbone
+    outputs to logits suitable for a classification task. For usage of
+    this model with pre-trained weights, see the `from_preset()` constructor.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
+
+    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:
+        backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
+        num_classes: int. Number of classes to predict.
+        preprocessor: A `keras_hub.models.XLMRobertaPreprocessor` or `None`. If
+            `None`, this model will not apply preprocessing, and inputs should
+            be preprocessed before calling the model.
+        activation: Optional `str` or callable. The activation function to use
+            on the model outputs. Set `activation="softmax"` to return output
+            probabilities. Defaults to `None`.
+        hidden_dim: int. The size of the pooler layer.
+        dropout: float. The dropout probability value, applied to the pooled
+            output, and after the first dense layer.
+
+    Examples:
+
+    Raw string data.
+    ```python
+    features = ["The quick brown fox jumped.", "نسيت الواجب"]
+    labels = [0, 3]
+
+    # Pretrained classifier.
+    classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
+        "xlm_roberta_base_multi",
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    classifier.predict(x=features, batch_size=2)
+
+    # Re-compile (e.g., with a new learning rate).
+    classifier.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+        jit_compile=True,
+    )
+    # Access backbone programmatically (e.g., to change `trainable`).
+    classifier.backbone.trainable = False
+    # Fit again.
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Preprocessed integer data.
+    ```python
+    features = {
+        "token_ids": np.ones(shape=(2, 12), dtype="int32"),
+        "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]] * 2),
+    }
+    labels = [0, 3]
+
+    # Pretrained classifier without preprocessing.
+    classifier = keras_hub.models.XLMRobertaClassifier.from_preset(
+        "xlm_roberta_base_multi",
+        num_classes=4,
+        preprocessor=None,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Custom backbone and vocabulary.
+    ```python
+    features = ["The quick brown fox jumped.", "نسيت الواجب"]
+    labels = [0, 3]
+
+    def train_sentencepiece(ds, vocab_size):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=ds.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=vocab_size,
+            model_type="WORD",
+            unk_id=0,
+            bos_id=1,
+            eos_id=2,
+        )
+        return bytes_io.getvalue()
+    ds = tf.data.Dataset.from_tensor_slices(
+        ["the quick brown fox", "the earth is round"]
+    )
+    proto = train_sentencepiece(ds, vocab_size=10)
+    tokenizer = keras_hub.models.XLMRobertaTokenizer(
+        proto=proto
+    )
+    preprocessor = keras_hub.models.XLMRobertaPreprocessor(
+        tokenizer,
+        sequence_length=128,
+    )
+    backbone = keras_hub.models.XLMRobertaBackbone(
+        vocabulary_size=250002,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_sequence_length=128,
+    )
+    classifier = keras_hub.models.XLMRobertaClassifier(
+        backbone=backbone,
+        preprocessor=preprocessor,
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = XLMRobertaBackbone
+    preprocessor_cls = XLMRobertaPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        activation=None,
+        hidden_dim=None,
+        dropout=0.0,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.pooled_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=backbone.dtype_policy,
+            name="pooled_dropout",
+        )
+        hidden_dim = hidden_dim or backbone.hidden_dim
+        self.pooled_dense = keras.layers.Dense(
+            hidden_dim,
+            activation="tanh",
+            dtype=backbone.dtype_policy,
+            name="pooled_dense",
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=backbone.dtype_policy,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=roberta_kernel_initializer(),
+            activation=activation,
+            dtype=backbone.dtype_policy,
+            name="logits",
+        )
+
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = self.pooled_dropout(x)
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.num_classes = num_classes
+        self.activation = keras.activations.get(activation)
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "activation": keras.activations.serialize(self.activation),
+                "hidden_dim": self.hidden_dim,
+                "dropout": self.dropout,
+            }
+        )
+        return config

keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm.py

@@ -0,0 +1,141 @@
+# 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.masked_lm_head import MaskedLMHead
+from keras_hub.src.models.masked_lm import MaskedLM
+from keras_hub.src.models.roberta.roberta_backbone import (
+    roberta_kernel_initializer,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_masked_lm_preprocessor import (
+    XLMRobertaMaskedLMPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.XLMRobertaMaskedLM")
+class XLMRobertaMaskedLM(MaskedLM):
+    """An end-to-end XLM-RoBERTa model for the masked language modeling task.
+
+    This model will train XLM-RoBERTa on a masked language modeling task.
+    The model will predict labels for a number of masked tokens in the
+    input data. For usage of this model with pre-trained weights, see the
+    `from_preset()` method.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case inputs can be raw string features during `fit()`, `predict()`,
+    and `evaluate()`. Inputs will be tokenized and dynamically masked during
+    training and evaluation. This is done by default when creating the model
+    with `from_preset()`.
+
+    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:
+        backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
+        preprocessor: A `keras_hub.models.XLMRobertaMaskedLMPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
+
+    Examples:
+
+    Raw string inputs and pretrained backbone.
+    ```python
+    # Create a dataset with raw string features. Labels are inferred.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+
+    # Pretrained language model
+    # on an MLM task.
+    masked_lm = keras_hub.models.XLMRobertaMaskedLM.from_preset(
+        "xlm_roberta_base_multi",
+    )
+    masked_lm.fit(x=features, batch_size=2)
+    ```
+
+    # Re-compile (e.g., with a new learning rate).
+    masked_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+        jit_compile=True,
+    )
+    # Access backbone programmatically (e.g., to change `trainable`).
+    masked_lm.backbone.trainable = False
+    # Fit again.
+    masked_lm.fit(x=features, batch_size=2)
+    ```
+
+    Preprocessed integer data.
+    ```python
+    # Create a preprocessed dataset where 0 is the mask token.
+    features = {
+        "token_ids": np.array([[1, 2, 0, 4, 0, 6, 7, 8]] * 2),
+        "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1]] * 2),
+        "mask_positions": np.array([[2, 4]] * 2)
+    }
+    # Labels are the original masked values.
+    labels = [[3, 5]] * 2
+
+    masked_lm = keras_hub.models.XLMRobertaMaskedLM.from_preset(
+        "xlm_roberta_base_multi",
+        preprocessor=None,
+    )
+
+    masked_lm.fit(x=features, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = XLMRobertaBackbone
+    preprocessor_cls = XLMRobertaMaskedLMPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        preprocessor=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=roberta_kernel_initializer(),
+            dtype=backbone.dtype_policy,
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
+        inputs = {
+            **backbone.input,
+            "mask_positions": keras.Input(
+                shape=(None,), dtype="int32", name="mask_positions"
+            ),
+        }
+        backbone_outputs = backbone(backbone.input)
+        outputs = self.masked_lm_head(
+            backbone_outputs, inputs["mask_positions"]
+        )
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )

keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm_preprocessor.py

@@ -0,0 +1,195 @@
+# 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 absl import logging
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
+    MaskedLMMaskGenerator,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_preprocessor import (
+    XLMRobertaPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.XLMRobertaMaskedLMPreprocessor")
+class XLMRobertaMaskedLMPreprocessor(XLMRobertaPreprocessor):
+    """XLM-RoBERTa preprocessing for the masked language modeling task.
+
+    This preprocessing layer will prepare inputs for a masked language modeling
+    task. It is primarily intended for use with the
+    `keras_hub.models.XLMRobertaMaskedLM` task model. Preprocessing will occur in
+    multiple steps.
+
+    1. Tokenize any number of input segments using the `tokenizer`.
+    2. Pack the inputs together with the appropriate `"<s>"`, `"</s>"` and
+      `"<pad>"` tokens, i.e., adding a single `"<s>"` at the start of the
+      entire sequence, `"</s></s>"` between each segment,
+      and a `"</s>"` at the end of the entire sequence.
+    3. Randomly select non-special tokens to mask, controlled by
+      `mask_selection_rate`.
+    4. Construct a `(x, y, sample_weight)` tuple suitable for training with a
+      `keras_hub.models.XLMRobertaMaskedLM` task model.
+
+    Args:
+        tokenizer: A `keras_hub.models.XLMRobertaTokenizer` instance.
+        sequence_length: int. The length of the packed inputs.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It supports an arbitrary number of segments.
+        mask_selection_rate: float. The probability an input token will be
+            dynamically masked.
+        mask_selection_length: int. The maximum number of masked tokens
+            in a given sample.
+        mask_token_rate: float. The probability the a selected token will be
+            replaced with the mask token.
+        random_token_rate: float. The probability the a selected token will be
+            replaced with a random token from the vocabulary. A selected token
+            will be left as is with probability
+            `1 - mask_token_rate - random_token_rate`.
+
+    Call arguments:
+        x: A tensor of single string sequences, or a tuple of multiple
+            tensor sequences to be packed together. Inputs may be batched or
+            unbatched. For single sequences, raw python inputs will be converted
+            to tensors. For multiple sequences, pass tensors directly.
+        y: Label data. Should always be `None` as the layer generates labels.
+        sample_weight: Label weights. Should always be `None` as the layer
+            generates label weights.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_hub.models.XLMRobertaMaskedLMPreprocessor.from_preset(
+        "xlm_roberta_base_multi"
+    )
+
+    # Tokenize and mask a single sentence.
+    preprocessor("The quick brown fox jumped.")
+    # Tokenize and mask a batch of single sentences.
+    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
+    # Tokenize and mask sentence pairs.
+    # In this case, always convert input to tensors before calling the layer.
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+    preprocessor((first, second))
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.XLMRobertaMaskedLMPreprocessor.from_preset(
+        "xlm_roberta_base_multi"
+    )
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+
+    # Map single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(first)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map sentence pairs.
+    ds = tf.data.Dataset.from_tensor_slices((first, second))
+    # Watch out for tf.data's default unpacking of tuples here!
+    # Best to invoke the `preprocessor` directly in this case.
+    ds = ds.map(
+        lambda first, second: preprocessor(x=(first, second)),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=512,
+        truncate="round_robin",
+        mask_selection_rate=0.15,
+        mask_selection_length=96,
+        mask_token_rate=0.8,
+        random_token_rate=0.1,
+        **kwargs,
+    ):
+        super().__init__(
+            tokenizer,
+            sequence_length=sequence_length,
+            truncate=truncate,
+            **kwargs,
+        )
+        self.mask_selection_rate = mask_selection_rate
+        self.mask_selection_length = mask_selection_length
+        self.mask_token_rate = mask_token_rate
+        self.random_token_rate = random_token_rate
+        self.masker = None
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # Defer masker creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.masker = MaskedLMMaskGenerator(
+            mask_selection_rate=self.mask_selection_rate,
+            mask_selection_length=self.mask_selection_length,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+            vocabulary_size=self.tokenizer.vocabulary_size(),
+            mask_token_id=self.tokenizer.mask_token_id,
+            unselectable_token_ids=[
+                self.tokenizer.start_token_id,
+                self.tokenizer.end_token_id,
+                self.tokenizer.pad_token_id,
+            ],
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config
+
+    def call(self, x, y=None, sample_weight=None):
+        if y is not None or sample_weight is not None:
+            logging.warning(
+                f"{self.__class__.__name__} generates `y` and `sample_weight` "
+                "based on your input data, but your data already contains `y` "
+                "or `sample_weight`. Your `y` and `sample_weight` will be "
+                "ignored."
+            )
+
+        x = super().call(x)
+        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
+        masker_outputs = self.masker(token_ids)
+        x = {
+            "token_ids": masker_outputs["token_ids"],
+            "padding_mask": padding_mask,
+            "mask_positions": masker_outputs["mask_positions"],
+        }
+        y = masker_outputs["mask_ids"]
+        sample_weight = masker_outputs["mask_weights"]
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)