# 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.

"""OPT Causal LM preprocessor layer."""

from absl import logging

from keras_nlp.api_export import keras_nlp_export
from keras_nlp.models.opt.opt_preprocessor import OPTPreprocessor
from keras_nlp.utils.keras_utils import pack_x_y_sample_weight


@keras_nlp_export("keras_nlp.models.OPTCausalLMPreprocessor")
class OPTCausalLMPreprocessor(OPTPreprocessor):
    """OPT Causal LM preprocessor.

    This preprocessing layer is primarily meant to be used with
    `keras_nlp.models.OPTCausalLM`. By default, it will take in batches of
    strings, and return outputs in a `(x, y, sample_weight)` format, where the
    `y` label is the next token id in the `x` sequence. For use with generation,
    pass `return_labels=False` in which case the output will simply be the
    encoded string features.

    Args:
        tokenizer: A `keras_nlp.models.OPTTokenizer` instance.
        sequence_length: The length of the packed inputs.
        add_start_token: If true, the preprocessor will append the tokenizer
            start token to each input sequence.
        add_end_token: If true, the preprocessor will append the tokenizer
            end token to each input sequence.

    Call arguments:
        x: A string `tf.Tensor` or list of python strings.
        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.
        sequence_length: Pass to override the configured `sequence_length` of
            the layer.
        add_start_token: Pass to override the configured value of
            `add_start_token` on the layer.
        add_end_token: Pass to override the configured value of
            `add_end_token` on the layer.
        return_labels: If `True`, the output `"token_ids"` will be offset by one
            and returned as labels. If `False` only features will be returned.

    Examples:
    ```python
    # Load the preprocessor from a preset.
    preprocessor = keras_nlp.models.OPTCausalLMPreprocessor.from_preset(
        "opt_125m_en"
    )

    # Tokenize and pack a single sentence.
    sentence = tf.constant("League of legends")
    preprocessor(sentence)
    # Same output.
    preprocessor("League of legends")

    # Tokenize a batch of sentences.
    sentences = tf.constant(["Taco tuesday", "Fish taco please!"])
    preprocessor(sentences)
    # Same output.
    preprocessor(["Taco tuesday", "Fish taco please!"])

    # Map a dataset to preprocess a single sentence.
    features = tf.constant(
        [
            "Avatar 2 is amazing!",
            "Well, I am not sure.",
        ]
    )
    labels = tf.constant([1, 0])
    ds = tf.data.Dataset.from_tensor_slices((features, labels))
    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)

    # Map a dataset to preprocess unlabled sentences.
    ds = tf.data.Dataset.from_tensor_slices(features)
    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
    """

    def call(
        self,
        x,
        y=None,
        sample_weight=None,
        sequence_length=None,
        add_start_token=None,
        add_end_token=None,
        return_labels=True,
    ):
        if y is not None or sample_weight is not None:
            logging.warning(
                "`OPTCausalLMPreprocessor` 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."
            )
        if return_labels:
            # Tokenize with one extra token to account for the truncation below.
            sequence_length = (sequence_length or self.sequence_length) + 1
        x = super().call(
            x,
            sequence_length=sequence_length,
            add_start_token=add_start_token,
            add_end_token=add_end_token,
        )
        if return_labels:
            token_ids, padding_mask = x["token_ids"], x["padding_mask"]
            # The last token does not have a next token, so we truncate it out.
            x = {
                "token_ids": token_ids[..., :-1],
                "padding_mask": padding_mask[..., :-1],
            }
            # Target `y` will be the next token.
            y = token_ids[..., 1:]
            sample_weight = padding_mask[..., 1:]
            return pack_x_y_sample_weight(x, y, sample_weight)
        else:
            return x
