pyxllib 0.3.96__py3-none-any.whl → 0.3.197__py3-none-any.whl

pyxlpr/ppocr/modeling/heads/rec_nrtr_head.py

@@ -1,826 +0,0 @@
-# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
-#
-# 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 math
-import paddle
-import copy
-from paddle import nn
-import paddle.nn.functional as F
-from paddle.nn import LayerList
-from paddle.nn.initializer import XavierNormal as xavier_uniform_
-from paddle.nn import Dropout, Linear, LayerNorm, Conv2D
-import numpy as np
-from pyxlpr.ppocr.modeling.heads.multiheadAttention import MultiheadAttention
-from paddle.nn.initializer import Constant as constant_
-from paddle.nn.initializer import XavierNormal as xavier_normal_
-
-zeros_ = constant_(value=0.)
-ones_ = constant_(value=1.)
-
-
-class Transformer(nn.Layer):
-    """A transformer model. User is able to modify the attributes as needed. The architechture
-    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
-    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
-    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
-    Processing Systems, pages 6000-6010.
-
-    Args:
-        d_model: the number of expected features in the encoder/decoder inputs (default=512).
-        nhead: the number of heads in the multiheadattention models (default=8).
-        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
-        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-        custom_encoder: custom encoder (default=None).
-        custom_decoder: custom decoder (default=None).
-
-    """
-
-    def __init__(self,
-                 d_model=512,
-                 nhead=8,
-                 num_encoder_layers=6,
-                 beam_size=0,
-                 num_decoder_layers=6,
-                 dim_feedforward=1024,
-                 attention_dropout_rate=0.0,
-                 residual_dropout_rate=0.1,
-                 custom_encoder=None,
-                 custom_decoder=None,
-                 in_channels=0,
-                 out_channels=0,
-                 scale_embedding=True):
-        super(Transformer, self).__init__()
-        self.out_channels = out_channels + 1
-        self.embedding = Embeddings(
-            d_model=d_model,
-            vocab=self.out_channels,
-            padding_idx=0,
-            scale_embedding=scale_embedding)
-        self.positional_encoding = PositionalEncoding(
-            dropout=residual_dropout_rate,
-            dim=d_model, )
-        if custom_encoder is not None:
-            self.encoder = custom_encoder
-        else:
-            if num_encoder_layers > 0:
-                encoder_layer = TransformerEncoderLayer(
-                    d_model, nhead, dim_feedforward, attention_dropout_rate,
-                    residual_dropout_rate)
-                self.encoder = TransformerEncoder(encoder_layer,
-                                                  num_encoder_layers)
-            else:
-                self.encoder = None
-
-        if custom_decoder is not None:
-            self.decoder = custom_decoder
-        else:
-            decoder_layer = TransformerDecoderLayer(
-                d_model, nhead, dim_feedforward, attention_dropout_rate,
-                residual_dropout_rate)
-            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
-
-        self._reset_parameters()
-        self.beam_size = beam_size
-        self.d_model = d_model
-        self.nhead = nhead
-        self.tgt_word_prj = nn.Linear(
-            d_model, self.out_channels, bias_attr=False)
-        w0 = np.random.normal(0.0, d_model**-0.5,
-                              (d_model, self.out_channels)).astype(np.float32)
-        self.tgt_word_prj.weight.set_value(w0)
-        self.apply(self._init_weights)
-
-    def _init_weights(self, m):
-
-        if isinstance(m, nn.Conv2D):
-            xavier_normal_(m.weight)
-            if m.bias is not None:
-                zeros_(m.bias)
-
-    def forward_train(self, src, tgt):
-        tgt = tgt[:, :-1]
-
-        tgt_key_padding_mask = self.generate_padding_mask(tgt)
-        tgt = self.embedding(tgt).transpose([1, 0, 2])
-        tgt = self.positional_encoding(tgt)
-        tgt_mask = self.generate_square_subsequent_mask(tgt.shape[0])
-
-        if self.encoder is not None:
-            src = self.positional_encoding(src.transpose([1, 0, 2]))
-            memory = self.encoder(src)
-        else:
-            memory = src.squeeze(2).transpose([2, 0, 1])
-        output = self.decoder(
-            tgt,
-            memory,
-            tgt_mask=tgt_mask,
-            memory_mask=None,
-            tgt_key_padding_mask=tgt_key_padding_mask,
-            memory_key_padding_mask=None)
-        output = output.transpose([1, 0, 2])
-        logit = self.tgt_word_prj(output)
-        return logit
-
-    def forward(self, src, targets=None):
-        """Take in and process masked source/target sequences.
-        Args:
-            src: the sequence to the encoder (required).
-            tgt: the sequence to the decoder (required).
-        Shape:
-            - src: :math:`(S, N, E)`.
-            - tgt: :math:`(T, N, E)`.
-        Examples:
-            >>> output = transformer_model(src, tgt)
-        """
-
-        if self.training:
-            max_len = targets[1].max()
-            tgt = targets[0][:, :2 + max_len]
-            return self.forward_train(src, tgt)
-        else:
-            if self.beam_size > 0:
-                return self.forward_beam(src)
-            else:
-                return self.forward_test(src)
-
-    def forward_test(self, src):
-        bs = paddle.shape(src)[0]
-        if self.encoder is not None:
-            src = self.positional_encoding(paddle.transpose(src, [1, 0, 2]))
-            memory = self.encoder(src)
-        else:
-            memory = paddle.transpose(paddle.squeeze(src, 2), [2, 0, 1])
-        dec_seq = paddle.full((bs, 1), 2, dtype=paddle.int64)
-        dec_prob = paddle.full((bs, 1), 1., dtype=paddle.float32)
-        for len_dec_seq in range(1, 25):
-            dec_seq_embed = paddle.transpose(self.embedding(dec_seq), [1, 0, 2])
-            dec_seq_embed = self.positional_encoding(dec_seq_embed)
-            tgt_mask = self.generate_square_subsequent_mask(
-                paddle.shape(dec_seq_embed)[0])
-            output = self.decoder(
-                dec_seq_embed,
-                memory,
-                tgt_mask=tgt_mask,
-                memory_mask=None,
-                tgt_key_padding_mask=None,
-                memory_key_padding_mask=None)
-            dec_output = paddle.transpose(output, [1, 0, 2])
-            dec_output = dec_output[:, -1, :]
-            word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=1)
-            preds_idx = paddle.argmax(word_prob, axis=1)
-            if paddle.equal_all(
-                    preds_idx,
-                    paddle.full(
-                        paddle.shape(preds_idx), 3, dtype='int64')):
-                break
-            preds_prob = paddle.max(word_prob, axis=1)
-            dec_seq = paddle.concat(
-                [dec_seq, paddle.reshape(preds_idx, [-1, 1])], axis=1)
-            dec_prob = paddle.concat(
-                [dec_prob, paddle.reshape(preds_prob, [-1, 1])], axis=1)
-        return [dec_seq, dec_prob]
-
-    def forward_beam(self, images):
-        ''' Translation work in one batch '''
-
-        def get_inst_idx_to_tensor_position_map(inst_idx_list):
-            ''' Indicate the position of an instance in a tensor. '''
-            return {
-                inst_idx: tensor_position
-                for tensor_position, inst_idx in enumerate(inst_idx_list)
-            }
-
-        def collect_active_part(beamed_tensor, curr_active_inst_idx,
-                                n_prev_active_inst, n_bm):
-            ''' Collect tensor parts associated to active instances. '''
-
-            beamed_tensor_shape = paddle.shape(beamed_tensor)
-            n_curr_active_inst = len(curr_active_inst_idx)
-            new_shape = (n_curr_active_inst * n_bm, beamed_tensor_shape[1],
-                         beamed_tensor_shape[2])
-
-            beamed_tensor = beamed_tensor.reshape([n_prev_active_inst, -1])
-            beamed_tensor = beamed_tensor.index_select(
-                curr_active_inst_idx, axis=0)
-            beamed_tensor = beamed_tensor.reshape(new_shape)
-
-            return beamed_tensor
-
-        def collate_active_info(src_enc, inst_idx_to_position_map,
-                                active_inst_idx_list):
-            # Sentences which are still active are collected,
-            # so the decoder will not run on completed sentences.
-
-            n_prev_active_inst = len(inst_idx_to_position_map)
-            active_inst_idx = [
-                inst_idx_to_position_map[k] for k in active_inst_idx_list
-            ]
-            active_inst_idx = paddle.to_tensor(active_inst_idx, dtype='int64')
-            active_src_enc = collect_active_part(
-                src_enc.transpose([1, 0, 2]), active_inst_idx,
-                n_prev_active_inst, n_bm).transpose([1, 0, 2])
-            active_inst_idx_to_position_map = get_inst_idx_to_tensor_position_map(
-                active_inst_idx_list)
-            return active_src_enc, active_inst_idx_to_position_map
-
-        def beam_decode_step(inst_dec_beams, len_dec_seq, enc_output,
-                             inst_idx_to_position_map, n_bm,
-                             memory_key_padding_mask):
-            ''' Decode and update beam status, and then return active beam idx '''
-
-            def prepare_beam_dec_seq(inst_dec_beams, len_dec_seq):
-                dec_partial_seq = [
-                    b.get_current_state() for b in inst_dec_beams if not b.done
-                ]
-                dec_partial_seq = paddle.stack(dec_partial_seq)
-                dec_partial_seq = dec_partial_seq.reshape([-1, len_dec_seq])
-                return dec_partial_seq
-
-            def predict_word(dec_seq, enc_output, n_active_inst, n_bm,
-                             memory_key_padding_mask):
-                dec_seq = paddle.transpose(self.embedding(dec_seq), [1, 0, 2])
-                dec_seq = self.positional_encoding(dec_seq)
-                tgt_mask = self.generate_square_subsequent_mask(
-                    paddle.shape(dec_seq)[0])
-                dec_output = self.decoder(
-                    dec_seq,
-                    enc_output,
-                    tgt_mask=tgt_mask,
-                    tgt_key_padding_mask=None,
-                    memory_key_padding_mask=memory_key_padding_mask, )
-                dec_output = paddle.transpose(dec_output, [1, 0, 2])
-                dec_output = dec_output[:,
-                                        -1, :]  # Pick the last step: (bh * bm) * d_h
-                word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=1)
-                word_prob = paddle.reshape(word_prob, [n_active_inst, n_bm, -1])
-                return word_prob
-
-            def collect_active_inst_idx_list(inst_beams, word_prob,
-                                             inst_idx_to_position_map):
-                active_inst_idx_list = []
-                for inst_idx, inst_position in inst_idx_to_position_map.items():
-                    is_inst_complete = inst_beams[inst_idx].advance(word_prob[
-                        inst_position])
-                    if not is_inst_complete:
-                        active_inst_idx_list += [inst_idx]
-
-                return active_inst_idx_list
-
-            n_active_inst = len(inst_idx_to_position_map)
-            dec_seq = prepare_beam_dec_seq(inst_dec_beams, len_dec_seq)
-            word_prob = predict_word(dec_seq, enc_output, n_active_inst, n_bm,
-                                     None)
-            # Update the beam with predicted word prob information and collect incomplete instances
-            active_inst_idx_list = collect_active_inst_idx_list(
-                inst_dec_beams, word_prob, inst_idx_to_position_map)
-            return active_inst_idx_list
-
-        def collect_hypothesis_and_scores(inst_dec_beams, n_best):
-            all_hyp, all_scores = [], []
-            for inst_idx in range(len(inst_dec_beams)):
-                scores, tail_idxs = inst_dec_beams[inst_idx].sort_scores()
-                all_scores += [scores[:n_best]]
-                hyps = [
-                    inst_dec_beams[inst_idx].get_hypothesis(i)
-                    for i in tail_idxs[:n_best]
-                ]
-                all_hyp += [hyps]
-            return all_hyp, all_scores
-
-        with paddle.no_grad():
-            #-- Encode
-            if self.encoder is not None:
-                src = self.positional_encoding(images.transpose([1, 0, 2]))
-                src_enc = self.encoder(src)
-            else:
-                src_enc = images.squeeze(2).transpose([0, 2, 1])
-
-            n_bm = self.beam_size
-            src_shape = paddle.shape(src_enc)
-            inst_dec_beams = [Beam(n_bm) for _ in range(1)]
-            active_inst_idx_list = list(range(1))
-            # Repeat data for beam search
-            src_enc = paddle.tile(src_enc, [1, n_bm, 1])
-            inst_idx_to_position_map = get_inst_idx_to_tensor_position_map(
-                active_inst_idx_list)
-            # Decode
-            for len_dec_seq in range(1, 25):
-                src_enc_copy = src_enc.clone()
-                active_inst_idx_list = beam_decode_step(
-                    inst_dec_beams, len_dec_seq, src_enc_copy,
-                    inst_idx_to_position_map, n_bm, None)
-                if not active_inst_idx_list:
-                    break  # all instances have finished their path to <EOS>
-                src_enc, inst_idx_to_position_map = collate_active_info(
-                    src_enc_copy, inst_idx_to_position_map,
-                    active_inst_idx_list)
-        batch_hyp, batch_scores = collect_hypothesis_and_scores(inst_dec_beams,
-                                                                1)
-        result_hyp = []
-        hyp_scores = []
-        for bs_hyp, score in zip(batch_hyp, batch_scores):
-            l = len(bs_hyp[0])
-            bs_hyp_pad = bs_hyp[0] + [3] * (25 - l)
-            result_hyp.append(bs_hyp_pad)
-            score = float(score) / l
-            hyp_score = [score for _ in range(25)]
-            hyp_scores.append(hyp_score)
-        return [
-            paddle.to_tensor(
-                np.array(result_hyp), dtype=paddle.int64),
-            paddle.to_tensor(hyp_scores)
-        ]
-
-    def generate_square_subsequent_mask(self, sz):
-        """Generate a square mask for the sequence. The masked positions are filled with float('-inf').
-            Unmasked positions are filled with float(0.0).
-        """
-        mask = paddle.zeros([sz, sz], dtype='float32')
-        mask_inf = paddle.triu(
-            paddle.full(
-                shape=[sz, sz], dtype='float32', fill_value='-inf'),
-            diagonal=1)
-        mask = mask + mask_inf
-        return mask
-
-    def generate_padding_mask(self, x):
-        padding_mask = paddle.equal(x, paddle.to_tensor(0, dtype=x.dtype))
-        return padding_mask
-
-    def _reset_parameters(self):
-        """Initiate parameters in the transformer model."""
-
-        for p in self.parameters():
-            if p.dim() > 1:
-                xavier_uniform_(p)
-
-
-class TransformerEncoder(nn.Layer):
-    """TransformerEncoder is a stack of N encoder layers
-    Args:
-        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
-        num_layers: the number of sub-encoder-layers in the encoder (required).
-        norm: the layer normalization component (optional).
-    """
-
-    def __init__(self, encoder_layer, num_layers):
-        super(TransformerEncoder, self).__init__()
-        self.layers = _get_clones(encoder_layer, num_layers)
-        self.num_layers = num_layers
-
-    def forward(self, src):
-        """Pass the input through the endocder layers in turn.
-        Args:
-            src: the sequnce to the encoder (required).
-            mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-        """
-        output = src
-
-        for i in range(self.num_layers):
-            output = self.layers[i](output,
-                                    src_mask=None,
-                                    src_key_padding_mask=None)
-
-        return output
-
-
-class TransformerDecoder(nn.Layer):
-    """TransformerDecoder is a stack of N decoder layers
-
-    Args:
-        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
-        num_layers: the number of sub-decoder-layers in the decoder (required).
-        norm: the layer normalization component (optional).
-
-    """
-
-    def __init__(self, decoder_layer, num_layers):
-        super(TransformerDecoder, self).__init__()
-        self.layers = _get_clones(decoder_layer, num_layers)
-        self.num_layers = num_layers
-
-    def forward(self,
-                tgt,
-                memory,
-                tgt_mask=None,
-                memory_mask=None,
-                tgt_key_padding_mask=None,
-                memory_key_padding_mask=None):
-        """Pass the inputs (and mask) through the decoder layer in turn.
-
-        Args:
-            tgt: the sequence to the decoder (required).
-            memory: the sequnce from the last layer of the encoder (required).
-            tgt_mask: the mask for the tgt sequence (optional).
-            memory_mask: the mask for the memory sequence (optional).
-            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
-            memory_key_padding_mask: the mask for the memory keys per batch (optional).
-        """
-        output = tgt
-        for i in range(self.num_layers):
-            output = self.layers[i](
-                output,
-                memory,
-                tgt_mask=tgt_mask,
-                memory_mask=memory_mask,
-                tgt_key_padding_mask=tgt_key_padding_mask,
-                memory_key_padding_mask=memory_key_padding_mask)
-
-        return output
-
-
-class TransformerEncoderLayer(nn.Layer):
-    """TransformerEncoderLayer is made up of self-attn and feedforward network.
-    This standard encoder layer is based on the paper "Attention Is All You Need".
-    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
-    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
-    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
-    in a different way during application.
-
-    Args:
-        d_model: the number of expected features in the input (required).
-        nhead: the number of heads in the multiheadattention models (required).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-
-    """
-
-    def __init__(self,
-                 d_model,
-                 nhead,
-                 dim_feedforward=2048,
-                 attention_dropout_rate=0.0,
-                 residual_dropout_rate=0.1):
-        super(TransformerEncoderLayer, self).__init__()
-        self.self_attn = MultiheadAttention(
-            d_model, nhead, dropout=attention_dropout_rate)
-
-        self.conv1 = Conv2D(
-            in_channels=d_model,
-            out_channels=dim_feedforward,
-            kernel_size=(1, 1))
-        self.conv2 = Conv2D(
-            in_channels=dim_feedforward,
-            out_channels=d_model,
-            kernel_size=(1, 1))
-
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.dropout1 = Dropout(residual_dropout_rate)
-        self.dropout2 = Dropout(residual_dropout_rate)
-
-    def forward(self, src, src_mask=None, src_key_padding_mask=None):
-        """Pass the input through the endocder layer.
-        Args:
-            src: the sequnce to the encoder layer (required).
-            src_mask: the mask for the src sequence (optional).
-            src_key_padding_mask: the mask for the src keys per batch (optional).
-        """
-        src2 = self.self_attn(
-            src,
-            src,
-            src,
-            attn_mask=src_mask,
-            key_padding_mask=src_key_padding_mask)
-        src = src + self.dropout1(src2)
-        src = self.norm1(src)
-
-        src = paddle.transpose(src, [1, 2, 0])
-        src = paddle.unsqueeze(src, 2)
-        src2 = self.conv2(F.relu(self.conv1(src)))
-        src2 = paddle.squeeze(src2, 2)
-        src2 = paddle.transpose(src2, [2, 0, 1])
-        src = paddle.squeeze(src, 2)
-        src = paddle.transpose(src, [2, 0, 1])
-
-        src = src + self.dropout2(src2)
-        src = self.norm2(src)
-        return src
-
-
-class TransformerDecoderLayer(nn.Layer):
-    """TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
-    This standard decoder layer is based on the paper "Attention Is All You Need".
-    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
-    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
-    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
-    in a different way during application.
-
-    Args:
-        d_model: the number of expected features in the input (required).
-        nhead: the number of heads in the multiheadattention models (required).
-        dim_feedforward: the dimension of the feedforward network model (default=2048).
-        dropout: the dropout value (default=0.1).
-
-    """
-
-    def __init__(self,
-                 d_model,
-                 nhead,
-                 dim_feedforward=2048,
-                 attention_dropout_rate=0.0,
-                 residual_dropout_rate=0.1):
-        super(TransformerDecoderLayer, self).__init__()
-        self.self_attn = MultiheadAttention(
-            d_model, nhead, dropout=attention_dropout_rate)
-        self.multihead_attn = MultiheadAttention(
-            d_model, nhead, dropout=attention_dropout_rate)
-
-        self.conv1 = Conv2D(
-            in_channels=d_model,
-            out_channels=dim_feedforward,
-            kernel_size=(1, 1))
-        self.conv2 = Conv2D(
-            in_channels=dim_feedforward,
-            out_channels=d_model,
-            kernel_size=(1, 1))
-
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        self.norm3 = LayerNorm(d_model)
-        self.dropout1 = Dropout(residual_dropout_rate)
-        self.dropout2 = Dropout(residual_dropout_rate)
-        self.dropout3 = Dropout(residual_dropout_rate)
-
-    def forward(self,
-                tgt,
-                memory,
-                tgt_mask=None,
-                memory_mask=None,
-                tgt_key_padding_mask=None,
-                memory_key_padding_mask=None):
-        """Pass the inputs (and mask) through the decoder layer.
-
-        Args:
-            tgt: the sequence to the decoder layer (required).
-            memory: the sequnce from the last layer of the encoder (required).
-            tgt_mask: the mask for the tgt sequence (optional).
-            memory_mask: the mask for the memory sequence (optional).
-            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
-            memory_key_padding_mask: the mask for the memory keys per batch (optional).
-
-        """
-        tgt2 = self.self_attn(
-            tgt,
-            tgt,
-            tgt,
-            attn_mask=tgt_mask,
-            key_padding_mask=tgt_key_padding_mask)
-        tgt = tgt + self.dropout1(tgt2)
-        tgt = self.norm1(tgt)
-        tgt2 = self.multihead_attn(
-            tgt,
-            memory,
-            memory,
-            attn_mask=memory_mask,
-            key_padding_mask=memory_key_padding_mask)
-        tgt = tgt + self.dropout2(tgt2)
-        tgt = self.norm2(tgt)
-
-        # default
-        tgt = paddle.transpose(tgt, [1, 2, 0])
-        tgt = paddle.unsqueeze(tgt, 2)
-        tgt2 = self.conv2(F.relu(self.conv1(tgt)))
-        tgt2 = paddle.squeeze(tgt2, 2)
-        tgt2 = paddle.transpose(tgt2, [2, 0, 1])
-        tgt = paddle.squeeze(tgt, 2)
-        tgt = paddle.transpose(tgt, [2, 0, 1])
-
-        tgt = tgt + self.dropout3(tgt2)
-        tgt = self.norm3(tgt)
-        return tgt
-
-
-def _get_clones(module, N):
-    return LayerList([copy.deepcopy(module) for i in range(N)])
-
-
-class PositionalEncoding(nn.Layer):
-    """Inject some information about the relative or absolute position of the tokens
-        in the sequence. The positional encodings have the same dimension as
-        the embeddings, so that the two can be summed. Here, we use sine and cosine
-        functions of different frequencies.
-    .. math::
-        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
-        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
-        \text{where pos is the word position and i is the embed idx)
-    Args:
-        d_model: the embed dim (required).
-        dropout: the dropout value (default=0.1).
-        max_len: the max. length of the incoming sequence (default=5000).
-    Examples:
-        >>> pos_encoder = PositionalEncoding(d_model)
-    """
-
-    def __init__(self, dropout, dim, max_len=5000):
-        super(PositionalEncoding, self).__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        pe = paddle.zeros([max_len, dim])
-        position = paddle.arange(0, max_len, dtype=paddle.float32).unsqueeze(1)
-        div_term = paddle.exp(
-            paddle.arange(0, dim, 2).astype('float32') *
-            (-math.log(10000.0) / dim))
-        pe[:, 0::2] = paddle.sin(position * div_term)
-        pe[:, 1::2] = paddle.cos(position * div_term)
-        pe = paddle.unsqueeze(pe, 0)
-        pe = paddle.transpose(pe, [1, 0, 2])
-        self.register_buffer('pe', pe)
-
-    def forward(self, x):
-        """Inputs of forward function
-        Args:
-            x: the sequence fed to the positional encoder model (required).
-        Shape:
-            x: [sequence length, batch size, embed dim]
-            output: [sequence length, batch size, embed dim]
-        Examples:
-            >>> output = pos_encoder(x)
-        """
-        x = x + self.pe[:paddle.shape(x)[0], :]
-        return self.dropout(x)
-
-
-class PositionalEncoding_2d(nn.Layer):
-    """Inject some information about the relative or absolute position of the tokens
-        in the sequence. The positional encodings have the same dimension as
-        the embeddings, so that the two can be summed. Here, we use sine and cosine
-        functions of different frequencies.
-    .. math::
-        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
-        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
-        \text{where pos is the word position and i is the embed idx)
-    Args:
-        d_model: the embed dim (required).
-        dropout: the dropout value (default=0.1).
-        max_len: the max. length of the incoming sequence (default=5000).
-    Examples:
-        >>> pos_encoder = PositionalEncoding(d_model)
-    """
-
-    def __init__(self, dropout, dim, max_len=5000):
-        super(PositionalEncoding_2d, self).__init__()
-        self.dropout = nn.Dropout(p=dropout)
-
-        pe = paddle.zeros([max_len, dim])
-        position = paddle.arange(0, max_len, dtype=paddle.float32).unsqueeze(1)
-        div_term = paddle.exp(
-            paddle.arange(0, dim, 2).astype('float32') *
-            (-math.log(10000.0) / dim))
-        pe[:, 0::2] = paddle.sin(position * div_term)
-        pe[:, 1::2] = paddle.cos(position * div_term)
-        pe = paddle.transpose(paddle.unsqueeze(pe, 0), [1, 0, 2])
-        self.register_buffer('pe', pe)
-
-        self.avg_pool_1 = nn.AdaptiveAvgPool2D((1, 1))
-        self.linear1 = nn.Linear(dim, dim)
-        self.linear1.weight.data.fill_(1.)
-        self.avg_pool_2 = nn.AdaptiveAvgPool2D((1, 1))
-        self.linear2 = nn.Linear(dim, dim)
-        self.linear2.weight.data.fill_(1.)
-
-    def forward(self, x):
-        """Inputs of forward function
-        Args:
-            x: the sequence fed to the positional encoder model (required).
-        Shape:
-            x: [sequence length, batch size, embed dim]
-            output: [sequence length, batch size, embed dim]
-        Examples:
-            >>> output = pos_encoder(x)
-        """
-        w_pe = self.pe[:paddle.shape(x)[-1], :]
-        w1 = self.linear1(self.avg_pool_1(x).squeeze()).unsqueeze(0)
-        w_pe = w_pe * w1
-        w_pe = paddle.transpose(w_pe, [1, 2, 0])
-        w_pe = paddle.unsqueeze(w_pe, 2)
-
-        h_pe = self.pe[:paddle.shape(x).shape[-2], :]
-        w2 = self.linear2(self.avg_pool_2(x).squeeze()).unsqueeze(0)
-        h_pe = h_pe * w2
-        h_pe = paddle.transpose(h_pe, [1, 2, 0])
-        h_pe = paddle.unsqueeze(h_pe, 3)
-
-        x = x + w_pe + h_pe
-        x = paddle.transpose(
-            paddle.reshape(x,
-                           [x.shape[0], x.shape[1], x.shape[2] * x.shape[3]]),
-            [2, 0, 1])
-
-        return self.dropout(x)
-
-
-class Embeddings(nn.Layer):
-    def __init__(self, d_model, vocab, padding_idx, scale_embedding):
-        super(Embeddings, self).__init__()
-        self.embedding = nn.Embedding(vocab, d_model, padding_idx=padding_idx)
-        w0 = np.random.normal(0.0, d_model**-0.5,
-                              (vocab, d_model)).astype(np.float32)
-        self.embedding.weight.set_value(w0)
-        self.d_model = d_model
-        self.scale_embedding = scale_embedding
-
-    def forward(self, x):
-        if self.scale_embedding:
-            x = self.embedding(x)
-            return x * math.sqrt(self.d_model)
-        return self.embedding(x)
-
-
-class Beam():
-    ''' Beam search '''
-
-    def __init__(self, size, device=False):
-
-        self.size = size
-        self._done = False
-        # The score for each translation on the beam.
-        self.scores = paddle.zeros((size, ), dtype=paddle.float32)
-        self.all_scores = []
-        # The backpointers at each time-step.
-        self.prev_ks = []
-        # The outputs at each time-step.
-        self.next_ys = [paddle.full((size, ), 0, dtype=paddle.int64)]
-        self.next_ys[0][0] = 2
-
-    def get_current_state(self):
-        "Get the outputs for the current timestep."
-        return self.get_tentative_hypothesis()
-
-    def get_current_origin(self):
-        "Get the backpointers for the current timestep."
-        return self.prev_ks[-1]
-
-    @property
-    def done(self):
-        return self._done
-
-    def advance(self, word_prob):
-        "Update beam status and check if finished or not."
-        num_words = word_prob.shape[1]
-
-        # Sum the previous scores.
-        if len(self.prev_ks) > 0:
-            beam_lk = word_prob + self.scores.unsqueeze(1).expand_as(word_prob)
-        else:
-            beam_lk = word_prob[0]
-
-        flat_beam_lk = beam_lk.reshape([-1])
-        best_scores, best_scores_id = flat_beam_lk.topk(self.size, 0, True,
-                                                        True)  # 1st sort
-        self.all_scores.append(self.scores)
-        self.scores = best_scores
-        # bestScoresId is flattened as a (beam x word) array,
-        # so we need to calculate which word and beam each score came from
-        prev_k = best_scores_id // num_words
-        self.prev_ks.append(prev_k)
-        self.next_ys.append(best_scores_id - prev_k * num_words)
-        # End condition is when top-of-beam is EOS.
-        if self.next_ys[-1][0] == 3:
-            self._done = True
-            self.all_scores.append(self.scores)
-
-        return self._done
-
-    def sort_scores(self):
-        "Sort the scores."
-        return self.scores, paddle.to_tensor(
-            [i for i in range(int(self.scores.shape[0]))], dtype='int32')
-
-    def get_the_best_score_and_idx(self):
-        "Get the score of the best in the beam."
-        scores, ids = self.sort_scores()
-        return scores[1], ids[1]
-
-    def get_tentative_hypothesis(self):
-        "Get the decoded sequence for the current timestep."
-        if len(self.next_ys) == 1:
-            dec_seq = self.next_ys[0].unsqueeze(1)
-        else:
-            _, keys = self.sort_scores()
-            hyps = [self.get_hypothesis(k) for k in keys]
-            hyps = [[2] + h for h in hyps]
-            dec_seq = paddle.to_tensor(hyps, dtype='int64')
-        return dec_seq
-
-    def get_hypothesis(self, k):
-        """ Walk back to construct the full hypothesis. """
-        hyp = []
-        for j in range(len(self.prev_ks) - 1, -1, -1):
-            hyp.append(self.next_ys[j + 1][k])
-            k = self.prev_ks[j][k]
-        return list(map(lambda x: x.item(), hyp[::-1]))