# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# 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
#
#     http://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 shutil

import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from paddlenlp.ops import (
    InferBartDecoding,
    InferGptDecoding,
    InferGptJDecoding,
    InferMBartDecoding,
    InferMIRODecoding,
    InferOptDecoding,
    InferPegasusDecoding,
    InferT5Decoding,
    InferTransformerDecoding,
    InferUnifiedDecoding,
)
from paddlenlp.transformers import (
    BartPretrainedModel,
    CodeGenPreTrainedModel,
    GPTChineseTokenizer,
    GPTJPretrainedModel,
    GPTPretrainedModel,
    GPTTokenizer,
    InferTransformerModel,
    MBartPretrainedModel,
    OPTPretrainedModel,
    PegasusPretrainedModel,
    PositionalEmbedding,
    T5PretrainedModel,
    TransformerModel,
    UnifiedTransformerPretrainedModel,
    UNIMOPretrainedModel,
    WordEmbedding,
    position_encoding_init,
)
from paddlenlp.utils.log import logger

from .encoder import enable_fast_encoder


class FasterTransformer(TransformerModel):
    """
    FasterTransformer is a fast version for generation with the Transformer
    model. It uses a custom op based on and enhancing NV FasterTransformer to
    do fast generation.

    Args:
        src_vocab_size (int):
            The size of source vocabulary.
        trg_vocab_size (int):
            The size of target vocabulary.
        max_length (int):
            The maximum length of input sequences.
        num_encoder_layers (int):
            The number of sub-layers to be stacked in the encoder.
        num_decoder_layers (int):
            The number of sub-layers to be stacked in the decoder.
        n_head (int):
            The number of head used in multi-head attention.
        d_model (int):
            The dimension for word embeddings, which is also the last dimension of
            the input and output of multi-head attention, position-wise feed-forward
            networks, encoder and decoder.
        d_inner_hid (int):
            Size of the hidden layer in position-wise feed-forward networks.
        dropout (float):
            Dropout rates. Used for pre-process, activation and inside attention.
        weight_sharing (bool):
            Whether to use weight sharing.
        attn_dropout (float):
            The dropout probability used in MHA to drop some attention target.
            If None, use the value of dropout. Defaults to None.
        act_dropout (float):
            The dropout probability used after FFN activition. If None, use
            the value of dropout. Defaults to None.
        bos_id (int, optional):
            The start token id and also is used as padding id. Defaults to 0.
        eos_id (int, optional):
            The end token id. Defaults to 1.
        pad_id (int, optional):
            The pad token id. Defaults to None. If it's None, the bos_id will be used as pad_id.
        decoding_strategy (str, optional):
            Indicating the strategy of decoding. It can be 'beam_search', 'beam_search_v2',
            'topk_sampling' and 'topp_sampling'. For beam search strategies,
            'v2' would select the top `beam_size * 2` beams and process the top
            `beam_size` alive and finish beams in them separately, while 'v1'
            would only select the top `beam_size` beams and mix up the alive and
            finish beams. 'v2' always searchs more and get better results, since
            the alive beams would always be `beam_size` while the number of alive
            beams in `v1` might decrease when meeting the end token. However,
            'v2' always generates longer results thus might do more calculation
            and be slower.
        beam_size (int, optional):
            The beam width for beam search. Defaults to 4.
        topk (int, optional):
            The number of highest probability tokens to keep for top-k sampling.
            Defaults to 4.
        topp (float, optional):
            The most probable tokens whose cumulative probability is not less than
            `topp` are kept for top-p sampling. Defaults to 4.
        max_out_len (int, optional):
            The maximum output length. Defaults to 256.
        diversity_rate (float, optional):
            Refer to `A Simple, Fast Diverse Decoding Algorithm for Neural Generation <https://arxiv.org/abs/1611.08562>`_
            for details. Bigger `diversity_rate` would lead to more diversity.
            if `diversity_rate == 0` is equivalent to naive BeamSearch. Default
            to 0 if not set.
        use_fp16_decoding(bool, optional):
            Whether to use fp16 for decoding.
        enable_fast_encoder(bool, optional):
            Whether to use the fast version of encoder. This is experimental option for now.
            Defaults to False.
        use_fp16_encoder(bool, optional):
            Whether to use fp16 for encoder. Only works when enable_fast_encoder is True.
            Defaults to False.
        rel_len(bool, optional):
            Indicating whether `max_out_len` in is the length relative to that
            of source text. Only works in `v2` temporarily. It is suggest to set
            a small `max_out_len` and use `rel_len=True`. Default to False if
            not set.
        alpha(float, optional):
            The power number in length penalty calculation. Only works in `v2`
            temporarily. Refer to `GNMT <https://arxiv.org/pdf/1609.08144.pdf>`_.
            Default to 0.6 if not set.
    """

    def __init__(
        self,
        src_vocab_size,
        trg_vocab_size,
        max_length,
        num_encoder_layers,
        num_decoder_layers,
        n_head,
        d_model,
        d_inner_hid,
        dropout,
        weight_sharing,
        attn_dropout=None,
        act_dropout=None,
        bos_id=0,
        eos_id=1,
        pad_id=None,
        decoding_strategy="beam_search",
        beam_size=4,
        topk=1,
        topp=0.0,
        max_out_len=256,
        diversity_rate=0.0,
        decoding_lib=None,
        use_fp16_decoding=False,
        enable_fast_encoder=False,
        use_fp16_encoder=False,
        rel_len=False,
        alpha=0.6,
    ):
        # if decoding_lib is None:
        #     raise ValueError(
        #         "The args decoding_lib must be set to use FasterTransformer. ")
        # elif not os.path.exists(decoding_lib):
        #     raise ValueError("The path to decoding lib is not exist.")

        args = dict(locals())
        args.pop("self")
        args.pop("__class__", None)
        self.decoding_strategy = args.pop("decoding_strategy")
        self.beam_size = args.pop("beam_size")
        self.topk = args.pop("topk")
        self.topp = args.pop("topp")
        self.max_out_len = args.pop("max_out_len")
        self.diversity_rate = args.pop("diversity_rate")
        self.decoding_lib = args.pop("decoding_lib")
        self.use_fp16_decoding = args.pop("use_fp16_decoding")
        self.enable_fast_encoder = args.pop("enable_fast_encoder")
        self.use_fp16_encoder = args.pop("use_fp16_encoder")
        self.rel_len = args.pop("rel_len")
        self.alpha = args.pop("alpha")
        self.dropout = dropout
        self.weight_sharing = weight_sharing
        self.trg_vocab_size = trg_vocab_size
        self.d_model = d_model
        self.bos_id = bos_id
        self.pad_id = pad_id if pad_id is not None else self.bos_id
        self.max_length = max_length
        super(FasterTransformer, self).__init__(**args)

        if self.enable_fast_encoder:
            logger.warning("enable_fast_encoder is an experimental option and subject to change.")
        elif self.use_fp16_encoder:
            self.use_fp16_encoder = False

        self.decoding_linear = nn.Linear(in_features=d_model, out_features=trg_vocab_size)

        if weight_sharing:
            self.trg_word_embedding = WordEmbedding(vocab_size=trg_vocab_size, emb_dim=d_model, bos_id=self.bos_id)
            self.trg_pos_embedding = PositionalEmbedding(emb_dim=d_model, max_length=max_length)

        self.decoding = InferTransformerDecoding(
            decoder=self.transformer.decoder,
            word_embedding=self.trg_word_embedding.word_embedding,
            positional_embedding=self.trg_pos_embedding.pos_encoder,
            linear=self.decoding_linear,
            num_decoder_layers=num_decoder_layers,
            n_head=n_head,
            d_model=d_model,
            bos_id=bos_id,
            eos_id=eos_id,
            decoding_strategy=decoding_strategy,
            beam_size=beam_size,
            topk=topk,
            topp=topp,
            max_out_len=max_out_len,
            diversity_rate=self.diversity_rate,
            decoding_lib=self.decoding_lib,
            use_fp16_decoding=self.use_fp16_decoding,
            rel_len=self.rel_len,
            alpha=self.alpha,
        )

    def forward(self, src_word, trg_word=None):
        src_max_len = paddle.shape(src_word)[-1]
        src_slf_attn_bias = (
            paddle.cast(src_word == self.pad_id, dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e9
        )
        src_pos = paddle.cast(src_word != self.pad_id, dtype=src_word.dtype) * paddle.arange(start=0, end=src_max_len)

        # Run encoder
        src_emb = self.src_word_embedding(src_word)
        src_pos_emb = self.src_pos_embedding(src_pos)
        src_emb = src_emb + src_pos_emb
        enc_input = F.dropout(src_emb, p=self.dropout, training=False) if self.dropout else src_emb

        if self.enable_fast_encoder and self.use_fp16_encoder:
            enc_input = paddle.cast(enc_input, dtype="float16")

        enc_output = self.transformer.encoder(enc_input, src_slf_attn_bias)

        if self.use_fp16_decoding and enc_output.dtype != paddle.float16:
            enc_output = paddle.cast(enc_output, dtype="float16")
        elif not self.use_fp16_decoding and enc_output.dtype != paddle.float32:
            enc_output = paddle.cast(enc_output, dtype="float32")

        mem_seq_lens = paddle.sum(paddle.cast(src_word != self.pad_id, dtype="int32"), dtype="int32", axis=1)
        ids = self.decoding(enc_output, mem_seq_lens, trg_word=trg_word)

        return ids

    def load(self, init_from_params=None, state_dict=None):
        # Load the trained model
        if init_from_params is None and state_dict is None:
            raise ValueError("Either init_from_params or state_dict must be given to load the infer model. ")

        if state_dict is None:
            state_dict = paddle.load(init_from_params, return_numpy=True)
        else:
            for state in state_dict:
                # NOTE: This API only used in dygraph, so paddle.Tensor is enough.
                if isinstance(state_dict[state], paddle.Tensor):
                    state_dict[state] = state_dict[state].numpy()

        # To set weight[padding_idx] to 0.
        state_dict["trg_word_embedding.word_embedding.weight"][self.bos_id] = [0] * self.d_model

        # Dealing with weight sharing.
        if self.weight_sharing:
            state_dict["decoding_linear.weight"] = np.transpose(state_dict["trg_word_embedding.word_embedding.weight"])
        else:
            state_dict["decoding_linear.weight"] = state_dict["linear.weight"]

        if self.decoding._fuse_qkv:
            for item in self.state_dict():
                if "decoder" in item and "self_attn.q_proj" in item:
                    num_layer = item.split(".")[3]
                    param_type = item.split(".")[-1]

                    state_dict["decoding.slf_q_" + param_type + "_" + num_layer] = np.concatenate(
                        (
                            state_dict[item],
                            state_dict["transformer.decoder.layers." + num_layer + ".self_attn.k_proj." + param_type],
                            state_dict["transformer.decoder.layers." + num_layer + ".self_attn.v_proj." + param_type],
                        ),
                        axis=-1,
                    )

        if self.use_fp16_decoding:
            for item in self.state_dict():
                if "decoder" in item or "decoding.slf" in item:
                    state_dict[item] = np.float16(state_dict[item])
            state_dict["decoding_linear.weight"] = np.float16(state_dict["decoding_linear.weight"])
            state_dict["trg_word_embedding.word_embedding.weight"] = np.float16(
                state_dict["trg_word_embedding.word_embedding.weight"]
            )
            state_dict["trg_pos_embedding.pos_encoder.weight"] = np.float16(
                state_dict["trg_pos_embedding.pos_encoder.weight"]
            )
            state_dict["decoding_linear.bias"] = np.zeros([self.trg_vocab_size], dtype="float16")

        self.load_dict(state_dict)

        if self.enable_fast_encoder:
            self = enable_fast_encoder(self, use_fp16=self.use_fp16_encoder)

    def export_params(self, init_from_params, place):
        """
        This method is used for load static graph from dygraph checkpoint
        or export inference model using static graph.
        Do NOT support faster encoder.

        Args:
            init_from_params (string):
                The path to dygraph checkpoint.
            place (paddle.Place):
                The place to execute static graph.

        Example:
            .. code-block::
                paddle.enable_static()
                place = "gpu"
                place = paddle.set_device(place)
                reader.adapt_vocab_size(args)

                test_program = paddle.static.Program()
                startup_program = paddle.static.Program()
                with paddle.static.program_guard(test_program, startup_program):
                    src_word = paddle.static.data(
                        name="src_word", shape=[None, None], dtype="int64")

                    # Define model
                    transformer = FasterTransformer(
                        src_vocab_size=args.src_vocab_size,
                        trg_vocab_size=args.trg_vocab_size,
                        max_length=args.max_length + 1,
                        num_encoder_layers=args.n_layer,
                        num_decoder_layers=args.n_layer,
                        n_head=args.n_head,
                        d_model=args.d_model,
                        d_inner_hid=args.d_inner_hid,
                        dropout=args.dropout,
                        weight_sharing=args.weight_sharing,
                        bos_id=args.bos_idx,
                        eos_id=args.eos_idx,
                        decoding_strategy=args.decoding_strategy,
                        beam_size=args.beam_size,
                        max_out_len=args.max_out_len,
                        decoding_lib=args.decoding_lib,
                        use_fp16_decoding=args.use_fp16_decoding,
                        rel_len=args.use_rel_len,
                        alpha=args.alpha)

                    finished_seq = transformer(src_word=src_word)

                test_program = test_program.clone(for_test=True)

                exe = paddle.static.Executor(place)
                exe.run(startup_program)

                # Load checkpoint.
                transformer.export_params(
                    init_from_params=os.path.join(args.init_from_params,
                                                "transformer.pdparams"),
                    place=place)

                paddle.static.save_inference_model(
                    os.path.join(args.inference_model_dir, "transformer"),
                    feed_vars=src_word,
                    fetch_vars=finished_seq,
                    executor=exe,
                    program=test_program)
        """
        # Load the trained model
        assert init_from_params, "Please set init_from_params to load the infer model."

        model_dict = paddle.load(init_from_params, return_numpy=True)

        # To set weight[padding_idx] to 0.
        model_dict["trg_word_embedding.word_embedding.weight"][self.bos_id] = [0] * self.d_model

        # Dealing with weight sharing.
        if self.weight_sharing:
            model_dict["decoding_linear.weight"] = np.transpose(model_dict["trg_word_embedding.word_embedding.weight"])
        else:
            model_dict["decoding_linear.weight"] = model_dict["linear.weight"]

        # To avoid a longer length than training, reset the size of position
        # encoding to max_length
        model_dict["encoder.pos_encoder.weight"] = position_encoding_init(self.max_length, self.d_model)
        model_dict["decoder.pos_encoder.weight"] = position_encoding_init(self.max_length, self.d_model)

        if self.decoding._fuse_qkv:
            for item in self.state_dict():
                if "decoder" in item and "self_attn.q_proj" in item:
                    num_layer = item.split(".")[3]
                    param_type = item.split(".")[-1]

                    model_dict["decoding.slf_q_" + param_type + "_" + num_layer] = np.concatenate(
                        (
                            model_dict[item],
                            model_dict["transformer.decoder.layers." + num_layer + ".self_attn.k_proj." + param_type],
                            model_dict["transformer.decoder.layers." + num_layer + ".self_attn.v_proj." + param_type],
                        ),
                        axis=-1,
                    )

        if self.use_fp16_decoding:
            for item in self.state_dict():
                if "decoder" in item or "decoding.slf" in item:
                    model_dict[item] = np.float16(model_dict[item])
            model_dict["decoding_linear.weight"] = np.float16(model_dict["decoding_linear.weight"])
            model_dict["trg_word_embedding.word_embedding.weight"] = np.float16(
                model_dict["trg_word_embedding.word_embedding.weight"]
            )
            model_dict["trg_pos_embedding.pos_encoder.weight"] = np.float16(
                model_dict["trg_pos_embedding.pos_encoder.weight"]
            )
            model_dict["decoding_linear.bias"] = np.zeros([self.trg_vocab_size], dtype="float16")

        for item in self.state_dict():
            param = self
            attr_list = item.split(".")
            for attr in attr_list:
                param = getattr(param, attr)
            param_name = param.name
            var = paddle.static.global_scope().find_var(param_name).get_tensor()
            var.set(model_dict[item], place)


class TransformerGenerator(paddle.nn.Layer):
    """
    The Transformer model for auto-regressive generation with beam search. It wraps
    `FasterTransformer` and `InferTransformerModel`, and automatically chioces using
    `FasterTransformer` (with jit building) or the slower verison `InferTransformerModel`.

    Args:
        src_vocab_size (int):
            The size of source vocabulary.
        trg_vocab_size (int):
            The size of target vocabulary.
        max_length (int):
            The maximum length of input sequences.
        num_encoder_layers (int):
            The number of sub-layers to be stacked in the encoder.
        num_decoder_layers (int):
            The number of sub-layers to be stacked in the decoder.
        n_head (int):
            The number of head used in multi-head attention.
        d_model (int):
            The dimension for word embeddings, which is also the last dimension of
            the input and output of multi-head attention, position-wise feed-forward
            networks, encoder and decoder.
        d_inner_hid (int):
            Size of the hidden layer in position-wise feed-forward networks.
        dropout (float):
            Dropout rates. Used for pre-process, activation and inside attention.
        weight_sharing (bool):
            Whether to use weight sharing.
        bos_id (int, optional):
            The start token id and also is used as padding id. Defaults to 0.
        eos_id (int, optional):
            The end token id. Defaults to 1.
        beam_size (int, optional):
            The beam width for beam search. Defaults to 4.
        max_out_len (int, optional):
            The maximum output length. Defaults to 256.
        activation (str, optional):
            The activation used in FFN. Defaults to "relu".
        normalize_before (bool, optional):
            Whether to apply pre-normalization. Defaults to True.
        kwargs:
            The key word arguments can be `output_time_major`, `use_ft`, `use_fp16_decoding`,
            `rel_len`, `alpha`:

            - `output_time_major(bool, optional)`: Indicate the data layout of predicted
            Tensor. If `False`, the data layout would be batch major with shape
            `[batch_size, seq_len, beam_size]`. If  `True`, the data layout would
            be time major with shape `[seq_len, batch_size, beam_size]`. Default
            to `False`.

            - `use_ft(bool, optional)`: Whether to use FastGeneration
            for decoding. Default to True if not set.

            - `use_fp16_decoding(bool, optional)`: Whether to use fp16
            for decoding.  Only works when using FastGeneration.

            - `beam_search_version(str, optional)`: Indicating the strategy of
            beam search. It can be 'v1' or 'v2'. 'v2' would select the top
            `beam_size * 2` beams and process the top `beam_size` alive and
            finish beams in them separately, while 'v1' would only select the
            top `beam_size` beams and mix up the alive and finish beams. 'v2' always
            searchs more and get better results, since the alive beams would
            always be `beam_size` while the number of alive beams in `v1` might
            decrease when meeting the end token. However, 'v2' always generates
            longer results thus might do more calculation and be slower.

            - `rel_len(bool, optional)`: Indicating whether `max_out_len` in is
            the length relative to that of source text. Only works in `v2` temporarily.
            It is suggest to set a small `max_out_len` and use `rel_len=True`.
            Default to False if not set.

            - `alpha(float, optional)`: The power number in length penalty
            calculation. Refer to `GNMT <https://arxiv.org/pdf/1609.08144.pdf>`_.
            Only works in `v2` temporarily. Default to 0.6 if not set.

            - diversity_rate(float, optional): Refer to `A Simple, Fast Diverse
            Decoding Algorithm for Neural Generation <https://arxiv.org/abs/1611.08562>`_
            for details. Bigger `diversity_rate` would lead to more diversity.
            if `diversity_rate == 0` is equivalent to naive BeamSearch. Default
            to 0 if not set. **NOTE**: Only works when using FastGeneration
            temporarily.
    """

    def __init__(
        self,
        src_vocab_size,
        trg_vocab_size,
        max_length,
        num_encoder_layers,
        num_decoder_layers,
        n_head,
        d_model,
        d_inner_hid,
        dropout,
        weight_sharing,
        bos_id=0,
        eos_id=1,
        pad_id=None,
        beam_size=4,
        max_out_len=256,
        activation="relu",
        normalize_before=True,
        **kwargs
    ):
        logger.warning("TransformerGenerator is an experimental API and subject to change.")
        # `kwargs` can include output_time_major, use_fp16_decoding, topk, topp.
        # The later three arguments can only work when using FastGeneration,
        # and expose topk, topp later.
        super(TransformerGenerator, self).__init__()
        self.d_model = d_model
        self.max_length = max_length
        self.output_time_major = kwargs.pop("output_time_major", True)
        # Only works for FastGeneration.
        # TODO: original version supports diversity rate.
        diversity_rate = kwargs.pop("diversity_rate", 0.0)
        use_fp16_decoding = kwargs.pop("use_fp16_decoding", False)
        use_ft = kwargs.pop("use_ft", True)
        beam_search_version = kwargs.pop("beam_search_version", "v1")
        rel_len = kwargs.pop("rel_len", False)
        alpha = kwargs.pop("alpha", 0.6)

        # TODO: Faster version needs to update attr to support custom
        # activation and normalize_before which are both aupport in cpp codes.
        if use_ft and activation == "relu" and normalize_before:
            try:
                decoding_strategy = "beam_search_v2" if beam_search_version == "v2" else "beam_search"
                self.transformer = FasterTransformer(
                    src_vocab_size=src_vocab_size,
                    trg_vocab_size=trg_vocab_size,
                    max_length=max_length,
                    num_encoder_layers=num_encoder_layers,
                    num_decoder_layers=num_decoder_layers,
                    n_head=n_head,
                    d_model=d_model,
                    d_inner_hid=d_inner_hid,
                    dropout=dropout,
                    weight_sharing=weight_sharing,
                    bos_id=bos_id,
                    eos_id=eos_id,
                    pad_id=pad_id,
                    beam_size=beam_size,
                    max_out_len=max_out_len,
                    diversity_rate=diversity_rate,
                    decoding_strategy=decoding_strategy,
                    use_fp16_decoding=use_fp16_decoding,
                    rel_len=rel_len,
                    alpha=alpha,
                )
            except Exception:
                logger.warning(
                    "Exception occurs when using FastGeneration. " "The original forward will be involved. "
                )
                if diversity_rate != 0:
                    logger.warning(
                        "diversity_rate would not work since it is only " "supported by FastGeneration temporarily."
                    )
                self.transformer = InferTransformerModel(
                    src_vocab_size=src_vocab_size,
                    trg_vocab_size=trg_vocab_size,
                    max_length=max_length,
                    num_encoder_layers=num_encoder_layers,
                    num_decoder_layers=num_decoder_layers,
                    n_head=n_head,
                    d_model=d_model,
                    d_inner_hid=d_inner_hid,
                    dropout=dropout,
                    weight_sharing=weight_sharing,
                    bos_id=bos_id,
                    eos_id=eos_id,
                    pad_id=pad_id,
                    beam_size=beam_size,
                    max_out_len=max_out_len,
                    output_time_major=self.output_time_major,
                    beam_search_version=beam_search_version,
                    activation=activation,
                    normalize_before=normalize_before,
                    rel_len=rel_len,
                    alpha=alpha,
                )
        else:
            if diversity_rate != 0:
                logger.warning(
                    "diversity_rate would not work since it is only " "supported by FastGeneration temporarily."
                )
            self.transformer = InferTransformerModel(
                src_vocab_size=src_vocab_size,
                trg_vocab_size=trg_vocab_size,
                max_length=max_length,
                num_encoder_layers=num_encoder_layers,
                num_decoder_layers=num_decoder_layers,
                n_head=n_head,
                d_model=d_model,
                d_inner_hid=d_inner_hid,
                dropout=dropout,
                weight_sharing=weight_sharing,
                bos_id=bos_id,
                eos_id=eos_id,
                pad_id=pad_id,
                beam_size=beam_size,
                max_out_len=max_out_len,
                output_time_major=self.output_time_major,
                beam_search_version=beam_search_version,
                activation=activation,
                normalize_before=normalize_before,
                rel_len=rel_len,
                alpha=alpha,
            )

    def forward(self, src_word, trg_word=None):
        r"""
        Performs decoding for transformer model.

        Args:
            src_word (Tensor):
                The ids of source sequence words. It is a tensor with shape
                `[batch_size, source_sequence_length]` and its data type can be
                int or int64.
            trg_word (Tensor):
                The ids of target sequence words. Normally, it should NOT be
                given. If it's given, force decoding with previous output token
                will be trigger. Defaults to None.

        Returns:
            Tensor:
                An int64 tensor shaped indicating the predicted ids. Its shape is
                `[batch_size, seq_len, beam_size]` or `[seq_len, batch_size, beam_size]`
                according to `output_time_major`. While, when using FastGeneration
                and beam search v2, the beam dimension would be doubled to include
                both the top `beam_size` alive and finish beams, thus the tensor
                shape is `[batch_size, seq_len, beam_size * 2]` or `[seq_len, batch_size, beam_size * 2]`.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.ops import TransformerGenerator

                transformer = TransformerGenerator(
                    src_vocab_size=30000,
                    trg_vocab_size=30000,
                    max_length=256,
                    num_encoder_layers=6,
                    num_decoder_layers=6,
                    n_head=8,
                    d_model=512,
                    d_inner_hid=2048,
                    dropout=0.1,
                    weight_sharing=True,
                    bos_id=0,
                    eos_id=1,
                    beam_size=4,
                    max_out_len=256)

                batch_size = 5
                seq_len = 10
                transformer(
                    src_word=paddle.randint(low=3, high=30000, shape=[batch_size, seq_len]))
        """
        out = self.transformer(src_word, trg_word=trg_word)
        # TODO(guosheng): FasterTransformer has an output with layout
        # `[seq_len, batch_size, beam_size]`. While the output layout of
        # original one is `[batch_size, seq_len, beam_size]`. Maybe we need
        # unify them later.
        if not self.output_time_major and isinstance(self.transformer, FasterTransformer):
            out = paddle.transpose(out, [1, 0, 2])
        return out

    def load(self, path=None, state_dict=None):
        if path is None and state_dict is None:
            raise ValueError("Either path or state_dict must be given to load the infer model. ")

        if isinstance(self.transformer, FasterTransformer):
            self.transformer.load(path, state_dict)
        else:
            if state_dict is None:
                state_dict = paddle.load(path)
            self.transformer.load_dict(state_dict)


class FasterOPT(OPTPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterOPT, self).__init__(model.config)
        self._model = model
        self.use_fp16_decoding = use_fp16_decoding
        self.decoding = InferOptDecoding(model=model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding)

    def forward(
        self,
        input_ids,
        seq_len=None,
        attention_mask=None,
        top_k=4,
        top_p=0.0,
        max_length=256,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        forced_eos_token_id=None,
        temperature=0,
        decode_strategy="sample",
        num_return_sequences=1,
        **model_kwargs
    ):
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, "int32")

        # change top_p to zero if not using top_p sampling for FT
        if decode_strategy == "greedy_search":
            top_p = 0.0
            top_k = 1
        if top_p == 1.0:
            top_p = 0.0
        if seq_len is None:
            seq_len = paddle.sum(paddle.cast(input_ids != pad_token_id, dtype="int32"), axis=-1, dtype="int32")

            if bos_token_id == pad_token_id and paddle.sum(paddle.any(input_ids == pad_token_id), dtype="int64") > 0:
                seq_len = seq_len + 1

        if num_return_sequences > 1:
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids, expand_size=num_return_sequences, seq_len=seq_len, attention_mask=attention_mask
            )
            seq_len = model_kwargs["seq_len"]
            attention_mask = model_kwargs.get("attention_mask", None)

        return self.decoding(
            input_ids,
            mem_seq_len=seq_len,
            attention_mask=attention_mask,
            topk=top_k,
            topp=top_p,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            forced_eos_token_id=forced_eos_token_id,
            temperature=temperature,
        )

    def export_params(self, state_to_load, place):
        for item in state_to_load:
            param_data = np.array(state_to_load[item])
            if self.use_fp16_decoding:
                param_data = np.float16(param_data)

            param = self
            attr_list = item.split(".")
            attr_list = ["decoding", "model"] + attr_list
            for attr in attr_list:
                param = getattr(param, attr)
            param_name = param.name
            var = paddle.static.global_scope().find_var(param_name).get_tensor()
            var.set(param_data, place)

    def save_resources(self, tokenizer, path):
        vocab_file = os.path.join(path, "vocab.txt")
        if isinstance(tokenizer, GPTTokenizer):
            with open(vocab_file, "w", encoding="utf-8") as f:
                for token in tokenizer.encoder:
                    f.write(token + "\n")
            merges_file = os.path.join(path, "merges.txt")
            shutil.copyfile(tokenizer._merges_file, merges_file)
        elif isinstance(tokenizer, GPTChineseTokenizer):
            tokenizer.save_resources(path)

    generate = forward


class FasterGPT(GPTPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterGPT, self).__init__(model.config)
        self._model = model
        self.use_fp16_decoding = use_fp16_decoding
        self.decoding = InferGptDecoding(model=model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding)

    def forward(
        self,
        input_ids,
        seq_len=None,
        attention_mask=None,
        top_k=4,
        top_p=0.0,
        max_length=256,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        forced_eos_token_id=None,
        temperature=0,
        decode_strategy="sample",
        num_return_sequences=1,
        **model_kwargs
    ):
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, "int32")

        # change top_p to zero if not using top_p sampling for FT
        if decode_strategy == "greedy_search":
            top_p = 0.0
            top_k = 1
        if top_p == 1.0:
            top_p = 0.0
        if seq_len is None:
            seq_len = paddle.sum(paddle.cast(input_ids != pad_token_id, dtype="int32"), axis=-1, dtype="int32")

            if bos_token_id == pad_token_id and paddle.sum(paddle.any(input_ids == pad_token_id), dtype="int64") > 0:
                seq_len = seq_len + 1

        if num_return_sequences > 1:
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids, expand_size=num_return_sequences, seq_len=seq_len, attention_mask=attention_mask
            )
            seq_len = model_kwargs["seq_len"]
            attention_mask = model_kwargs.get("attention_mask", None)

        return self.decoding(
            input_ids,
            mem_seq_len=seq_len,
            attention_mask=attention_mask,
            topk=top_k,
            topp=top_p,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            forced_eos_token_id=forced_eos_token_id,
            temperature=temperature,
        )

    def export_params(self, state_to_load, place):
        for item in state_to_load:
            param_data = np.array(state_to_load[item])
            if self.use_fp16_decoding:
                param_data = np.float16(param_data)

            param = self
            attr_list = item.split(".")
            attr_list = ["decoding", "model"] + attr_list
            for attr in attr_list:
                param = getattr(param, attr)
            param_name = param.name
            var = paddle.static.global_scope().find_var(param_name).get_tensor()
            var.set(param_data, place)

    def save_resources(self, tokenizer, path):
        vocab_file = os.path.join(path, "vocab.txt")
        if isinstance(tokenizer, GPTTokenizer):
            with open(vocab_file, "w", encoding="utf-8") as f:
                for token in tokenizer.encoder:
                    f.write(token + "\n")
            merges_file = os.path.join(path, "merges.txt")
            shutil.copyfile(tokenizer._merges_file, merges_file)
        elif isinstance(tokenizer, GPTChineseTokenizer):
            tokenizer.save_resources(path)

    generate = forward


class FasterUnifiedTransformer(UnifiedTransformerPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterUnifiedTransformer, self).__init__(model.config)
        self._model = model
        self._use_fp16_decoding = use_fp16_decoding
        self.vocab_size = model.lm_head.decoder_bias.shape[0]
        self.unk_token_id = self._model.config.unk_token_id
        self.mask_token_id = self._model.config.mask_token_id
        self.bos_token_id = self._model.config.bos_token_id
        self.pad_token_id = self._model.config.pad_token_id
        self.logits_mask = self.generate_logits_mask(use_fp16_decoding)
        self._n_head = self._model.config.num_attention_heads
        self._hidden_dims = self._model.config.hidden_size
        self._normalize_before = self._model.config.normalize_before
        self._size_per_head = self._hidden_dims // self._n_head
        self._n_layer = self._model.config.num_hidden_layers
        self._hidden_act = self._model.config.hidden_act

        self.decoding = InferUnifiedDecoding(
            model=self._model,
            decoding_lib=decoding_lib,
            use_fp16_decoding=use_fp16_decoding,
            logits_mask=self.logits_mask,
            n_head=self._n_head,
            hidden_dims=self._hidden_dims,
            size_per_head=self._size_per_head,
            n_layer=self._n_layer,
            unk_id=self.unk_token_id,
            mask_id=self.mask_token_id,
            normalize_before=self._normalize_before,
            hidden_act=self._hidden_act,
        )

    def prepare_inputs_for_generation(
        self, input_ids, token_type_ids, attention_mask, seq_len, position_ids=None, role_ids=None, **kwargs
    ):
        input_ids = input_ids[:, :-1]
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, dtype="int32")

        if token_type_ids.dtype == paddle.int64:
            token_type_ids = paddle.cast(token_type_ids, dtype="int32")
        decoder_type_ids = token_type_ids[:, -1:]
        token_type_ids = token_type_ids[:, :-1]

        # TODO(guosheng): attention_mask of UnifiedTransformer uses 0/-INF
        # and is 4D. While now we want to use 1/0 to unify all models and
        # tokenizers.
        attention_mask = attention_mask[:, :, :-1, :-1] if attention_mask.ndim == 4 else attention_mask[:, :-1, :-1]
        attention_mask = paddle.cast(attention_mask == 0, dtype="float16" if self._use_fp16_decoding else "float32")

        seq_len = seq_len - 1
        if seq_len.dtype == paddle.int64:
            seq_len = paddle.cast(seq_len, dtype="int32")

        if position_ids is not None:
            if position_ids.dtype == paddle.int64:
                position_ids = paddle.cast(position_ids, dtype="int32")
            decoder_position_ids = position_ids[:, -1:]
            position_ids = position_ids[:, :-1]
        else:
            decoder_position_ids = None

        field_values = {}
        if role_ids is not None:
            if role_ids.dtype == paddle.int64:
                role_ids = paddle.cast(role_ids, dtype="int32")
            decoder_role_ids = role_ids[:, -1:]
            role_ids = role_ids[:, :-1]
        else:
            decoder_role_ids = None

        field_values["input_ids"] = input_ids
        field_values["token_type_ids"] = token_type_ids
        field_values["attention_mask"] = attention_mask
        field_values["seq_len"] = seq_len
        field_values["decoder_type_ids"] = decoder_type_ids
        field_values["position_ids"] = position_ids
        field_values["decoder_position_ids"] = decoder_position_ids
        field_values["role_ids"] = role_ids
        field_values["decoder_role_ids"] = decoder_role_ids

        return field_values

    def generate_logits_mask(self, use_fp16_decoding):
        # pre-process distribution
        logits_mask = np.zeros(shape=[self.vocab_size], dtype=np.float32)

        if use_fp16_decoding:
            logits_mask[self.unk_token_id] = -1e4
            logits_mask[self.bos_token_id] = -1e4
            logits_mask[self.pad_token_id] = -1e4
        else:
            logits_mask[self.unk_token_id] = -1e9
            logits_mask[self.bos_token_id] = -1e9
            logits_mask[self.pad_token_id] = -1e9

        logits_mask_t = paddle.assign(logits_mask)
        if use_fp16_decoding:
            return paddle.cast(logits_mask_t, dtype="float16")
        else:
            return logits_mask_t

    def forward(
        self,
        input_ids,
        token_type_ids,
        attention_mask,
        seq_len=None,
        role_ids=None,
        position_ids=None,
        max_length=128,
        min_length=0,
        top_k=4,
        top_p=0.0,
        decode_strategy="sampling",
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        num_beams=4,
        diversity_rate=0.0,
        temperature=1.0,
        num_return_sequences=1,
        length_penalty=0.6,
        early_stopping=False,
        forced_eos_token_id=None,
        **model_kwargs
    ):

        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(
                paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, keepdim=True, dtype="int32"
            )
        if decode_strategy.startswith("beam_search"):
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_beams,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
                role_ids=role_ids,
            )
        elif decode_strategy == "sampling":
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_return_sequences,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
                role_ids=role_ids,
            )
        elif decode_strategy == "greedy_search":
            model_kwargs = {
                "token_type_ids": token_type_ids,
                "position_ids": position_ids,
                "attention_mask": attention_mask,
                "seq_len": seq_len,
                "role_ids": role_ids,
            }
        else:
            raise ValueError("Only greedy search, beam search and sampling are supported. ")

        model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)

        seq_len = model_inputs.pop("seq_len")
        decoder_type_ids = model_inputs.pop("decoder_type_ids")
        role_ids = model_inputs.pop("role_ids", None)
        decoder_role_ids = model_inputs.pop("decoder_role_ids", None)
        position_ids = model_inputs.pop("position_ids", None)
        decoder_position_ids = model_inputs.pop("decoder_position_ids", None)

        return self.decoding(
            input_ids=model_inputs["input_ids"],
            attn_mask=model_inputs["attention_mask"],
            memory_seq_lens=seq_len,
            type_id=model_inputs["token_type_ids"],
            decoder_type_id=decoder_type_ids,
            role_id=role_ids,
            decoder_role_id=decoder_role_ids,
            position_id=position_ids,
            decoder_position_id=decoder_position_ids,
            beam_size=num_beams,
            diversity_rate=diversity_rate,
            topk=top_k,
            topp=top_p,
            decoding_strategy=decode_strategy,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            temperature=temperature,
            length_penalty=length_penalty,
            pos_bias=True,
            forced_eos_token_id=forced_eos_token_id,
            early_stopping=early_stopping,
            min_length=min_length,
        )

    generate = forward


class FasterUNIMOText(UNIMOPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False, **kwargs):
        super(FasterUNIMOText, self).__init__(model.config)
        self._model = model
        self._use_fp16_decoding = use_fp16_decoding
        self.unk_token_id = self._model.config.unk_token_id
        self.mask_token_id = self._model.config.mask_token_id
        self.bos_token_id = self._model.config.bos_token_id
        self.pad_token_id = self._model.config.pad_token_id
        self.vocab_size = model.lm_head.decoder_bias.shape[0]

        self.logits_mask = self.generate_logits_mask(use_fp16_decoding)
        self._n_head = self._model.config.num_attention_heads
        self._hidden_dims = self._model.config.hidden_size
        self._normalize_before = self._model.config.normalize_before
        self._size_per_head = self._hidden_dims // self._n_head
        self._n_layer = self._model.config.num_hidden_layers
        self._hidden_act = self._model.config.hidden_act
        self.trans_out = kwargs.get("trans_out", False)

        self.decoding = InferUnifiedDecoding(
            model=self._model,
            decoding_lib=decoding_lib,
            use_fp16_decoding=use_fp16_decoding,
            logits_mask=self.logits_mask,
            n_head=self._n_head,
            hidden_dims=self._hidden_dims,
            size_per_head=self._size_per_head,
            n_layer=self._n_layer,
            unk_id=self.unk_token_id,
            mask_id=self.mask_token_id,
            normalize_before=self._normalize_before,
            hidden_act=self._hidden_act,
        )

    def prepare_inputs_for_generation(self, input_ids, token_type_ids, attention_mask, **kwargs):
        input_ids = input_ids[:, :-1]
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, dtype="int32")

        if token_type_ids.dtype == paddle.int64:
            token_type_ids = paddle.cast(token_type_ids, dtype="int32")
        decoder_type_ids = token_type_ids[:, -1:]
        token_type_ids = token_type_ids[:, :-1]

        attention_mask = attention_mask[:, :, :-1, :-1]
        attention_mask = paddle.cast(attention_mask == 0, dtype="float16" if self._use_fp16_decoding else "float32")

        seq_len = kwargs.get("seq_len") - 1
        if seq_len.dtype == paddle.int64:
            seq_len = paddle.cast(seq_len, dtype="int32")

        return {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "attention_mask": attention_mask,
            "seq_len": seq_len,
            "decoder_type_ids": decoder_type_ids,
        }

    def generate_logits_mask(self, use_fp16_decoding):
        # pre-process distribution
        logits_mask = np.zeros(shape=[self.vocab_size], dtype=np.float32)

        if use_fp16_decoding:
            logits_mask[self.unk_token_id] = -1e4
            logits_mask[self.bos_token_id] = -1e4
            logits_mask[self.pad_token_id] = -1e4
        else:
            logits_mask[self.unk_token_id] = -1e9
            logits_mask[self.bos_token_id] = -1e9
            logits_mask[self.pad_token_id] = -1e9

        logits_mask_t = paddle.assign(logits_mask)
        if use_fp16_decoding:
            return paddle.cast(logits_mask_t, dtype="float16")
        else:
            return logits_mask_t

    def forward(
        self,
        input_ids,
        token_type_ids,
        attention_mask,
        seq_len=None,
        max_length=128,
        min_length=0,
        top_k=4,
        top_p=0.0,
        num_beams=4,
        decode_strategy="sampling",
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        diversity_rate=0.0,
        temperature=1.0,
        num_return_sequences=1,
        length_penalty=0.6,
        early_stopping=False,
        forced_eos_token_id=None,
        position_ids=None,
        **model_kwargs
    ):

        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, dtype="int32")
        if decode_strategy.startswith("beam_search"):
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_beams,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
            )
        elif decode_strategy == "sampling":
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_return_sequences,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
            )
        elif decode_strategy == "greedy_search":
            model_kwargs = {
                "token_type_ids": token_type_ids,
                "position_ids": position_ids,
                "attention_mask": attention_mask,
                "seq_len": seq_len,
            }
        else:
            raise ValueError("Only greedy search, beam search and sampling are supported. ")

        model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
        seq_len = model_inputs.pop("seq_len")
        decoder_type_ids = model_inputs.pop("decoder_type_ids")

        ids, output_scores = self.decoding(
            input_ids=model_inputs["input_ids"],
            attn_mask=model_inputs["attention_mask"],
            memory_seq_lens=seq_len,
            type_id=model_inputs["token_type_ids"],
            decoder_type_id=decoder_type_ids,
            beam_size=num_beams,
            diversity_rate=diversity_rate,
            topk=top_k,
            topp=top_p,
            decoding_strategy=decode_strategy,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            temperature=temperature,
            length_penalty=length_penalty,
            forced_eos_token_id=forced_eos_token_id,
            pos_bias=False,
            early_stopping=early_stopping,
            min_length=min_length,
        )
        if self.trans_out:
            if decode_strategy.startswith("beam_search"):
                ids = ids.transpose([1, 2, 0])
            else:
                ids = ids.transpose([1, 0])
        return ids, output_scores

    generate = forward


class FasterMIRO(UNIMOPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False, **kwargs):
        super(FasterMIRO, self).__init__(model.config)
        self._model = model
        self._use_fp16_decoding = use_fp16_decoding
        self.unk_token_id = self._model.config.unk_token_id
        self.mask_token_id = self._model.config.mask_token_id
        self.bos_token_id = self._model.config.bos_token_id
        self.pad_token_id = self._model.config.pad_token_id
        self.vocab_size = model.lm_head.decoder_bias.shape[0]

        self.logits_mask = self.generate_logits_mask(use_fp16_decoding)
        self._n_head = self._model.config.num_attention_heads
        self._hidden_dims = self._model.config.hidden_size
        self._normalize_before = self._model.config.normalize_before
        self._size_per_head = self._hidden_dims // self._n_head
        self._n_layer = self._model.config.num_hidden_layers
        self._hidden_act = self._model.config.hidden_act
        self.trans_out = kwargs.get("trans_out", False)

        self.decoding = InferMIRODecoding(
            model=self._model,
            decoding_lib=decoding_lib,
            use_fp16_decoding=use_fp16_decoding,
            logits_mask=self.logits_mask,
            n_head=self._n_head,
            hidden_dims=self._hidden_dims,
            size_per_head=self._size_per_head,
            n_layer=self._n_layer,
            unk_id=self.unk_token_id,
            mask_id=self.mask_token_id,
            normalize_before=self._normalize_before,
            hidden_act=self._hidden_act,
        )

    def prepare_inputs_for_generation(self, input_ids, token_type_ids, attention_mask, **kwargs):
        input_ids = input_ids[:, :-1]
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, dtype="int32")

        if token_type_ids.dtype == paddle.int64:
            token_type_ids = paddle.cast(token_type_ids, dtype="int32")
        decoder_type_ids = token_type_ids[:, -1:]
        token_type_ids = token_type_ids[:, :-1]

        attention_mask = attention_mask[:, :, :-1, :-1]
        attention_mask = paddle.cast(attention_mask == 0, dtype="float16" if self._use_fp16_decoding else "float32")

        seq_len = kwargs.get("seq_len") - 1
        if seq_len.dtype == paddle.int64:
            seq_len = paddle.cast(seq_len, dtype="int32")

        return {
            "input_ids": input_ids,
            "token_type_ids": token_type_ids,
            "attention_mask": attention_mask,
            "seq_len": seq_len,
            "decoder_type_ids": decoder_type_ids,
        }

    def generate_logits_mask(self, use_fp16_decoding):
        # pre-process distribution
        logits_mask = np.zeros(shape=[self.vocab_size], dtype=np.float32)

        if use_fp16_decoding:
            logits_mask[self.unk_token_id] = -1e4
            logits_mask[self.bos_token_id] = -1e4
            logits_mask[self.pad_token_id] = -1e4
        else:
            logits_mask[self.unk_token_id] = -1e9
            logits_mask[self.bos_token_id] = -1e9
            logits_mask[self.pad_token_id] = -1e9

        logits_mask_t = paddle.assign(logits_mask)
        if use_fp16_decoding:
            return paddle.cast(logits_mask_t, dtype="float16")
        else:
            return logits_mask_t

    def forward(
        self,
        input_ids,
        token_type_ids,
        attention_mask,
        seq_len=None,
        max_length=128,
        min_length=0,
        top_k=4,
        top_p=0.0,
        num_beams=4,
        decode_strategy="sampling",
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        diversity_rate=0.0,
        temperature=1.0,
        num_return_sequences=1,
        length_penalty=0.6,
        early_stopping=False,
        forced_eos_token_id=None,
        position_ids=None,
        **model_kwargs
    ):

        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, dtype="int32")
        if decode_strategy.startswith("beam_search"):
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_beams,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
            )
        elif decode_strategy == "sampling":
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids,
                expand_size=num_return_sequences,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                seq_len=seq_len,
            )
        elif decode_strategy == "greedy_search":
            model_kwargs = {
                "token_type_ids": token_type_ids,
                "position_ids": position_ids,
                "attention_mask": attention_mask,
                "seq_len": seq_len,
            }
        else:
            raise ValueError("Only greedy search, beam search and sampling are supported. ")

        model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
        seq_len = model_inputs.pop("seq_len")
        decoder_type_ids = model_inputs.pop("decoder_type_ids")

        ids, output_scores = self.decoding(
            input_ids=model_inputs["input_ids"],
            attn_mask=model_inputs["attention_mask"],
            memory_seq_lens=seq_len,
            type_id=model_inputs["token_type_ids"],
            decoder_type_id=decoder_type_ids,
            beam_size=num_beams,
            diversity_rate=diversity_rate,
            topk=top_k,
            topp=top_p,
            decoding_strategy=decode_strategy,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            temperature=temperature,
            length_penalty=length_penalty,
            forced_eos_token_id=forced_eos_token_id,
            pos_bias=False,
            early_stopping=early_stopping,
            min_length=min_length,
        )
        if self.trans_out:
            if decode_strategy.startswith("beam_search"):
                ids = ids.transpose([1, 2, 0])
            else:
                ids = ids.transpose([1, 0])
        return ids, output_scores

    generate = forward


class FasterBART(BartPretrainedModel):
    enable_faster_encoder_func = enable_fast_encoder

    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False, enable_fast_encoder=True):
        super(FasterBART, self).__init__(model.config)
        self.use_fp16_decoding = use_fp16_decoding
        self._model = model
        if use_fp16_decoding:
            weight_attr = paddle.ParamAttr(initializer=nn.initializer.Assign(model.bart.encoder.embed_tokens.weight))
            model.bart.encoder.embed_tokens = nn.Embedding(
                *model.bart.encoder.embed_tokens.weight.shape, weight_attr=weight_attr
            )
        self.encoder = model.bart.get_encoder()
        self.decoder = model.bart.get_decoder()
        self.pad_token_id = model.bart.config["pad_token_id"]
        self.enable_fast_encoder = enable_fast_encoder

        self.decoding = InferBartDecoding(
            model=self._model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding
        )
        if self.enable_fast_encoder:
            # Must use `enable_fast_encoder` in `__init__` when dygraph to static graph.
            self.encoder = FasterBART.enable_faster_encoder_func(self.encoder)

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def forward(
        self,
        input_ids=None,
        encoder_output=None,
        seq_len=None,
        num_beams=4,
        top_k=1,
        top_p=0.0,
        temperature=1.0,
        decode_strategy="beam_search",
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        decoder_start_token_id=None,
        min_length=0,
        max_length=20,
        diversity_rate=0.0,
        length_penalty=0.6,
        num_return_sequences=1,
        early_stopping=False,
        forced_eos_token_id=None,
        **model_kwargs
    ):

        if encoder_output is None:
            assert input_ids is not None, "You have to specify either input_ids or encoder_output."
            encoder_output = self.prepare_encoder_decoder_kwargs_for_generation(input_ids, model_kwargs)[
                "encoder_output"
            ]
        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(
                paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, keepdim=True, dtype="int32"
            )
        if self.use_fp16_decoding:
            encoder_output = paddle.cast(encoder_output, "float16")
        if decode_strategy.startswith("beam_search") and num_beams > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_beams, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        elif decode_strategy == "sampling" and num_return_sequences > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_return_sequences, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        if decoder_start_token_id is not None:
            bos_token_id = decoder_start_token_id

        return self.decoding(
            enc_output=encoder_output,
            memory_seq_lens=seq_len,
            beam_size=num_beams,
            top_k=top_k,
            decoding_strategy=decode_strategy,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            top_p=top_p,
            max_out_len=max_length,
            min_out_len=min_length,
            temperature=temperature,
            diversity_rate=diversity_rate,
            alpha=length_penalty,
            early_stopping=early_stopping,
            forced_eos_token_id=forced_eos_token_id,
        )

    generate = forward


class FasterMBART(MBartPretrainedModel):
    enable_faster_encoder_func = enable_fast_encoder

    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False, enable_fast_encoder=False):
        super(FasterMBART, self).__init__(model.config)
        self.use_fp16_decoding = use_fp16_decoding
        self._model = model
        if use_fp16_decoding:
            weight_attr = paddle.ParamAttr(initializer=nn.initializer.Assign(model.mbart.encoder.embed_tokens.weight))
            model.mbart.encoder.embed_tokens = nn.Embedding(
                *model.mbart.encoder.embed_tokens.weight.shape, weight_attr=weight_attr
            )
        self.encoder = model.mbart.get_encoder()
        self.decoder = model.mbart.get_decoder()
        self.pad_token_id = model.mbart.config["pad_token_id"]
        self.enable_fast_encoder = enable_fast_encoder

        self.decoding = InferMBartDecoding(
            model=self._model,
            decoding_lib=decoding_lib,
            use_fp16_decoding=use_fp16_decoding,
            hidden_act=model.mbart.config["activation_function"],
        )

        if self.enable_fast_encoder:
            # Must use `enable_fast_encoder` in `__init__` when dygraph to static graph.
            self.encoder = FasterMBART.enable_faster_encoder_func(self.encoder)

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def forward(
        self,
        input_ids=None,
        encoder_output=None,
        seq_len=None,
        forced_bos_token_id=None,
        num_beams=4,
        top_k=1,
        top_p=0.0,
        decode_strategy="beam_search_v3",
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        decoder_start_token_id=None,
        max_length=256,
        diversity_rate=0.0,
        length_penalty=0.6,
        temperature=1.0,
        num_return_sequences=1,
        early_stopping=False,
        forced_eos_token_id=None,
        **model_kwargs
    ):

        bos_token_id = bos_token_id if bos_token_id is not None else getattr(self._model, "bos_token_id", None)
        eos_token_id = eos_token_id if eos_token_id is not None else getattr(self._model, "eos_token_id", None)
        pad_token_id = pad_token_id if pad_token_id is not None else getattr(self._model, "pad_token_id", None)
        decoder_start_token_id = (
            decoder_start_token_id
            if decoder_start_token_id is not None
            else getattr(self._model, "decoder_start_token_id", None)
        )

        # (gongenlei) Not enable_fast_encoder temporarily
        if encoder_output is None:
            assert input_ids is not None, "You have to specify either input_ids or encoder_output."
            encoder_output = self.prepare_encoder_decoder_kwargs_for_generation(input_ids, model_kwargs)[
                "encoder_output"
            ]
        batch_size = paddle.shape(encoder_output)[0]
        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, dtype="int32")
        if self.use_fp16_decoding:
            encoder_output = paddle.cast(encoder_output, "float16")
        if decode_strategy.startswith("beam_search") and num_beams > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_beams, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        elif decode_strategy == "sampling" and num_return_sequences > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_return_sequences, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        if decoder_start_token_id is not None:
            bos_token_id = decoder_start_token_id

        if not isinstance(forced_bos_token_id, type(input_ids)):
            if forced_bos_token_id is not None:
                if decode_strategy == "sampling":
                    forced_bos_token_id = paddle.full(
                        [batch_size * num_return_sequences, 1], forced_bos_token_id, dtype="int32"
                    )
                else:
                    forced_bos_token_id = paddle.full([batch_size, 1], forced_bos_token_id, dtype="int32")
            else:
                forced_bos_token_id = paddle.zeros([0])
        elif decode_strategy == "sampling":
            num_samples = paddle.shape(encoder_output)[0]
            forced_bos_token_id = paddle.expand(forced_bos_token_id, shape=[num_samples, 1])

        return self.decoding(
            enc_output=encoder_output,
            memory_seq_lens=seq_len,
            beam_size=num_beams,
            trg_word=forced_bos_token_id,
            top_k=top_k,
            top_p=top_p,
            decoding_strategy=decode_strategy,
            diversity_rate=diversity_rate,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            alpha=length_penalty,
            temperature=temperature,
            early_stopping=early_stopping,
        )

    generate = forward


class FasterGPTJ(GPTJPretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterGPTJ, self).__init__(model.config)
        self._model = model
        self.use_fp16_decoding = use_fp16_decoding
        self.decoding = InferGptJDecoding(model=model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding)

    def forward(
        self,
        input_ids,
        seq_len=None,
        attention_mask=None,
        top_k=4,
        top_p=0.0,
        min_length=0,
        max_length=256,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        forced_eos_token_id=None,
        temperature=0,
        repetition_penalty=1.0,
        decode_strategy="sampling",
        num_return_sequences=1,
        **model_kwargs
    ):
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, "int32")

        # change top_p to zero if not using top_p sampling for FT
        if decode_strategy == "greedy_search":
            top_p = 0.0
            top_k = 1
        if top_p == 1.0:
            top_p = 0.0
        if seq_len is None:
            seq_len = paddle.sum(paddle.cast(input_ids != pad_token_id, dtype="int32"), axis=-1, dtype="int32")

        if num_return_sequences > 1:
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids, expand_size=num_return_sequences, seq_len=seq_len, attention_mask=attention_mask
            )
            seq_len = model_kwargs["seq_len"]
            attention_mask = model_kwargs.get("attention_mask", None)

        return self.decoding(
            input_ids,
            mem_seq_len=seq_len,
            attention_mask=attention_mask,
            topk=top_k,
            topp=top_p,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            forced_eos_token_id=forced_eos_token_id,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            min_length=min_length,
        )

    generate = forward


class FasterCodeGen(CodeGenPreTrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterCodeGen, self).__init__(model.config)
        self._model = model
        self.use_fp16_decoding = use_fp16_decoding
        self.decoding = InferGptJDecoding(
            model=model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding, transpose_qkv=True
        )

    def forward(
        self,
        input_ids,
        seq_len=None,
        attention_mask=None,
        top_k=4,
        top_p=0.0,
        min_length=0,
        max_length=256,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        forced_eos_token_id=None,
        temperature=0,
        repetition_penalty=1.0,
        decode_strategy="sampling",
        num_return_sequences=1,
        **model_kwargs
    ):
        if input_ids.dtype == paddle.int64:
            input_ids = paddle.cast(input_ids, "int32")

        # change top_p to zero if not using top_p sampling for FT
        if decode_strategy == "greedy_search":
            top_p = 0.0
            top_k = 1
        if top_p == 1.0:
            top_p = 0.0
        if seq_len is None:
            seq_len = paddle.sum(paddle.cast(input_ids != pad_token_id, dtype="int32"), axis=-1, dtype="int32")

        if num_return_sequences > 1:
            input_ids, model_kwargs = self.expand_inputs_for_generation(
                input_ids, expand_size=num_return_sequences, seq_len=seq_len, attention_mask=attention_mask
            )
            seq_len = model_kwargs["seq_len"]
            attention_mask = model_kwargs.get("attention_mask", None)

        return self.decoding(
            input_ids,
            mem_seq_len=seq_len,
            attention_mask=attention_mask,
            topk=top_k,
            topp=top_p,
            max_out_len=max_length,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            forced_eos_token_id=forced_eos_token_id,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            min_length=min_length,
        )

    generate = forward


class FasterPegasus(PegasusPretrainedModel):
    enable_faster_encoder_func = enable_fast_encoder

    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False, enable_fast_encoder=False, **kwargs):
        super(FasterPegasus, self).__init__(model.config)
        self.use_fp16_decoding = use_fp16_decoding
        self._model = model
        self.encoder = model.get_encoder()
        self.decoder = model.get_decoder()
        self.pad_token_id = model.pegasus.config["pad_token_id"]
        self.enable_fast_encoder = enable_fast_encoder
        self.trans_out = kwargs.get("trans_out", False)

        self.decoding = InferPegasusDecoding(
            model=self._model,
            decoding_lib=decoding_lib,
            use_fp16_decoding=use_fp16_decoding,
            hidden_act=model.pegasus.config["activation_function"],
        )

        # TODO(gongenlei): Support faster_encoder
        # if self.enable_fast_encoder:
        #     # Must use `enable_fast_encoder` in `__init__` when dygraph to static graph.
        #     self.encoder = FasterPegasus.enable_faster_encoder_func(self.encoder)

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def forward(
        self,
        input_ids=None,
        encoder_output=None,
        seq_len=None,
        min_length=0,
        max_length=256,
        num_beams=4,
        decode_strategy="beam_search_v3",
        decoder_start_token_id=None,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        diversity_rate=0.0,
        length_penalty=0.6,
        top_k=1,
        top_p=0.0,
        temperature=1.0,
        num_return_sequences=1,
        early_stopping=False,
        forced_bos_token_id=None,
        forced_eos_token_id=None,
        **model_kwargs
    ):

        bos_token_id = bos_token_id if bos_token_id is not None else getattr(self._model, "bos_token_id", None)
        eos_token_id = eos_token_id if eos_token_id is not None else getattr(self._model, "eos_token_id", None)
        pad_token_id = pad_token_id if pad_token_id is not None else getattr(self._model, "pad_token_id", None)
        decoder_start_token_id = (
            decoder_start_token_id
            if decoder_start_token_id is not None
            else getattr(self._model, "decoder_start_token_id", None)
        )

        if encoder_output is None:
            assert input_ids is not None, "You have to specify either input_ids or encoder_output."
            encoder_output = self.prepare_encoder_decoder_kwargs_for_generation(input_ids, model_kwargs)[
                "encoder_output"
            ]

        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, dtype="int32")
        if self.use_fp16_decoding:
            encoder_output = paddle.cast(encoder_output, "float16")
        if decode_strategy.startswith("beam_search") and num_beams > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_beams, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        elif decode_strategy == "sampling" and num_return_sequences > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_return_sequences, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        if decoder_start_token_id is not None:
            bos_token_id = decoder_start_token_id

        ids = self.decoding(
            enc_output=encoder_output,
            memory_seq_lens=seq_len,
            beam_size=num_beams,
            top_k=top_k,
            top_p=top_p,
            decoding_strategy=decode_strategy,
            max_out_len=max_length,
            min_out_len=min_length,
            diversity_rate=diversity_rate,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            alpha=length_penalty,
            temperature=temperature,
            early_stopping=early_stopping,
            forced_eos_token_id=forced_eos_token_id,
        )

        if self.trans_out:
            if decode_strategy.startswith("beam_search"):
                ids = ids.transpose([1, 2, 0])
            else:
                ids = ids.transpose([1, 0])

        return ids

    generate = forward


class FasterT5(T5PretrainedModel):
    def __init__(self, model, decoding_lib=None, use_fp16_decoding=False):
        super(FasterT5, self).__init__(model.config)
        self.use_fp16_decoding = use_fp16_decoding
        self._model = model
        if use_fp16_decoding:
            weight_attr = paddle.ParamAttr(initializer=nn.initializer.Assign(model.encoder.embed_tokens.weight))
            model.encoder.embed_tokens = nn.Embedding(
                *model.encoder.embed_tokens.weight.shape, weight_attr=weight_attr
            )
        self.encoder = model.t5.get_encoder()
        self.decoder = model.t5.get_decoder()
        self.pad_token_id = model.t5.config["pad_token_id"]

        self.decoding = InferT5Decoding(
            model=self._model, decoding_lib=decoding_lib, use_fp16_decoding=use_fp16_decoding
        )

    def get_encoder(self):
        return self.encoder

    def get_decoder(self):
        return self.decoder

    def forward(
        self,
        input_ids=None,
        encoder_output=None,
        seq_len=None,
        max_length=128,
        min_length=0,
        top_k=4,
        top_p=0.0,
        num_beams=4,
        decode_strategy="sampling",
        decoder_start_token_id=None,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=None,
        diversity_rate=0.0,
        temperature=1.0,
        num_return_sequences=1,
        length_penalty=0.6,
        early_stopping=False,
        forced_eos_token_id=None,
        **model_kwargs
    ):

        bos_token_id = bos_token_id if bos_token_id is not None else getattr(self._model, "bos_token_id", None)
        eos_token_id = eos_token_id if eos_token_id is not None else getattr(self._model, "eos_token_id", None)
        pad_token_id = pad_token_id if pad_token_id is not None else getattr(self._model, "pad_token_id", None)

        if encoder_output is None:
            assert input_ids is not None, "You have to specify either input_ids or encoder_output."
            encoder_output = self.prepare_encoder_decoder_kwargs_for_generation(input_ids, model_kwargs)[
                "encoder_output"
            ]

            if isinstance(encoder_output, (list, tuple)):
                encoder_output = encoder_output[0]

        if seq_len is None:
            assert input_ids is not None, "You have to specify either input_ids when generating seq_len."
            seq_len = paddle.sum(paddle.cast(input_ids != self.pad_token_id, dtype="int32"), axis=-1, dtype="int32")
        if self.use_fp16_decoding:
            encoder_output = paddle.cast(encoder_output, "float16")
        if decode_strategy.startswith("beam_search") and num_beams > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_beams, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        elif decode_strategy == "sampling" and num_return_sequences > 1:
            encoder_output, expanded_kwargs = self.expand_inputs_for_generation(
                encoder_output, expand_size=num_return_sequences, seq_len=seq_len
            )
            seq_len = expanded_kwargs["seq_len"]
        if decoder_start_token_id is not None:
            bos_token_id = decoder_start_token_id

        return self.decoding(
            enc_output=encoder_output,
            memory_seq_lens=seq_len,
            beam_size=num_beams,
            top_k=top_k,
            top_p=top_p,
            decoding_strategy=decode_strategy,
            max_out_len=max_length,
            diversity_rate=diversity_rate,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            alpha=length_penalty,
            temperature=temperature,
            early_stopping=early_stopping,
        )

    generate = forward
