# Copyright (c) 2022 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 numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.utils import map_structure

__all__ = [
    "position_encoding_init",
    "WordEmbedding",
    "PositionalEmbedding",
    "CrossEntropyCriterion",
    "TransformerDecodeCell",
    "TransformerBeamSearchDecoder",
    "TransformerModel",
]


def position_encoding_init(n_position, d_pos_vec, dtype="float64"):
    """
    Generates the initial values for the sinusoidal position encoding table.
    This method follows the implementation in tensor2tensor, but is slightly
    different from the description in "Attention Is All You Need".

    Args:
        n_position (int):
            The largest position for sequences, that is, the maximum length
            of source or target sequences.
        d_pos_vec (int):
            The size of positional embedding vector.
        dtype (str, optional):
            The output `numpy.array`'s data type. Defaults to "float32".

    Returns:
        numpy.array:
            The embedding table of sinusoidal position encoding with shape
            `[n_position, d_pos_vec]`.

    Example:
        .. code-block::

            from paddlenlp.transformers import position_encoding_init

            max_length = 256
            emb_dim = 512
            pos_table = position_encoding_init(max_length, emb_dim)
    """
    channels = d_pos_vec
    position = np.arange(n_position)
    num_timescales = channels // 2
    log_timescale_increment = np.log(float(1e4) / float(1)) / (num_timescales - 1)
    inv_timescales = np.exp(np.arange(num_timescales) * -log_timescale_increment)
    scaled_time = np.expand_dims(position, 1) * np.expand_dims(inv_timescales, 0)
    signal = np.concatenate([np.sin(scaled_time), np.cos(scaled_time)], axis=1)
    signal = np.pad(signal, [[0, 0], [0, np.mod(channels, 2)]], "constant")
    position_enc = signal
    return position_enc.astype(dtype)


class WordEmbedding(nn.Layer):
    """
    Word Embedding layer of Transformer.

    This layer automatically constructs a 2D embedding matrix based on the
    input the size of vocabulary (`vocab_size`) and the size of each embedding
    vector (`emb_dim`). This layer lookups embeddings vector of ids provided
    by input `word`.

    After the embedding, those weights are multiplied by `sqrt(d_model)` which is
    `sqrt(emb_dim)` in the interface.

    .. math::

        Out = embedding(word) * sqrt(emb\_dim)

    Args:
        vocab_size (int):
            The size of vocabulary.
        emb_dim (int):
            Dimensionality of each embedding vector.
        bos_id (int, optional):
            The start token id and also is used as padding id. Defaults to 0.
    """

    def __init__(self, vocab_size, emb_dim, bos_id=0):
        super(WordEmbedding, self).__init__()
        self.emb_dim = emb_dim

        self.word_embedding = nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=emb_dim,
            padding_idx=bos_id,
            weight_attr=paddle.ParamAttr(initializer=nn.initializer.Normal(0.0, emb_dim**-0.5)),
        )

    def forward(self, word):
        r"""
        Computes word embedding.

        Args:
            word (Tensor):
                The input ids which indicates the sequences' words with shape
                `[batch_size, sequence_length]` whose data type can be
                int or int64.

        Returns:
            Tensor:
                The (scaled) embedding tensor of shape
                `(batch_size, sequence_length, emb_dim)` whose data type can be
                float32 or float64.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import WordEmbedding

                word_embedding = WordEmbedding(
                    vocab_size=30000,
                    emb_dim=512,
                    bos_id=0)

                batch_size = 5
                sequence_length = 10
                src_words = paddle.randint(low=3, high=30000, shape=[batch_size, sequence_length])
                src_emb = word_embedding(src_words)
        """
        word_emb = self.emb_dim**0.5 * self.word_embedding(word)
        return word_emb


class PositionalEmbedding(nn.Layer):
    """
    This layer produces sinusoidal positional embeddings of any length.
    While in `forward()` method, this layer lookups embeddings vector of
    ids provided by input `pos`.

    Args:
        emb_dim (int):
            The size of each embedding vector.
        max_length (int):
            The maximum length of sequences.
    """

    def __init__(self, emb_dim, max_length):
        super(PositionalEmbedding, self).__init__()
        self.emb_dim = emb_dim

        self.pos_encoder = nn.Embedding(num_embeddings=max_length, embedding_dim=self.emb_dim)
        self.pos_encoder.weight.set_value(
            position_encoding_init(max_length, self.emb_dim, dtype=paddle.get_default_dtype())
        )

    def forward(self, pos):
        r"""
        Computes positional embedding.

        Args:
            pos (Tensor):
                The input position ids with shape `[batch_size, sequence_length]` whose
                data type can be int or int64.

        Returns:
            Tensor:
                The positional embedding tensor of shape
                `(batch_size, sequence_length, emb_dim)` whose data type can be
                float32 or float64.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import PositionalEmbedding

                pos_embedding = PositionalEmbedding(
                    emb_dim=512,
                    max_length=256)

                batch_size = 5
                pos = paddle.tile(paddle.arange(start=0, end=50), repeat_times=[batch_size, 1])
                pos_emb = pos_embedding(pos)
        """
        pos_emb = self.pos_encoder(pos)
        pos_emb.stop_gradient = True
        return pos_emb


class CrossEntropyCriterion(nn.Layer):
    """
    Computes the cross entropy loss for given input with or without label smoothing.

    Args:
        label_smooth_eps (float, optional):
            The weight used to mix up the original ground-truth distribution
            and the fixed distribution. Defaults to None. If given, label smoothing
            will be applied on `label`.
        pad_idx (int, optional):
            The token id used to pad variant sequence. Defaults to 0.
    """

    def __init__(self, label_smooth_eps=None, pad_idx=0):
        super(CrossEntropyCriterion, self).__init__()
        self.label_smooth_eps = label_smooth_eps
        self.pad_idx = pad_idx

    def forward(self, predict, label):
        r"""
        Computes cross entropy loss with or without label smoothing.

        Args:
            predict (Tensor):
                The predict results of `TransformerModel` with shape
                `[batch_size, sequence_length, vocab_size]` whose data type can
                be float32 or float64.
            label (Tensor):
                The label for corresponding results with shape
                `[batch_size, sequence_length, 1]`.

        Returns:
            tuple:
                A tuple with items: (`sum_cost`, `avg_cost`, `token_num`).

                With the corresponding fields:

                - `sum_cost` (Tensor):
                    The sum of loss of current batch whose data type can be float32, float64.
                - `avg_cost` (Tensor):
                    The average loss of current batch whose data type can be float32, float64.
                    The relation between `sum_cost` and `avg_cost` can be described as:

                    .. math:

                        avg_cost = sum_cost / token_num

                - `token_num` (Tensor):
                    The number of tokens of current batch.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import CrossEntropyCriterion

                criterion = CrossEntropyCriterion(label_smooth_eps=0.1, pad_idx=0)
                batch_size = 1
                seq_len = 2
                vocab_size = 30000
                predict = paddle.rand(shape=[batch_size, seq_len, vocab_size])
                label = paddle.randint(
                    low=3,
                    high=vocab_size,
                    shape=[batch_size, seq_len, 1])

                criterion(predict, label)
        """
        weights = paddle.cast(label != self.pad_idx, dtype=paddle.get_default_dtype())
        if self.label_smooth_eps:
            label = paddle.squeeze(label, axis=[2])
            label = F.label_smooth(
                label=F.one_hot(x=label, num_classes=predict.shape[-1]), epsilon=self.label_smooth_eps
            )
            if paddle.get_default_dtype() != "float32":
                label = paddle.cast(label, dtype=paddle.get_default_dtype())

        cost = F.cross_entropy(
            input=predict, label=label, reduction="none", soft_label=True if self.label_smooth_eps else False
        )
        weighted_cost = cost * weights
        sum_cost = paddle.sum(weighted_cost)
        token_num = paddle.sum(weights)
        token_num.stop_gradient = True
        avg_cost = sum_cost / token_num
        return sum_cost, avg_cost, token_num


class TransformerDecodeCell(nn.Layer):
    """
    This layer wraps a Transformer decoder combined with embedding
    layer and output layer to produce logits from ids and position.

    Args:
        decoder (callable):
            Can be a `paddle.nn.TransformerDecoder` instance. Or a wrapper that includes an
            embedding layer accepting ids and positions and includes an
            output layer transforming decoder output to logits.
        word_embedding (callable, optional):
            Can be a `WordEmbedding` instance or a callable that accepts ids as
            arguments and return embeddings. It can be None if `decoder`
            includes a embedding layer. Defaults to None.
        pos_embedding (callable, optional):
            Can be a `PositionalEmbedding` instance or a callable that accepts position
            as arguments and return embeddings. It can be None if `decoder`
            includes a positional embedding layer. Defaults to None.
        linear (callable, optional):
            Can be a `paddle.nn.Linear` instance or a callable to transform decoder
            output to logits.
        dropout (float, optional):
            The dropout rate for the results of `word_embedding` and `pos_embedding`.
            Defaults to 0.1.
    """

    def __init__(self, decoder, word_embedding=None, pos_embedding=None, linear=None, dropout=0.1):
        super(TransformerDecodeCell, self).__init__()
        self.decoder = decoder
        self.word_embedding = word_embedding
        self.pos_embedding = pos_embedding
        self.linear = linear
        self.dropout = dropout

    def forward(self, inputs, states, static_cache, trg_src_attn_bias, memory, **kwargs):
        r"""
        Produces logits.

        Args:
            inputs (Tensor|tuple|list):
                A tuple/list includes target ids and positions. If `word_embedding` is None,
                then it should be a Tensor which means the input for decoder.
            states (list):
                It is a list and each element of the list is an instance
                of `paddle.nn.MultiheadAttention.Cache` for corresponding decoder
                layer. It can be produced by `paddle.nn.TransformerDecoder.gen_cache`.
            static_cache (list):
                It is a list and each element of the list is an instance of
                `paddle.nn.MultiheadAttention.StaticCache` for corresponding
                decoder layer. It can be produced by `paddle.nn.TransformerDecoder.gen_cache`.
            trg_src_attn_bias (Tensor):
                A tensor used in self attention to prevents attention to some unwanted
                positions, usually the subsequent positions. It is a tensor with shape
                broadcasted to `[batch_size, n_head, target_length, target_length]`,
                where the unwanted positions have `-INF` values and the others
                have 0 values. The data type should be float32 or float64. It can
                be None when nothing wanted or needed to be prevented attention to.
            memory (Tensor):
                The output of Transformer encoder. It is a tensor with shape
                `[batch_size, source_length, d_model]` and its data type can be
                float32 or float64.

        Returns:
            tuple:
                A tuple with items: `(outputs, new_states)`

                With the corresponding fields:

                - `outputs` (Tensor):
                    A float32 or float64 3D tensor representing logits shaped
                    `[batch_size, sequence_length, vocab_size]`.
                - `new_states` (Tensor):
                    This output has the same structure and data type with `states`
                    while the length is one larger since concatenating the
                    intermediate results of current step.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import TransformerDecodeCell
                from paddlenlp.transformers import TransformerBeamSearchDecoder

                def decoder():
                    # do decoder
                    pass

                cell = TransformerDecodeCell(decoder())

                self.decode = TransformerBeamSearchDecoder(
                    cell, start_token=0, end_token=1, beam_size=4,
                    var_dim_in_state=2)
        """

        if states and static_cache:
            states = list(zip(states, static_cache))

        if self.word_embedding:
            if not isinstance(inputs, (list, tuple)):
                inputs = inputs

            word_emb = self.word_embedding(inputs[0])
            pos_emb = self.pos_embedding(inputs[1])
            word_emb = word_emb + pos_emb
            inputs = F.dropout(word_emb, p=self.dropout, training=False) if self.dropout else word_emb

            cell_outputs, new_states = self.decoder(inputs, memory, None, trg_src_attn_bias, states)
        else:
            cell_outputs, new_states = self.decoder(inputs, memory, None, trg_src_attn_bias, states)

        if self.linear:
            cell_outputs = self.linear(cell_outputs)

        new_states = [cache[0] for cache in new_states]

        return cell_outputs, new_states


class TransformerBeamSearchDecoder(nn.decode.BeamSearchDecoder):
    """
    This layer is a subclass of `BeamSearchDecoder` to make
    beam search adapt to Transformer decoder.

    Args:
        cell (`TransformerDecodeCell`):
            An instance of `TransformerDecoderCell`.
        start_token (int):
            The start token id.
        end_token (int):
            The end token id.
        beam_size (int):
            The beam width used in beam search.
        var_dim_in_state (int):
            Indicate which dimension of states is variant.
    """

    def __init__(self, cell, start_token, end_token, beam_size, var_dim_in_state):
        super(TransformerBeamSearchDecoder, self).__init__(cell, start_token, end_token, beam_size)
        self.cell = cell
        self.var_dim_in_state = var_dim_in_state

    def _merge_batch_beams_with_var_dim(self, c):
        # Init length of cache is 0, and it increases with decoding carrying on,
        # thus need to reshape elaborately
        var_dim_in_state = self.var_dim_in_state + 1  # count in beam dim
        c = paddle.transpose(c, list(range(var_dim_in_state, len(c.shape))) + list(range(0, var_dim_in_state)))
        c = paddle.reshape(
            c,
            [0] * (len(c.shape) - var_dim_in_state)
            + [self.batch_size * self.beam_size]
            + [int(size) for size in c.shape[-var_dim_in_state + 2 :]],
        )
        c = paddle.transpose(
            c,
            list(range((len(c.shape) + 1 - var_dim_in_state), len(c.shape)))
            + list(range(0, (len(c.shape) + 1 - var_dim_in_state))),
        )
        return c

    def _split_batch_beams_with_var_dim(self, c):
        var_dim_size = c.shape[self.var_dim_in_state]
        c = paddle.reshape(
            c,
            [-1, self.beam_size]
            + [int(size) for size in c.shape[1 : self.var_dim_in_state]]
            + [var_dim_size]
            + [int(size) for size in c.shape[self.var_dim_in_state + 1 :]],
        )
        return c

    @staticmethod
    def tile_beam_merge_with_batch(t, beam_size):
        r"""
        Tiles the batch dimension of a tensor. Specifically, this function takes
        a tensor t shaped `[batch_size, s0, s1, ...]` composed of minibatch
        entries `t[0], ..., t[batch_size - 1]` and tiles it to have a shape
        `[batch_size * beam_size, s0, s1, ...]` composed of minibatch entries
        `t[0], t[0], ..., t[1], t[1], ...` where each minibatch entry is repeated
        `beam_size` times.

        Args:
            t (list|tuple):
                A list of tensor with shape `[batch_size, ...]`.
            beam_size (int):
                The beam width used in beam search.

        Returns:
            Tensor:
                A tensor with shape `[batch_size * beam_size, ...]`, whose
                data type is same as `t`.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import TransformerBeamSearchDecoder

                t = paddle.rand(shape=[10, 10])
                TransformerBeamSearchDecoder.tile_beam_merge_with_batch(t, beam_size=4)
        """
        return map_structure(lambda x: nn.decode.BeamSearchDecoder.tile_beam_merge_with_batch(x, beam_size), t)

    def step(self, time, inputs, states, **kwargs):
        # Steps for decoding.
        # Compared to RNN, Transformer has 3D data at every decoding step
        inputs = paddle.reshape(inputs, [-1, 1])  # token
        pos = paddle.ones_like(inputs) * time  # pos

        cell_states = map_structure(self._merge_batch_beams_with_var_dim, states.cell_states)

        cell_outputs, next_cell_states = self.cell((inputs, pos), cell_states, **kwargs)

        # Squeeze to adapt to BeamSearchDecoder which use 2D logits
        cell_outputs = map_structure(lambda x: paddle.squeeze(x, [1]) if len(x.shape) == 3 else x, cell_outputs)
        cell_outputs = map_structure(self._split_batch_beams, cell_outputs)
        next_cell_states = map_structure(self._split_batch_beams_with_var_dim, next_cell_states)

        beam_search_output, beam_search_state = self._beam_search_step(
            time=time, logits=cell_outputs, next_cell_states=next_cell_states, beam_state=states
        )

        if kwargs.get("trg_word", None) is not None:
            if paddle.in_dynamic_mode():
                if kwargs.get("trg_word").shape[1] > time:
                    beam_search_output, beam_search_state = self.force_decoding(
                        beam_search_output, beam_search_state, kwargs.get("trg_word"), kwargs.get("trg_length"), time
                    )
            else:

                def condition(trg_word, time):
                    return trg_word.shape[1] > time

                def default_fn(beam_search_output, beam_search_state):
                    return beam_search_output, beam_search_state

                from functools import partial

                beam_search_output, beam_search_state = paddle.static.nn.case(
                    [
                        (
                            condition(kwargs.get("trg_word"), time),
                            partial(
                                self.force_decoding,
                                beam_search_output=beam_search_output,
                                beam_search_state=beam_search_state,
                                trg_word=kwargs.get("trg_word"),
                                trg_length=kwargs.get("trg_length"),
                                time=time,
                            ),
                        )
                    ],
                    default=partial(
                        default_fn, beam_search_output=beam_search_output, beam_search_state=beam_search_state
                    ),
                )

        next_inputs, finished = (beam_search_output.predicted_ids, beam_search_state.finished)

        return (beam_search_output, beam_search_state, next_inputs, finished)

    def force_decoding(self, beam_search_output, beam_search_state, trg_word, trg_length, time):
        batch_size = beam_search_output.predicted_ids.shape[0]
        beam_size = beam_search_output.predicted_ids.shape[1]

        ids_dtype = beam_search_output.predicted_ids.dtype
        scores_dtype = beam_search_output.scores.dtype
        parent_ids = paddle.zeros(shape=[batch_size, 1], dtype=ids_dtype)
        scores = paddle.ones(shape=[batch_size, beam_size], dtype=scores_dtype) * -10e9
        scores = paddle.scatter(
            scores.flatten(),
            paddle.arange(0, batch_size * beam_size, step=beam_size, dtype=scores_dtype),
            paddle.zeros([batch_size]),
        ).reshape([batch_size, beam_size])

        force_position = paddle.unsqueeze(trg_length > time, [1])
        # NOTE: When the date type of the input of paddle.tile is bool
        # and enable static mode, its stop_gradient must be True .
        force_position.stop_gradient = True
        force_position = paddle.tile(force_position, [1, beam_size])
        crt_trg_word = paddle.slice(trg_word, axes=[1], starts=[time], ends=[time + 1])
        crt_trg_word = paddle.tile(crt_trg_word, [1, beam_size])

        predicted_ids = paddle.where(force_position, crt_trg_word, beam_search_output.predicted_ids)
        scores = paddle.where(force_position, scores, beam_search_output.scores)
        parent_ids = paddle.where(force_position, parent_ids, beam_search_output.parent_ids)

        cell_states = beam_search_state.cell_states
        log_probs = paddle.where(force_position, scores, beam_search_state.log_probs)
        finished = beam_search_state.finished
        lengths = beam_search_state.lengths

        return self.OutputWrapper(scores, predicted_ids, parent_ids), self.StateWrapper(
            cell_states, log_probs, finished, lengths
        )


class TransformerModel(nn.Layer):
    """
    The Transformer model.

    This model is a Paddle `paddle.nn.Layer <https://www.paddlepaddle.org.cn/documentation
    /docs/zh/api/paddle/nn/Layer_cn.html>`__ subclass. Use it as a regular Paddle Layer
    and refer to the Paddle documentation for all matter related to general usage and behavior.

    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 activation. If None, use
            the value of dropout. Defaults to None.
        bos_id (int, optional):
            The start token id and also be 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.
        activation (str, optional):
            The activation used in FFN. Defaults to "relu".
    """

    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,
        activation="relu",
    ):
        super(TransformerModel, self).__init__()
        self.trg_vocab_size = trg_vocab_size
        self.emb_dim = d_model
        self.bos_id = bos_id
        self.eos_id = eos_id
        self.pad_id = pad_id if pad_id is not None else self.bos_id
        self.dropout = dropout

        self.src_word_embedding = WordEmbedding(vocab_size=src_vocab_size, emb_dim=d_model, bos_id=self.bos_id)
        self.src_pos_embedding = PositionalEmbedding(emb_dim=d_model, max_length=max_length)
        if weight_sharing:
            assert (
                src_vocab_size == trg_vocab_size
            ), "Vocabularies in source and target should be same for weight sharing."
            self.trg_word_embedding = self.src_word_embedding
            self.trg_pos_embedding = self.src_pos_embedding
        else:
            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.transformer = paddle.nn.Transformer(
            d_model=d_model,
            nhead=n_head,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dim_feedforward=d_inner_hid,
            dropout=dropout,
            attn_dropout=attn_dropout,
            act_dropout=act_dropout,
            activation=activation,
            normalize_before=True,
        )

        if weight_sharing:
            self.linear = lambda x: paddle.matmul(
                x=x, y=self.trg_word_embedding.word_embedding.weight, transpose_y=True
            )
        else:
            self.linear = nn.Linear(in_features=d_model, out_features=trg_vocab_size, bias_attr=False)

    def forward(self, src_word, trg_word):
        r"""
        The Transformer forward methods. The input are source/target sequences, and
        returns logits.

        Args:
            src_word (Tensor):
                The ids of source sequences 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 sequences words. It is a tensor with shape
                `[batch_size, target_sequence_length]` and its data type can be
                int or int64.

        Returns:
            Tensor:
                Output tensor of the final layer of the model whose data
                type can be float32 or float64 with shape
                `[batch_size, sequence_length, vocab_size]`.

        Example:
            .. code-block::

                import paddle
                from paddlenlp.transformers import TransformerModel

                transformer = TransformerModel(
                    src_vocab_size=30000,
                    trg_vocab_size=30000,
                    max_length=257,
                    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)

                batch_size = 5
                seq_len = 10
                predict = transformer(
                    src_word=paddle.randint(low=3, high=30000, shape=[batch_size, seq_len]),
                    trg_word=paddle.randint(low=3, high=30000, shape=[batch_size, seq_len]))
        """
        src_max_len = src_word.shape[-1]
        trg_max_len = trg_word.shape[-1]
        src_slf_attn_bias = (
            paddle.cast(src_word == self.bos_id, dtype=paddle.get_default_dtype()).unsqueeze([1, 2]) * -1e9
        )
        src_slf_attn_bias.stop_gradient = True
        trg_slf_attn_bias = self.transformer.generate_square_subsequent_mask(trg_max_len)
        trg_slf_attn_bias.stop_gradient = True
        trg_src_attn_bias = src_slf_attn_bias
        src_pos = paddle.cast(src_word != self.bos_id, dtype=src_word.dtype) * paddle.arange(
            start=0, end=src_max_len, dtype=src_word.dtype
        )
        trg_pos = paddle.cast(trg_word != self.bos_id, dtype=src_word.dtype) * paddle.arange(
            start=0, end=trg_max_len, dtype=trg_word.dtype
        )
        with paddle.static.amp.fp16_guard():
            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=self.training) if self.dropout else src_emb

            trg_emb = self.trg_word_embedding(trg_word)
            trg_pos_emb = self.trg_pos_embedding(trg_pos)
            trg_emb = trg_emb + trg_pos_emb
            dec_input = F.dropout(trg_emb, p=self.dropout, training=self.training) if self.dropout else trg_emb

            dec_output = self.transformer(
                enc_input,
                dec_input,
                src_mask=src_slf_attn_bias,
                tgt_mask=trg_slf_attn_bias,
                memory_mask=trg_src_attn_bias,
            )

            predict = self.linear(dec_output)

        return predict
