# Copyright 2023 The KerasNLP 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.

"""XLM-RoBERTa backbone model."""

import copy

from keras_nlp.src.api_export import keras_nlp_export
from keras_nlp.src.models.roberta import roberta_backbone
from keras_nlp.src.models.xlm_roberta.xlm_roberta_presets import backbone_presets
from keras_nlp.src.utils.python_utils import classproperty


@keras_nlp_export("keras_nlp.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.

    Examples:
    ```python
    input_data = {
        "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
        "padding_mask": tf.constant(
            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
    }

    # Pretrained XLM-R encoder.
    model = keras_nlp.models.XLMRobertaBackbone.from_preset(
        "xlm_roberta_base_multi",
    )
    model(input_data)

    # Randomly initialized XLM-R model with custom config.
    model = keras_nlp.models.XLMRobertaBackbone(
        vocabulary_size=250002,
        num_layers=4,
        num_heads=4,
        hidden_dim=256,
        intermediate_dim=512,
        max_sequence_length=128
    )
    model(input_data)
    ```
    """

    @classproperty
    def presets(cls):
        return copy.deepcopy(backbone_presets)

