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, Tuple
  
  import torch
  import torch.nn as nn
  from fairseq import utils
  from fairseq.dataclass.utils import gen_parser_from_dataclass
  from fairseq.distributed import fsdp_wrap
  from fairseq.models import FairseqEncoderDecoderModel
  from fairseq.models.transformer import (
      TransformerEncoderBase,
      TransformerDecoderBase,
      TransformerConfig,
  )
  from torch import Tensor
  
  
  class TransformerModelBase(FairseqEncoderDecoderModel):
      """
      Transformer model from `"Attention Is All You Need" (Vaswani, et al, 2017)
      <https://arxiv.org/abs/1706.03762>`_.
  
      Args:
          encoder (TransformerEncoder): the encoder
          decoder (TransformerDecoder): the decoder
  
      The Transformer model provides the following named architectures and
      command-line arguments:
  
      .. argparse::
          :ref: fairseq.models.transformer_parser
          :prog:
      """
  
      def __init__(self, cfg, encoder, decoder):
          super().__init__(encoder, decoder)
          self.cfg = cfg
          self.supports_align_args = True
  
      @staticmethod
      def add_args(parser):
          """Add model-specific arguments to the parser."""
          # we want to build the args recursively in this case.
          gen_parser_from_dataclass(
              parser, TransformerConfig(), delete_default=False, with_prefix=""
          )
  
      @classmethod
      def build_model(cls, cfg, task):
          """Build a new model instance."""
  
          # --  TODO T96535332
          #  bug caused by interaction between OmegaConf II and argparsing
          cfg.decoder.input_dim = int(cfg.decoder.input_dim)
          cfg.decoder.output_dim = int(cfg.decoder.output_dim)
          # --
  
          if cfg.encoder.layers_to_keep:
              cfg.encoder.layers = len(cfg.encoder.layers_to_keep.split(","))
          if cfg.decoder.layers_to_keep:
              cfg.decoder.layers = len(cfg.decoder.layers_to_keep.split(","))
  
          src_dict, tgt_dict = task.source_dictionary, task.target_dictionary
  
          if cfg.share_all_embeddings:
              if src_dict != tgt_dict:
                  raise ValueError("--share-all-embeddings requires a joined dictionary")
              if cfg.encoder.embed_dim != cfg.decoder.embed_dim:
                  raise ValueError(
                      "--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim"
                  )
              if cfg.decoder.embed_path and (
                  cfg.decoder.embed_path != cfg.encoder.embed_path
              ):
                  raise ValueError(
                      "--share-all-embeddings not compatible with --decoder-embed-path"
                  )
              encoder_embed_tokens = cls.build_embedding(
                  cfg, src_dict, cfg.encoder.embed_dim, cfg.encoder.embed_path
              )
              decoder_embed_tokens = encoder_embed_tokens
              cfg.share_decoder_input_output_embed = True
          else:
              encoder_embed_tokens = cls.build_embedding(
                  cfg, src_dict, cfg.encoder.embed_dim, cfg.encoder.embed_path
              )
              decoder_embed_tokens = cls.build_embedding(
                  cfg, tgt_dict, cfg.decoder.embed_dim, cfg.decoder.embed_path
              )
          if cfg.offload_activations:
              cfg.checkpoint_activations = True  # offloading implies checkpointing
          encoder = cls.build_encoder(cfg, src_dict, encoder_embed_tokens)
          decoder = cls.build_decoder(cfg, tgt_dict, decoder_embed_tokens)
          if not cfg.share_all_embeddings:
              # fsdp_wrap is a no-op when --ddp-backend != fully_sharded
              encoder = fsdp_wrap(encoder, min_num_params=cfg.min_params_to_wrap)
              decoder = fsdp_wrap(decoder, min_num_params=cfg.min_params_to_wrap)
          return cls(cfg, encoder, decoder)
  
      @classmethod
      def build_embedding(cls, cfg, dictionary, embed_dim, path=None):
          num_embeddings = len(dictionary)
          padding_idx = dictionary.pad()
  
          emb = Embedding(num_embeddings, embed_dim, padding_idx)
          # if provided, load from preloaded dictionaries
          if path:
              embed_dict = utils.parse_embedding(path)
              utils.load_embedding(embed_dict, dictionary, emb)
          return emb
  
      @classmethod
      def build_encoder(cls, cfg, src_dict, embed_tokens):
          return TransformerEncoderBase(cfg, src_dict, embed_tokens)
  
      @classmethod
      def build_decoder(cls, cfg, tgt_dict, embed_tokens):
          return TransformerDecoderBase(
              cfg,
              tgt_dict,
              embed_tokens,
              no_encoder_attn=cfg.no_cross_attention,
          )
  
      # TorchScript doesn't support optional arguments with variable length (**kwargs).
      # Current workaround is to add union of all arguments in child classes.
      def forward(
          self,
          src_tokens,
          src_lengths,
          prev_output_tokens,
          return_all_hiddens: bool = True,
          features_only: bool = False,
          alignment_layer: Optional[int] = None,
          alignment_heads: Optional[int] = None,
      ):
          """
          Run the forward pass for an encoder-decoder model.
  
          Copied from the base class, but without ``**kwargs``,
          which are not supported by TorchScript.
          """
          # With probability `word_dropout_mixture`, perform word dropout.
          # Otherwise, keep the batch the same.
          if self.training and torch.rand(1) < self.encoder.word_dropout_mixture:
              batch_size, src_len = src_tokens.shape
              _, tgt_len = prev_output_tokens.shape
              # encoder_word_dropout_mask.shape = (batch_size, src_len, src_len),
              # where each dropped out token will be -1e4 and each kept token 
              # will be 1.
              # cross_attn_mask.shape = (batch_size, tgt_len, src_len), where
              # each dropped out token will be -1e4 and each kept token will be
              # 1.
              encoder_word_dropout_mask = torch.eye(src_len)[None].repeat(
                  [batch_size, 1, 1]).to(prev_output_tokens.device)
              cross_attn_word_dropout_mask = torch.zeros(
                  [batch_size, tgt_len, src_len]).to(encoder_word_dropout_mask)
              if self.encoder.word_dropout_type == 'uniform_length':
                  candidates = torch.randperm(src_len)
                  num_tokens_to_keep = torch.randint(low=0, high=src_len + 1, size=[1])
                  encoder_word_dropout_mask[:, :, candidates[:num_tokens_to_keep]] = 1.
                  cross_attn_word_dropout_mask[:, :, candidates[:num_tokens_to_keep]] = 1.
              elif self.encoder.word_dropout_type == 'inverse_length':
                  # dropout_prob = 1 - (1 / src_len), so keep_prob = 1 / src_len
                  src_tokens_to_keep = torch.bernoulli((1 / src_len) + torch.zeros(
                      [batch_size, src_len])).to(encoder_word_dropout_mask)
                  encoder_word_dropout_mask += src_tokens_to_keep[:, None, :]
                  encoder_word_dropout_mask = torch.clamp(
                      encoder_word_dropout_mask, 0., 1.)
                  cross_attn_word_dropout_mask += src_tokens_to_keep[:, None, :]
              else:
                  raise ValueError("Dropout type must be 'uniform_length' or "
                                   "'inverse_length'")
              # To avoid nan errors, we use -1e4 instead of -inf.
              encoder_word_dropout_mask = torch.where(
                  encoder_word_dropout_mask == 0., -1e4, 0.)
              cross_attn_word_dropout_mask = torch.where(
                  cross_attn_word_dropout_mask == 0., -1e4, 0.)
          else:
              encoder_word_dropout_mask = None
              cross_attn_word_dropout_mask = None
  
          encoder_out = self.encoder(
              src_tokens, src_lengths=src_lengths, return_all_hiddens=return_all_hiddens,
              word_dropout_mask=encoder_word_dropout_mask,
          )
          decoder_out = self.decoder(
              prev_output_tokens,
              encoder_out=encoder_out,
              features_only=features_only,
              alignment_layer=alignment_layer,
              alignment_heads=alignment_heads,
              src_lengths=src_lengths,
              return_all_hiddens=return_all_hiddens,
              cross_attn_word_dropout_mask=cross_attn_word_dropout_mask,
          )
          return decoder_out
  
      # Since get_normalized_probs is in the Fairseq Model which is not scriptable,
      # I rewrite the get_normalized_probs from Base Class to call the
      # helper function in the Base Class.
      @torch.jit.export
      def get_normalized_probs(
          self,
          net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
          log_probs: bool,
          sample: Optional[Dict[str, Tensor]] = None,
      ):
          """Get normalized probabilities (or log probs) from a net's output."""
          return self.get_normalized_probs_scriptable(net_output, log_probs, sample)
  
  
  def Embedding(num_embeddings, embedding_dim, padding_idx):
      m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
      nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
      nn.init.constant_(m.weight[padding_idx], 0)
      return m
  