text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from typing import Dict, List, Optional
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from fairseq import utils
  from fairseq.modules import LayerNorm, MultiheadAttention
  from fairseq.modules.quant_noise import quant_noise
  from torch import Tensor
  
  
  class TransformerEncoderLayer(nn.Module):
      """Encoder layer block.
  
      In the original paper each operation (multi-head attention or FFN) is
      postprocessed with: `dropout -> add residual -> layernorm`. In the
      tensor2tensor code they suggest that learning is more robust when
      preprocessing each layer with layernorm and postprocessing with:
      `dropout -> add residual`. We default to the approach in the paper, but the
      tensor2tensor approach can be enabled by setting
      *args.encoder_normalize_before* to ``True``.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
      """
  
      def __init__(self, args):
          super().__init__()
          self.embed_dim = args.encoder_embed_dim
          self.quant_noise = getattr(args, "quant_noise_pq", 0)
          self.quant_noise_block_size = getattr(args, "quant_noise_pq_block_size", 8)
          self.self_attn = self.build_self_attention(self.embed_dim, args)
          self.self_attn_layer_norm = LayerNorm(self.embed_dim)
          self.dropout = args.dropout
          self.activation_fn = utils.get_activation_fn(
              activation=getattr(args, "activation_fn", "relu")
          )
          self.activation_dropout = getattr(args, "activation_dropout", 0)
          if self.activation_dropout == 0:
              # for backwards compatibility with models that use args.relu_dropout
              self.activation_dropout = getattr(args, "relu_dropout", 0)
          self.normalize_before = args.encoder_normalize_before
          self.fc1 = self.build_fc1(
              self.embed_dim, args.encoder_ffn_embed_dim, self.quant_noise, self.quant_noise_block_size
          )
          self.fc2 = self.build_fc2(
              args.encoder_ffn_embed_dim, self.embed_dim, self.quant_noise, self.quant_noise_block_size
          )
  
          self.final_layer_norm = LayerNorm(self.embed_dim)
  
      def build_fc1(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), p=q_noise, block_size=qn_block_size)
  
      def build_fc2(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), p=q_noise, block_size=qn_block_size)
  
      def build_self_attention(self, embed_dim, args):
          return MultiheadAttention(
              embed_dim,
              args.encoder_attention_heads,
              dropout=args.attention_dropout,
              self_attention=True,
              q_noise=self.quant_noise,
              qn_block_size=self.quant_noise_block_size,
              relative_pos_type=("unmasked"
                                 if getattr(args, "use_relative_pos_embeddings", False)
                                 else None),
              max_relative_pos=getattr(args, "max_relative_pos", 128),
              heads_share_embeddings=getattr(args, "heads_share_embeddings", False),
              add_pos_embeddings_to_values=getattr(args, "add_pos_embeddings_to_values", False)
          )
  
      def upgrade_state_dict_named(self, state_dict, name):
          """
          Rename layer norm states from `...layer_norms.0.weight` to
          `...self_attn_layer_norm.weight` and `...layer_norms.1.weight` to
          `...final_layer_norm.weight`
          """
          layer_norm_map = {"0": "self_attn_layer_norm", "1": "final_layer_norm"}
          for old, new in layer_norm_map.items():
              for m in ("weight", "bias"):
                  k = "{}.layer_norms.{}.{}".format(name, old, m)
                  if k in state_dict:
                      state_dict["{}.{}.{}".format(name, new, m)] = state_dict[k]
                      del state_dict[k]
  
      def forward(
          self,
          x,
          encoder_padding_mask,
          attn_mask: Optional[Tensor] = None,
          need_attn: bool = False,
          need_head_weights: bool = False,
          sentence_position = None,
          disable_add_1stblock = False,
          only_weighted_sum = False):
                  #####
          """
          Args:
              x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
              encoder_padding_mask (ByteTensor): binary ByteTensor of shape
                  `(batch, src_len)` where padding elements are indicated by ``1``.
              attn_mask (ByteTensor): binary tensor of shape (T_tgt, T_src), where
              T_tgt is the length of query, while T_src is the length of key,
              though here both query and key is x here,
              attn_mask[t_tgt, t_src] = 1 means when calculating embedding
              for t_tgt, t_src is excluded (or masked out), =0 means it is
              included in attention
              need_attn (bool, optional): return attention weights.
              need_head_weights (bool, optional): return attention weights
                  for each head (default: return average over heads).
          Returns:
              encoded output of shape `(seq_len, batch, embed_dim)`
          """
          residual = x
          if self.normalize_before:
              x = self.self_attn_layer_norm(x)
          if attn_mask is not None:
              attn_mask = attn_mask.masked_fill(attn_mask.to(torch.bool), -1e8)
          if need_head_weights:
              need_attn = True
  
          # anything in original attn_mask = 1, becomes -1e8
          # anything in original attn_mask = 0, becomes 0
          # Note that we cannot use -inf here, because at some edge cases,
          # the attention weight (before softmax) for some padded element in query
          # will become -inf, which results in NaN in model parameters
          # TODO: to formally solve this problem, we need to change fairseq's
          # MultiheadAttention. We will do this later on.
  
          x, attn = self.self_attn(
              query=x,
              key=x,
              value=x,
              key_padding_mask=encoder_padding_mask,
              attn_mask=attn_mask,
              sentence_position = sentence_position,
              need_weights=True,
              need_head_weights=need_head_weights,
          )
          ###
          x = F.dropout(x, p=self.dropout, training=self.training)
          #####
          if not disable_add_1stblock:       
              x = residual + x
          if not self.normalize_before:
              x = self.self_attn_layer_norm(x)
          if only_weighted_sum:
              if need_attn:
                  return x, attn
              else:     
                  return x, None
  
          residual = x
          if self.normalize_before:
              x = self.final_layer_norm(x)
  
          x = self.activation_fn(self.fc1(x))
          x = F.dropout(x, p=float(self.activation_dropout), training=self.training)
          x = self.fc2(x)
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          if not self.normalize_before:
              x = self.final_layer_norm(x)
          if need_attn:
              return x, attn
          return x, None
  
  
  class TransformerDecoderLayer(nn.Module):
      """Decoder layer block.
  
      In the original paper each operation (multi-head attention, encoder
      attention or FFN) is postprocessed with: `dropout -> add residual ->
      layernorm`. In the tensor2tensor code they suggest that learning is more
      robust when preprocessing each layer with layernorm and postprocessing with:
      `dropout -> add residual`. We default to the approach in the paper, but the
      tensor2tensor approach can be enabled by setting
      *args.decoder_normalize_before* to ``True``.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
          no_encoder_attn (bool, optional): whether to attend to encoder outputs
              (default: False).
      """
  
      def __init__(
          self, args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False
      ):
          super().__init__()
          self.embed_dim = args.decoder_embed_dim
          self.dropout = args.dropout
          self.quant_noise = getattr(args, "quant_noise_pq", 0)
          self.quant_noise_block_size = getattr(args, "quant_noise_pq_block_size", 8)
  
          self.cross_self_attention = getattr(args, "cross_self_attention", False)
  
          self.self_attn = self.build_self_attention(
              self.embed_dim,
              args,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
          )
          self.activation_fn = utils.get_activation_fn(
              activation=getattr(args, "activation_fn", "relu")
          )
          self.activation_dropout = getattr(args, "activation_dropout", 0)
          if self.activation_dropout == 0:
              # for backwards compatibility with models that use args.relu_dropout
              self.activation_dropout = getattr(args, "relu_dropout", 0)
          self.normalize_before = args.decoder_normalize_before
  
          # use layerNorm rather than FusedLayerNorm for exporting.
          # char_inputs can be used to determint this.
          # TODO  remove this once we update apex with the fix
          export = getattr(args, "char_inputs", False)
          self.self_attn_layer_norm = LayerNorm(self.embed_dim, export=export)
  
          if no_encoder_attn:
              self.encoder_attn = None
              self.encoder_attn_layer_norm = None
          else:
              self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
              self.encoder_attn_layer_norm = LayerNorm(self.embed_dim, export=export)
  
          self.fc1 = self.build_fc1(
              self.embed_dim, args.decoder_ffn_embed_dim, self.quant_noise, self.quant_noise_block_size
          )
          self.fc2 = self.build_fc2(
              args.decoder_ffn_embed_dim, self.embed_dim, self.quant_noise, self.quant_noise_block_size
          )
  
          self.final_layer_norm = LayerNorm(self.embed_dim, export=export)
          self.need_attn = True
  
          self.onnx_trace = False
  
      def build_fc1(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), q_noise, qn_block_size)
  
      def build_fc2(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), q_noise, qn_block_size)
  
      def build_self_attention(self, embed_dim, args, add_bias_kv=False, add_zero_attn=False):
          return MultiheadAttention(
              embed_dim,
              args.decoder_attention_heads,
              dropout=args.attention_dropout,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
              self_attention=not getattr(args, "cross_self_attention", False),
              q_noise=self.quant_noise,
              qn_block_size=self.quant_noise_block_size,
              # relative_pos_type=("masked"
              #                    if getattr(args, "use_relative_pos_embeddings", False)
              #                    else None),
              relative_pos_type = None, ### 
              max_relative_pos=getattr(args, "max_relative_pos", 128),
              heads_share_embeddings=getattr(args, "heads_share_embeddings", False),
              add_pos_embeddings_to_values=getattr(args, "add_pos_embeddings_to_values", False)
          )
  
      def build_encoder_attention(self, embed_dim, args):
          return MultiheadAttention(
              embed_dim,
              args.decoder_attention_heads,
              kdim=getattr(args, "encoder_embed_dim", None),
              vdim=getattr(args, "encoder_embed_dim", None),
              dropout=args.attention_dropout,
              encoder_decoder_attention=True,
              q_noise=self.quant_noise,
              qn_block_size=self.quant_noise_block_size,
          )
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      def forward(
          self,
          x,
          encoder_out: Optional[torch.Tensor] = None,
          encoder_padding_mask: Optional[torch.Tensor] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          prev_self_attn_state: Optional[List[torch.Tensor]] = None,
          prev_attn_state: Optional[List[torch.Tensor]] = None,
          self_attn_mask: Optional[torch.Tensor] = None,
          self_attn_padding_mask: Optional[torch.Tensor] = None,
          need_attn: bool = False,
          need_head_weights: bool = False,
      ):
          """
          Args:
              x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
              encoder_padding_mask (ByteTensor, optional): binary
                  ByteTensor of shape `(batch, src_len)` where padding
                  elements are indicated by ``1``.
              need_attn (bool, optional): return attention weights
              need_head_weights (bool, optional): return attention weights
                  for each head (default: return average over heads).
  
          Returns:
              encoded output of shape `(seq_len, batch, embed_dim)`
          """
          if need_head_weights:
              need_attn = True
  
          residual = x
          if self.normalize_before:
              x = self.self_attn_layer_norm(x)
          if prev_self_attn_state is not None:
              prev_key, prev_value = prev_self_attn_state[:2]
              saved_state: Dict[str, Optional[Tensor]] = {
                  "prev_key": prev_key,
                  "prev_value": prev_value,
              }
              if len(prev_self_attn_state) >= 3:
                  saved_state["prev_key_padding_mask"] = prev_self_attn_state[2]
              assert incremental_state is not None
              self.self_attn._set_input_buffer(incremental_state, saved_state)
          _self_attn_input_buffer = self.self_attn._get_input_buffer(incremental_state)
          if self.cross_self_attention and not (
              incremental_state is not None
              and _self_attn_input_buffer is not None
              and "prev_key" in _self_attn_input_buffer
          ):
              if self_attn_mask is not None:
                  assert encoder_out is not None
                  self_attn_mask = torch.cat(
                      (x.new_zeros(x.size(0), encoder_out.size(0)), self_attn_mask), dim=1
                  )
              if self_attn_padding_mask is not None:
                  if encoder_padding_mask is None:
                      assert encoder_out is not None
                      encoder_padding_mask = self_attn_padding_mask.new_zeros(
                          encoder_out.size(1), encoder_out.size(0)
                      )
                  self_attn_padding_mask = torch.cat(
                      (encoder_padding_mask, self_attn_padding_mask), dim=1
                  )
              assert encoder_out is not None
              y = torch.cat((encoder_out, x), dim=0)
          else:
              y = x
  
          x, attn = self.self_attn(
              query=x,
              key=y,
              value=y,
              key_padding_mask=self_attn_padding_mask,
              incremental_state=incremental_state,
              need_weights=False,
              attn_mask=self_attn_mask,
          )
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          if not self.normalize_before:
              x = self.self_attn_layer_norm(x)
  
          if self.encoder_attn is not None:
              residual = x
              if self.normalize_before:
                  x = self.encoder_attn_layer_norm(x)
              if prev_attn_state is not None:
                  prev_key, prev_value = prev_attn_state[:2]
                  saved_state: Dict[str, Optional[Tensor]] = {
                      "prev_key": prev_key,
                      "prev_value": prev_value,
                  }
                  if len(prev_attn_state) >= 3:
                      saved_state["prev_key_padding_mask"] = prev_attn_state[2]
                  assert incremental_state is not None
                  self.encoder_attn._set_input_buffer(incremental_state, saved_state)
  
              x, attn = self.encoder_attn(
                  query=x,
                  key=encoder_out,
                  value=encoder_out,
                  key_padding_mask=encoder_padding_mask,
                  incremental_state=incremental_state,
                  static_kv=True,
                  need_weights=need_attn or (not self.training and self.need_attn),
                  need_head_weights=need_head_weights,
              )
              x = F.dropout(x, p=self.dropout, training=self.training)
              x = residual + x
              if not self.normalize_before:
                  x = self.encoder_attn_layer_norm(x)
  
          residual = x
          if self.normalize_before:
              x = self.final_layer_norm(x)
  
          x = self.activation_fn(self.fc1(x))
          x = F.dropout(x, p=float(self.activation_dropout), training=self.training)
          x = self.fc2(x)
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          if not self.normalize_before:
              x = self.final_layer_norm(x)
          if self.onnx_trace and incremental_state is not None:
              saved_state = self.self_attn._get_input_buffer(incremental_state)
              assert saved_state is not None
              if self_attn_padding_mask is not None:
                  self_attn_state = [
                      saved_state["prev_key"],
                      saved_state["prev_value"],
                      saved_state["prev_key_padding_mask"],
                  ]
              else:
                  self_attn_state = [saved_state["prev_key"], saved_state["prev_value"]]
              return x, attn, self_attn_state
          return x, attn, None
  
      def make_generation_fast_(self, need_attn: bool = False, **kwargs):
          self.need_attn = need_attn
  
      @torch.jit.export
      def reorder_incremental_state(
          self,
          incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
          new_order: Tensor,
      ):
          """Scriptable reorder incremental state in transformer layers."""
          self.self_attn.reorder_incremental_state(incremental_state, new_order)
  
          if self.encoder_attn is not None:
              self.encoder_attn.reorder_incremental_state(incremental_state, new_order)
  
  
  def Linear(in_features, out_features, bias=True):
      m = nn.Linear(in_features, out_features, bias)
      nn.init.xavier_uniform_(m.weight)
      if bias:
          nn.init.constant_(m.bias, 0.0)
      return m
  