text: """
  Implementation of "Attention is All You Need"
  """
  
  import torch
  import torch.nn as nn
  import numpy as np
  
  from models.encoder import PositionalEncoding
  from models.neural import MultiHeadedAttention, PositionwiseFeedForward, DecoderState
  
  MAX_SIZE = 5000
  
  
  class TransformerDecoderLayer(nn.Module):
      """
      Args:
        d_model (int): the dimension of keys/values/queries in
                         MultiHeadedAttention, also the input size of
                         the first-layer of the PositionwiseFeedForward.
        heads (int): the number of heads for MultiHeadedAttention.
        d_ff (int): the second-layer of the PositionwiseFeedForward.
        dropout (float): dropout probability(0-1.0).
        self_attn_type (string): type of self-attention scaled-dot, average
      """
  
      def __init__(self, d_model, heads, d_ff, dropout):
          super(TransformerDecoderLayer, self).__init__()
  
          self.self_attn = MultiHeadedAttention(
              heads, d_model, dropout=dropout)
  
          self.context_attn = MultiHeadedAttention(heads, d_model, dropout=dropout)
          self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
          self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-6)
          self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-6)
          self.drop = nn.Dropout(dropout)
          mask = self._get_attn_subsequent_mask(MAX_SIZE)
          # Register self.mask as a buffer in TransformerDecoderLayer, so
          # it gets TransformerDecoderLayer's cuda behavior automatically.
          self.register_buffer('mask', mask)
  
      def forward(self, inputs, memory_bank, src_pad_mask, tgt_pad_mask, previous_input=None,
                  layer_cache=None, requires_att=False, ignore_memory_attn=False, self_attn=True):
          """
          Args:
              inputs (`FloatTensor`): `[batch_size x 1 x model_dim]`
              memory_bank (`FloatTensor`): `[batch_size x src_len x model_dim]`
              src_pad_mask (`LongTensor`): `[batch_size x 1 x src_len]`
              tgt_pad_mask (`LongTensor`): `[batch_size x 1 x 1]`
  
          Returns:
              (`FloatTensor`, `FloatTensor`, `FloatTensor`):
  
              * output `[batch_size x 1 x model_dim]`
              * attn `[batch_size x 1 x src_len]`
              * all_input `[batch_size x current_step x model_dim]`
  
          """
          if self_attn:
              dec_mask = torch.gt(tgt_pad_mask +
                                  self.mask[:, :tgt_pad_mask.size(1),
                                            :tgt_pad_mask.size(1)], 0)
          else:
              dec_mask = tgt_pad_mask
  
          input_norm = self.layer_norm_1(inputs)
          all_input = input_norm
          if previous_input is not None:
              all_input = torch.cat((previous_input, input_norm), dim=1)
              dec_mask = None
  
          query, _ = self.self_attn(all_input, all_input, input_norm,
                                    mask=dec_mask,
                                    layer_cache=layer_cache,
                                    type="self")
  
          query = self.drop(query) + inputs
  
          if not ignore_memory_attn:
              query_norm = self.layer_norm_2(query)
              mid, att = self.context_attn(memory_bank, memory_bank, query_norm,
                                           mask=src_pad_mask,
                                           layer_cache=layer_cache,
                                           type="context",
                                           requires_att=requires_att)
              mid = self.drop(mid) + query
  
              output = self.feed_forward(mid)
          else:
              output = self.feed_forward(query)
              att = None
  
          return output, all_input, att
          # return output
  
      def _get_attn_subsequent_mask(self, size):
          """
          Get an attention mask to avoid using the subsequent info.
  
          Args:
              size: int
  
          Returns:
              (`LongTensor`):
  
              * subsequent_mask `[1 x size x size]`
          """
          attn_shape = (1, size, size)
          subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
          subsequent_mask = torch.from_numpy(subsequent_mask)
          return subsequent_mask
  
  
  class TransformerDecoder(nn.Module):
      """
      The Transformer decoder from "Attention is All You Need".
  
  
      .. mermaid::
  
         graph BT
            A[input]
            B[multi-head self-attn]
            BB[multi-head src-attn]
            C[feed forward]
            O[output]
            A --> B
            B --> BB
            BB --> C
            C --> O
  
  
      Args:
         num_layers (int): number of encoder layers.
         d_model (int): size of the model
         heads (int): number of heads
         d_ff (int): size of the inner FF layer
         dropout (float): dropout parameters
         embeddings (:obj:`onmt.modules.Embeddings`):
            embeddings to use, should have positional encodings
         attn_type (str): if using a seperate copy attention
      """
  
      def __init__(self, num_layers, d_model, heads, d_ff, dropout,
                   embeddings=None):
          super(TransformerDecoder, self).__init__()
  
          # Basic attributes.
          self.decoder_type = 'transformer'
          self.num_layers = num_layers
          self.hidden_size = d_model
  
          if embeddings is not None:
              self.embeddings = embeddings
              if self.embeddings.embedding_dim != d_model:
                  self.emb_to_hid = nn.Linear(self.embeddings.embedding_dim, d_model)
              self.pos_emb = PositionalEncoding(dropout, d_model)
  
          # Build TransformerDecoder.
          self.transformer_layers = nn.ModuleList(
              [TransformerDecoderLayer(d_model, heads, d_ff, dropout)
               for _ in range(num_layers)])
  
          self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
  
      def forward(self, tgt, memory_bank, state, init_tokens=None,
                  step=None, cache=None, memory_masks=None, tgt_masks=None,
                  requires_att=False, ignore_memory_attn=False, self_attn=True):
  
          tgt_batch, tgt_len = tgt.size()
  
          # Run the forward pass of the TransformerDecoder.
          emb = self.embeddings(tgt)
          if emb.size(-1) != self.hidden_size:
              emb = self.emb_to_hid(emb)
          if init_tokens is not None:
              emb = torch.cat([init_tokens.unsqueeze(1), emb[:, 1:, :]], 1)
          assert emb.dim() == 3  # len x batch x embedding_dim
  
          if ignore_memory_attn:
              assert memory_bank.dim() == 2
              output = self.pos_emb(emb, step, memory_bank.unsqueeze(1))
          else:
              output = self.pos_emb(emb, step)
  
          if tgt_masks is not None:
              tgt_pad_mask = tgt_masks.unsqueeze(1).expand(tgt_batch, tgt_len, tgt_len)
          else:
              assert tgt.dim() == 2
              padding_idx = self.embeddings.padding_idx
              tgt_pad_mask = tgt.data.eq(padding_idx).unsqueeze(1) \
                  .expand(tgt_batch, tgt_len, tgt_len)
  
          src_memory_bank = memory_bank
          if ignore_memory_attn:
              src_pad_mask = None
          elif memory_masks is not None:
              src_batch = memory_masks.size(0)
              src_len = memory_masks.size(-1)
              src_pad_mask = memory_masks.unsqueeze(1).expand(src_batch, tgt_len, src_len)
          else:
              src_batch = memory_bank.size(0)
              src_len = memory_bank.size(1)
              src_pad_mask = tgt_pad_mask.new_zeros([src_batch, tgt_len, src_len])
  
          if state.cache is None:
              saved_inputs = []
  
          for i in range(self.num_layers):
              prev_layer_input = None
              if state.cache is None:
                  if state.previous_input is not None:
                      prev_layer_input = state.previous_layer_inputs[i]
              output, all_input, last_layer_att \
                  = self.transformer_layers[i](
                      output, src_memory_bank,
                      src_pad_mask, tgt_pad_mask,
                      previous_input=prev_layer_input,
                      layer_cache=state.cache["layer_{}".format(i)]
                      if state.cache is not None else None,
                      requires_att=False if i < self.num_layers-1 else requires_att,
                      ignore_memory_attn=ignore_memory_attn,
                      self_attn=self_attn)
              if state.cache is None:
                  saved_inputs.append(all_input)
  
          if state.cache is None:
              saved_inputs = torch.stack(saved_inputs)
  
          output = self.layer_norm(output)
  
          # Process the result and update the attentions.
  
          if state.cache is None:
              state = state.update_state(tgt, saved_inputs)
  
          if requires_att and last_layer_att is not None:
              return output, state, {"copy": last_layer_att}
          else:
              return output, state, None
  
      def init_decoder_state(self, src, memory_bank, enc_hidden=None,
                             with_cache=False):
          """ Init decoder state """
          state = TransformerDecoderState(src)
          if with_cache:
              state._init_cache(memory_bank, self.num_layers)
          return state
  
  
  class TransformerDecoderState(DecoderState):
      """ Transformer Decoder state base class """
  
      def __init__(self, src):
          """
          Args:
              src (FloatTensor): a sequence of source words tensors
                      with optional feature tensors, of size (len x batch).
          """
          self.src = src
          self.previous_input = None
          self.previous_layer_inputs = None
          self.cache = None
  
      @property
      def _all(self):
          """
          Contains attributes that need to be updated in self.beam_update().
          """
          if (self.previous_input is not None
                  and self.previous_layer_inputs is not None):
              return (self.previous_input,
                      self.previous_layer_inputs,
                      self.src)
          else:
              return (self.src,)
  
      def detach(self):
          if self.previous_input is not None:
              self.previous_input = self.previous_input.detach()
          if self.previous_layer_inputs is not None:
              self.previous_layer_inputs = self.previous_layer_inputs.detach()
          self.src = self.src.detach()
  
      def update_state(self, new_input, previous_layer_inputs):
          state = TransformerDecoderState(self.src)
          state.previous_input = new_input
          state.previous_layer_inputs = previous_layer_inputs
          return state
  
      def _init_cache(self, memory_bank, num_layers):
          self.cache = {}
  
          for layer in range(num_layers):
              layer_cache = {
                  "memory_keys": None,
                  "memory_values": None
              }
              layer_cache["self_keys"] = None
              layer_cache["self_values"] = None
              self.cache["layer_{}".format(layer)] = layer_cache
  
      def repeat_beam_size_times(self, beam_size):
          """ Repeat beam_size times along batch dimension. """
          self.src = self.src.data.repeat(1, beam_size, 1)
  
      def map_batch_fn(self, fn):
          def _recursive_map(struct, batch_dim=0):
              for k, v in struct.items():
                  if v is not None:
                      if isinstance(v, dict):
                          _recursive_map(v)
                      else:
                          struct[k] = fn(v, batch_dim)
  
          self.src = fn(self.src, 0)
          if self.cache is not None:
              _recursive_map(self.cache)
  