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.
  
  import math
  from typing import Any, Dict, List, Optional, Tuple, NamedTuple
  
  import torch
  import torch.nn as nn
  from torch.nn import Parameter
  from fairseq import options, utils
  from fairseq.models import (
      FairseqEncoder,
      FairseqEncoderDecoderModel,
      FairseqIncrementalDecoder,
      register_model,
      register_model_architecture,
  )
  from fairseq.models.bart.hub_interface import BARTHubInterface
  from fairseq.modules import (
      AdaptiveSoftmax,
      FairseqDropout,
      FairseqFeatureDropout,
      LayerDropModuleList,
      LayerNorm,
      PositionalEmbedding,
      SinusoidalPositionalEmbedding,
      LunaDecoderLayer,
      LunaEncoderLayer,
  )
  from fairseq.modules.quant_noise import quant_noise as apply_quant_noise_
  from torch import Tensor
  
  
  DEFAULT_MAX_SOURCE_POSITIONS = 1024
  DEFAULT_MAX_TARGET_POSITIONS = 1024
  
  EncoderOut = NamedTuple(
      "EncoderOut",
      [
          ("encoder_out", Tensor),  # T x B x C
          ("encoder_projected_out", Tensor),  # L x B x C
          ("encoder_padding_mask", Optional[Tensor]),  # B x T
          ("encoder_projected_padding_mask", Optional[Tensor]),  # B x L
          ("encoder_embedding", Optional[Tensor]),  # B x T x C
          ("encoder_states", Optional[List[Tensor]]),  # List[T x B x C]
          ("src_tokens", Optional[Tensor]),  # B x T
          ("src_lengths", Optional[Tensor]),  # B x 1
      ],
  )
  
  
  @register_model("luna")
  class LunaModel(FairseqEncoderDecoderModel):
      """
  
      Args:
          encoder (LunaEncoder): the encoder
          decoder (LunaDecoder): the decoder
  
      The Luna model provides the following named architectures and
      command-line arguments:
  
      .. argparse::
          :ref: fairseq.models.luna_parser
          :prog:
      """
  
      def __init__(self, args, encoder, decoder):
          super().__init__(encoder, decoder)
          self.args = args
          self.supports_align_args = True
  
      @staticmethod
      def add_args(parser):
          """Add model-specific arguments to the parser."""
          # fmt: off
          parser.add_argument('--activation-fn',
                              choices=utils.get_available_activation_fns(),
                              help='activation function to use')
          parser.add_argument('--dropout', type=float, metavar='D',
                              help='dropout probability')
          parser.add_argument('--word-dropout', type=float, metavar='D',
                              help='dropout probability of words')
          parser.add_argument('--attention-dropout', type=float, metavar='D',
                              help='dropout probability for attention weights')
          parser.add_argument('--activation-dropout', '--relu-dropout', type=float, metavar='D',
                              help='dropout probability after activation in FFN.')
          parser.add_argument('--encoder-embed-path', type=str, metavar='STR',
                              help='path to pre-trained encoder embedding')
          parser.add_argument('--encoder-embed-dim', type=int, metavar='N',
                              help='encoder embedding dimension')
          parser.add_argument('--encoder-ffn-embed-dim', type=int, metavar='N',
                              help='encoder embedding dimension for FFN')
          parser.add_argument('--encoder-layers', type=int, metavar='N',
                              help='num encoder layers')
          parser.add_argument('--encoder-attention-heads', type=int, metavar='N',
                              help='num encoder attention heads')
          parser.add_argument('--encoder-projected-attention-heads', type=int, metavar='N',
                              help='num encoder projected attention heads')
          parser.add_argument('--encoder-learned-pos', action='store_true',
                              help='use learned positional embeddings in the encoder')
          parser.add_argument('--encoder-normalize-before', action='store_true',
                              help='apply layernorm before each encoder block')
          parser.add_argument('--decoder-embed-path', type=str, metavar='STR',
                              help='path to pre-trained decoder embedding')
          parser.add_argument('--decoder-embed-dim', type=int, metavar='N',
                              help='decoder embedding dimension')
          parser.add_argument('--decoder-ffn-embed-dim', type=int, metavar='N',
                              help='decoder embedding dimension for FFN')
          parser.add_argument('--decoder-layers', type=int, metavar='N',
                              help='num decoder layers')
          parser.add_argument('--decoder-attention-heads', type=int, metavar='N',
                              help='num decoder attention heads')
          parser.add_argument('--decoder-projected-attention-heads', type=int, metavar='N',
                              help='num decoder projected attention heads')
          parser.add_argument('--decoder-learned-pos', action='store_true',
                              help='use learned positional embeddings in the decoder')
          parser.add_argument('--decoder-output-dim', type=int, metavar='N',
                              help='decoder output dimension (extra linear layer '
                                   'if different from decoder embed dim')
          parser.add_argument('--decoder-normalize-before', action='store_true',
                              help='apply layernorm before each decoder block')
          parser.add_argument('--projection-length', type=int, metavar='N',
                              help='projected length of encoder as key')
          parser.add_argument('--fix-projection-length', action='store_true',
                              help='fix projection length for all input sequences')
          parser.add_argument('--share-decoder-input-output-embed', action='store_true',
                              help='share decoder input and output embeddings')
          parser.add_argument('--share-all-embeddings', action='store_true',
                              help='share encoder, decoder and output embeddings'
                                   ' (requires shared dictionary and embed dim)')
          parser.add_argument('--no-token-positional-embeddings', default=False, action='store_true',
                              help='if set, disables positional embeddings (outside self attention)')
          parser.add_argument('--adaptive-softmax-cutoff', metavar='EXPR',
                              help='comma separated list of adaptive softmax cutoff points. '
                                   'Must be used with adaptive_loss criterion'),
          parser.add_argument('--adaptive-softmax-dropout', type=float, metavar='D',
                              help='sets adaptive softmax dropout for the tail projections')
          parser.add_argument('--no-scale-embedding', action='store_true',
                              help='if True, dont scale embeddings')
          parser.add_argument('--layernorm-embedding', action='store_true',
                              help='add layernorm to embedding')
          # args for "Reducing Transformer Depth on Demand with Structured Dropout" (Fan et al., 2019)
          parser.add_argument('--encoder-layerdrop', type=float, metavar='D', default=0,
                              help='LayerDrop probability for encoder')
          parser.add_argument('--decoder-layerdrop', type=float, metavar='D', default=0,
                              help='LayerDrop probability for decoder')
          parser.add_argument('--encoder-layers-to-keep', default=None,
                              help='which layers to *keep* when pruning as a comma-separated list')
          parser.add_argument('--decoder-layers-to-keep', default=None,
                              help='which layers to *keep* when pruning as a comma-separated list')
          # args for Training with Quantization Noise for Extreme Model Compression ({Fan*, Stock*} et al., 2020)
          parser.add_argument('--quant-noise-pq', type=float, metavar='D', default=0,
                              help='iterative PQ quantization noise at training time')
          parser.add_argument('--quant-noise-pq-block-size', type=int, metavar='D', default=8,
                              help='block size of quantization noise at training time')
          parser.add_argument('--quant-noise-scalar', type=float, metavar='D', default=0,
                              help='scalar quantization noise and scalar quantization at training time')
          # fmt: on
      @classmethod
      def from_pretrained(
          cls,
          model_name_or_path,
          checkpoint_file='model.pt',
          data_name_or_path='.',
          bpe='gpt2',
          **kwargs,
      ):
          from fairseq import hub_utils
          x = hub_utils.from_pretrained(
              model_name_or_path,
              checkpoint_file,
              data_name_or_path,
              archive_map=cls.hub_models(),
              bpe=bpe,
              load_checkpoint_heads=True,
              **kwargs,
          )
          return BARTHubInterface(x['args'], x['task'], x['models'][0])
  
      @classmethod
      def build_model(cls, args, task):
          """Build a new model instance."""
  
          # make sure all arguments are present in older models
          base_architecture(args)
  
          if args.encoder_layers_to_keep:
              args.encoder_layers = len(args.encoder_layers_to_keep.split(","))
          if args.decoder_layers_to_keep:
              args.decoder_layers = len(args.decoder_layers_to_keep.split(","))
  
          if getattr(args, "max_source_positions", None) is None:
              args.max_source_positions = DEFAULT_MAX_SOURCE_POSITIONS
          if getattr(args, "max_target_positions", None) is None:
              args.max_target_positions = DEFAULT_MAX_TARGET_POSITIONS
  
          src_dict, tgt_dict = task.source_dictionary, task.target_dictionary
  
          if args.share_all_embeddings:
              if src_dict != tgt_dict:
                  raise ValueError("--share-all-embeddings requires a joined dictionary")
  
              if args.encoder_embed_dim != args.decoder_embed_dim:
                  raise ValueError("--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim")
  
              if args.decoder_embed_path and (args.decoder_embed_path != args.encoder_embed_path):
                  raise ValueError("--share-all-embeddings not compatible with --decoder-embed-path")
  
              encoder_embed_tokens = cls.build_embedding(args, src_dict, args.encoder_embed_dim, args.encoder_embed_path)
              decoder_embed_tokens = encoder_embed_tokens
              args.share_decoder_input_output_embed = True
          else:
              encoder_embed_tokens = cls.build_embedding(args, src_dict, args.encoder_embed_dim, args.encoder_embed_path)
              decoder_embed_tokens = cls.build_embedding(args, tgt_dict, args.decoder_embed_dim, args.decoder_embed_path)
  
          encoder = cls.build_encoder(args, src_dict, encoder_embed_tokens)
          decoder = cls.build_decoder(args, tgt_dict, decoder_embed_tokens)
          return cls(args, encoder, decoder)
  
      @classmethod
      def build_embedding(cls, args, 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, args, src_dict, embed_tokens):
          return LunaEncoder(args, src_dict, embed_tokens)
  
      @classmethod
      def build_decoder(cls, args, tgt_dict, embed_tokens):
          return LunaDecoder(args, tgt_dict, embed_tokens)
  
      # 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 = False,
          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.
          """
          encoder_out = self.encoder(src_tokens, src_lengths=src_lengths, return_all_hiddens=return_all_hiddens)
          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,
          )
          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)
  
  
  class LunaEncoder(FairseqEncoder):
      """
      Luna encoder consisting of *args.encoder_layers* layers. Each layer
      is a :class:`LunaEncoderLayer`.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
          dictionary (~fairseq.data.Dictionary): encoding dictionary
          embed_tokens (torch.nn.Embedding): input embedding
      """
  
      def __init__(self, args, dictionary, embed_tokens):
          super().__init__(dictionary)
          self.register_buffer("version", torch.Tensor([3]))
  
          self.dropout_module = FairseqDropout(args.dropout, module_name=self.__class__.__name__)
          self.dropword_module = FairseqFeatureDropout(args.word_dropout, module_name=self.__class__.__name__)
          self.encoder_layerdrop = args.encoder_layerdrop
  
          embed_dim = embed_tokens.embedding_dim
          assert embed_dim == args.encoder_embed_dim, 'encoder embedding dim mismatch.'
          self.padding_idx = embed_tokens.padding_idx
          self.max_source_positions = args.max_source_positions
  
          self.embed_tokens = embed_tokens
  
          self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
  
          self.embed_positions = (
              PositionalEmbedding(
                  args.max_source_positions,
                  embed_dim,
                  self.padding_idx,
                  learned=args.encoder_learned_pos,
              )
              if not args.no_token_positional_embeddings
              else None
          )
  
          assert not args.layernorm_embedding or not args.encoder_normalize_before
  
          if args.layernorm_embedding:
              self.layernorm_embedding = LayerNorm(embed_dim)
              self.layernorm_porjected_embedding = LayerNorm(embed_dim)
          else:
              self.layernorm_embedding = None
              self.layernorm_porjected_embedding = None
  
          self.proj_len = args.projection_length
          self.dynamic_projection = not args.fix_projection_length
          self.projected_embeddings = Parameter(torch.Tensor(self.proj_len, embed_dim))
          nn.init.normal_(self.projected_embeddings, mean=0., std=embed_dim ** -0.5)
          if not args.no_token_positional_embeddings and not args.encoder_learned_pos:
              projected_positions = get_sinusoidal_positional_embedding(self.proj_len, embed_dim)
          else:
              projected_positions = None
          self.register_buffer('projected_positions', projected_positions)
  
          if self.encoder_layerdrop > 0.0:
              self.layers = LayerDropModuleList(p=self.encoder_layerdrop)
          else:
              self.layers = nn.ModuleList([])
          self.layers.extend([self.build_encoder_layer(i, args) for i in range(args.encoder_layers)])
          self.num_layers = len(self.layers)
  
          if args.encoder_normalize_before:
              self.layer_norm = LayerNorm(embed_dim)
              self.proj_layer_norm = LayerNorm(embed_dim)
          else:
              self.layer_norm = None
              self.proj_layer_norm = None
  
      def build_encoder_layer(self, layer_id, args):
          return LunaEncoderLayer(args, layer_id)
  
      def forward_embedding(self, src_tokens):
          # embed tokens and positions
          x = embed = self.embed_scale * self.embed_tokens(src_tokens)
          x = self.dropword_module(x)
          if self.embed_positions is not None:
              x = x + self.embed_positions(src_tokens)
          if self.layernorm_embedding is not None:
              x = self.layernorm_embedding(x)
  
          return x, embed
  
      def forward_projected_embedding(self, src_lengths):
          max_len = src_lengths.max() if self.dynamic_projection else self.proj_len
          px = proj_embed = self.embed_scale * self.projected_embeddings[:max_len]
          if self.projected_positions is not None:
              px = px + self.projected_positions[:max_len]
          if self.layernorm_porjected_embedding is not None:
              px = self.layernorm_porjected_embedding(px)
  
          return px, proj_embed
  
      def forward(self, src_tokens, src_lengths, return_all_hiddens: bool = False):
          """
          Args:
              src_tokens (LongTensor): tokens in the source language of shape
                  `(batch, src_len)`
              src_lengths (torch.LongTensor): lengths of each source sentence of
                  shape `(batch)`
              return_all_hiddens (bool, optional): also return all of the
                  intermediate hidden states (default: False).
  
          Returns:
              namedtuple:
                  - **encoder_out** (Tensor): the last encoder layer's output of
                    shape `(src_len, batch, embed_dim)`
                  - **encoder_padding_mask** (ByteTensor): the positions of
                    padding elements of shape `(batch, src_len)`
                  - **encoder_embedding** (Tensor): the (scaled) embedding lookup
                    of shape `(batch, src_len, embed_dim)`
                  - **encoder_states** (List[Tensor]): all intermediate
                    hidden states of shape `(src_len, batch, embed_dim)`.
                    Only populated if *return_all_hiddens* is True.
          """
          x, encoder_embedding = self.forward_embedding(src_tokens)
          px, projected_embedding = self.forward_projected_embedding(src_lengths)
  
          bsz = x.size(0)
          len, dim = px.size()
  
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
          # L x C -> L x B x C
          px = px.unsqueeze(1).expand(len, bsz, dim)
  
          x = self.dropout_module(x)
          px = self.dropout_module(px)
  
          # compute padding mask
          encoder_padding_mask = src_tokens.eq(self.padding_idx)
          if self.dynamic_projection:
              pidx= torch.arange(len).unsqueeze(0).to(x.device)
              encoder_projected_padding_mask = pidx.ge(src_lengths.unsqueeze(1))
          else:
              encoder_projected_padding_mask = None
  
          encoder_states = [] if return_all_hiddens else None
  
          # encoder layers
          for layer in self.layers:
              x, px = layer(x, px, encoder_padding_mask, encoder_projected_padding_mask)
              if return_all_hiddens:
                  assert encoder_states is not None
                  encoder_states.append((x, px))
  
          if self.layer_norm is not None:
              x = self.layer_norm(x)
              px = self.proj_layer_norm(px)
  
          return EncoderOut(
              encoder_out=x,  # T x B x C
              encoder_projected_out=px,  # L x B x C
              encoder_padding_mask=encoder_padding_mask,  # B x T
              encoder_projected_padding_mask=encoder_projected_padding_mask,  # B x L
              encoder_embedding=encoder_embedding,  # B x T x C
              encoder_states=encoder_states,  # List[T x B x C]
              src_tokens=None,
              src_lengths=None,
          )
  
      @torch.jit.export
      def reorder_encoder_out(self, encoder_out: EncoderOut, new_order):
          """
          Reorder encoder output according to *new_order*.
  
          Args:
              encoder_out: output from the ``forward()`` method
              new_order (LongTensor): desired order
  
          Returns:
              *encoder_out* rearranged according to *new_order*
          """
          """
          Since encoder_padding_mask and encoder_embedding are both of type
          Optional[Tensor] in EncoderOut, they need to be copied as local
          variables for Torchscript Optional refinement
          """
          encoder_padding_mask: Optional[Tensor] = encoder_out.encoder_padding_mask
          encoder_projected_padding_mask: Optional[Tensor] = encoder_out.encoder_projected_padding_mask
          encoder_embedding: Optional[Tensor] = encoder_out.encoder_embedding
  
          new_encoder_out = (
              encoder_out.encoder_out
              if encoder_out.encoder_out is None
              else encoder_out.encoder_out.index_select(1, new_order)
          )
          new_encoder_projected_out = (
              encoder_out.encoder_projected_out
              if encoder_out.encoder_projected_out is None
              else encoder_out.encoder_projected_out.index_select(1, new_order)
          )
          new_encoder_padding_mask = (
              encoder_padding_mask
              if encoder_padding_mask is None
              else encoder_padding_mask.index_select(0, new_order)
          )
          new_encoder_projected_padding_mask = (
              encoder_projected_padding_mask
              if encoder_projected_padding_mask is None
              else encoder_projected_padding_mask.index_select(0, new_order)
          )
          new_encoder_embedding = (
              encoder_embedding
              if encoder_embedding is None
              else encoder_embedding.index_select(0, new_order)
          )
  
          src_tokens = encoder_out.src_tokens
          if src_tokens is not None:
              src_tokens = src_tokens.index_select(0, new_order)
  
          src_lengths = encoder_out.src_lengths
          if src_lengths is not None:
              src_lengths = src_lengths.index_select(0, new_order)
  
          encoder_states = encoder_out.encoder_states
          if encoder_states is not None:
              for idx, state in enumerate(encoder_states):
                  encoder_states[idx] = state.index_select(1, new_order)
  
          return EncoderOut(
              encoder_out=new_encoder_out,  # T x B x C
              encoder_projected_out=new_encoder_projected_out, # L x B x C
              encoder_padding_mask=new_encoder_padding_mask,  # B x T
              encoder_projected_padding_mask=new_encoder_projected_padding_mask,  # B x L
              encoder_embedding=new_encoder_embedding,  # B x T x C
              encoder_states=encoder_states,  # List[T x B x C]
              src_tokens=src_tokens,  # B x T
              src_lengths=src_lengths,  # B x 1
          )
  
      def max_positions(self):
          """Maximum input length supported by the encoder."""
          if self.embed_positions is None:
              return self.max_source_positions
          return min(self.max_source_positions, self.embed_positions.max_positions)
  
      def upgrade_state_dict_named(self, state_dict, name):
          """Upgrade a (possibly old) state dict for new versions of fairseq."""
          if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
              weights_key = "{}.embed_positions.weights".format(name)
              if weights_key in state_dict:
                  print("deleting {0}".format(weights_key))
                  del state_dict[weights_key]
              state_dict[
                  "{}.embed_positions._float_tensor".format(name)
              ] = torch.FloatTensor(1)
          for i in range(self.num_layers):
              # update layer norms
              self.layers[i].upgrade_state_dict_named(
                  state_dict, "{}.layers.{}".format(name, i)
              )
  
          version_key = "{}.version".format(name)
          if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) < 2:
              # earlier checkpoints did not normalize after the stack of layers
              self.layer_norm = None
              self.normalize = False
              state_dict[version_key] = torch.Tensor([1])
          return state_dict
  
  
  class LunaDecoder(FairseqIncrementalDecoder):
      """
      Luna decoder consisting of *args.decoder_layers* layers. Each layer
      is a :class:`LunaDecoderLayer`.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
          dictionary (~fairseq.data.Dictionary): decoding dictionary
          embed_tokens (torch.nn.Embedding): output embedding
      """
  
      def __init__(self, args, dictionary, embed_tokens):
          self.args = args
          super().__init__(dictionary)
          self.register_buffer("version", torch.Tensor([3]))
          self._future_mask = torch.empty(0)
  
          self.dropout_module = FairseqDropout(args.dropout, module_name=self.__class__.__name__)
          self.dropword_module = FairseqFeatureDropout(args.word_dropout, module_name=self.__class__.__name__)
          self.decoder_layerdrop = args.decoder_layerdrop
          self.share_input_output_embed = args.share_decoder_input_output_embed
  
          input_embed_dim = embed_tokens.embedding_dim
          embed_dim = args.decoder_embed_dim
          self.output_embed_dim = args.decoder_output_dim
  
          self.padding_idx = embed_tokens.padding_idx
          self.max_target_positions = args.max_target_positions
  
          self.embed_tokens = embed_tokens
  
          self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)
  
          self.project_in_dim = (
              Linear(input_embed_dim, embed_dim, bias=False)
              if embed_dim != input_embed_dim
              else None
          )
  
          self.embed_positions = (
              PositionalEmbedding(
                  args.max_target_positions,
                  embed_dim,
                  self.padding_idx,
                  learned=args.decoder_learned_pos,
              )
              if not args.no_token_positional_embeddings
              else None
          )
  
          assert not args.layernorm_embedding or not args.decoder_normalize_before
  
          if args.layernorm_embedding:
              self.layernorm_embedding = LayerNorm(embed_dim)
          else:
              self.layernorm_embedding = None
  
          if self.decoder_layerdrop > 0.0:
              self.layers = LayerDropModuleList(p=self.decoder_layerdrop)
          else:
              self.layers = nn.ModuleList([])
          self.layers.extend([self.build_decoder_layer(i, args) for i in range(args.decoder_layers)])
          self.num_layers = len(self.layers)
  
          if args.decoder_normalize_before:
              self.layer_norm = LayerNorm(embed_dim)
              self.proj_layer_norm = LayerNorm(embed_dim)
          else:
              self.layer_norm = None
              self.proj_layer_norm = None
  
          self.project_out_dim = (
              Linear(embed_dim, self.output_embed_dim, bias=False)
              if embed_dim != self.output_embed_dim and not args.tie_adaptive_weights else None
          )
  
          self.adaptive_softmax = None
          self.output_projection = None
          if args.adaptive_softmax_cutoff is not None:
              self.adaptive_softmax = AdaptiveSoftmax(
                  len(dictionary),
                  self.output_embed_dim,
                  options.eval_str_list(args.adaptive_softmax_cutoff, type=int),
                  dropout=args.adaptive_softmax_dropout,
                  adaptive_inputs=embed_tokens if args.tie_adaptive_weights else None,
                  factor=args.adaptive_softmax_factor,
                  tie_proj=args.tie_adaptive_proj,
              )
          elif self.share_input_output_embed:
              self.output_projection = nn.Linear(
                  self.embed_tokens.weight.shape[1],
                  self.embed_tokens.weight.shape[0],
                  bias=False,
              )
              self.output_projection.weight = self.embed_tokens.weight
          else:
              self.output_projection = nn.Linear(
                  self.output_embed_dim, len(dictionary), bias=False
              )
              nn.init.normal_(
                  self.output_projection.weight, mean=0, std=self.output_embed_dim ** -0.5
              )
  
      def build_decoder_layer(self, layer_id, args):
          return LunaDecoderLayer(args, layer_id)
  
      def forward(
          self,
          prev_output_tokens,
          encoder_out: Optional[EncoderOut] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          features_only: bool = False,
          alignment_layer: Optional[int] = None,
          alignment_heads: Optional[int] = None,
          src_lengths: Optional[Any] = None,
          return_all_hiddens: bool = False,
      ):
          """
          Args:
              prev_output_tokens (LongTensor): previous decoder outputs of shape
                  `(batch, tgt_len)`, for teacher forcing
              encoder_out (optional): output from the encoder, used for
                  encoder-side attention
              incremental_state (dict): dictionary used for storing state during
                  :ref:`Incremental decoding`
              features_only (bool, optional): only return features without
                  applying output layer (default: False).
  
          Returns:
              tuple:
                  - the decoder's output of shape `(batch, tgt_len, vocab)`
                  - a dictionary with any model-specific outputs
          """
          x, extra = self.extract_features(
              prev_output_tokens,
              encoder_out=encoder_out,
              incremental_state=incremental_state,
              alignment_layer=alignment_layer,
              alignment_heads=alignment_heads,
          )
          if not features_only:
              x = self.output_layer(x)
          return x, extra
  
      def extract_features(
          self,
          prev_output_tokens,
          encoder_out: Optional[EncoderOut] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          full_context_alignment: bool = False,
          alignment_layer: Optional[int] = None,
          alignment_heads: Optional[int] = None,
      ):
          return self.extract_features_scriptable(
              prev_output_tokens,
              encoder_out,
              incremental_state,
              full_context_alignment,
              alignment_layer,
              alignment_heads,
          )
  
      """
      A scriptable subclass of this class has an extract_features method and calls
      super().extract_features, but super() is not supported in torchscript. Aa copy of
      this function is made to be used in the subclass instead.
      """
  
      def extract_features_scriptable(
          self,
          prev_output_tokens,
          encoder_out: Optional[EncoderOut] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          full_context_alignment: bool = False,
          alignment_layer: Optional[int] = None,
          alignment_heads: Optional[int] = None,
      ):
          """
          Similar to *forward* but only return features.
  
          Includes several features from "Jointly Learning to Align and
          Translate with Transformer Models" (Garg et al., EMNLP 2019).
  
          Args:
              full_context_alignment (bool, optional): don't apply
                  auto-regressive mask to self-attention (default: False).
              alignment_layer (int, optional): return mean alignment over
                  heads at this layer (default: last layer).
              alignment_heads (int, optional): only average alignment over
                  this many heads (default: all heads).
  
          Returns:
              tuple:
                  - the decoder's features of shape `(batch, tgt_len, embed_dim)`
                  - a dictionary with any model-specific outputs
          """
          if alignment_layer is None:
              alignment_layer = self.num_layers - 1
  
          # embed positions
          positions = (
              self.embed_positions(prev_output_tokens, incremental_state=incremental_state)
              if self.embed_positions is not None
              else None
          )
  
          static_px = False
          if incremental_state is not None:
              prev_output_tokens = prev_output_tokens[:, -1:]
              if positions is not None:
                  positions = positions[:, -1:]
              projected_input_buffer = self._get_projected_input_buffer(incremental_state)
              if "prev_projected_input" in projected_input_buffer:
                  static_px = True
          else:
              projected_input_buffer = None
  
          # embed tokens and positions
          x = self.embed_tokens(prev_output_tokens) * self.embed_scale
          x = self.dropword_module(x)
          if self.project_in_dim is not None:
              x = self.project_in_dim(x)
          if positions is not None:
              x = x + positions
  
          if self.layernorm_embedding is not None:
              x = self.layernorm_embedding(x)
  
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
          x = self.dropout_module(x)
  
          if not static_px:
              # L x B x C
              px = encoder_out.encoder_projected_out
          else:
              px = None
  
          self_attn_padding_mask: Optional[Tensor] = None
          if prev_output_tokens.eq(self.padding_idx).any():
              self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
  
          # decoder layers
          attn: Optional[Tensor] = None
          inner_states: List[Optional[Tensor]] = [x]
          for idx, layer in enumerate(self.layers):
              x, px, layer_attn, _ = layer(x, px,
                                           encoder_out.encoder_out if encoder_out is not None else None,
                                           encoder_out.encoder_padding_mask if encoder_out is not None else None,
                                           encoder_out.encoder_projected_padding_mask if encoder_out is not None else None,
                                           incremental_state,
                                           self_attn_padding_mask=self_attn_padding_mask,
                                           need_attn=bool(idx == alignment_layer),
                                           need_head_weights=bool(idx == alignment_layer))
              inner_states.append(x)
              if layer_attn is not None and idx == alignment_layer:
                  attn = layer_attn.float().to(x)
  
          if attn is not None:
              if alignment_heads is not None:
                  attn = attn[:alignment_heads]
  
              # average probabilities over heads
              attn = attn.mean(dim=0)
  
          if self.layer_norm is not None:
              x = self.layer_norm(x)
              if not static_px:
                  px = self.proj_layer_norm(px)
  
          # T x B x C -> B x T x C
          x = x.transpose(0, 1)
          # L x B x C -> B x L x C
          if not static_px:
              px = px.transpose(0, 1)
              if projected_input_buffer is not None:
                  projected_input_buffer["prev_projected_input"] = px
                  incremental_state = self._set_projected_input_buffer(incremental_state, projected_input_buffer)
          else:
              px = projected_input_buffer["prev_projected_input"]
  
          if self.project_out_dim is not None:
              x = self.project_out_dim(x)
  
          return x, {"attn": [attn], "inner_states": inner_states}
  
      def _get_projected_input_buffer(
          self,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
      ) -> Dict[str, Optional[Tensor]]:
          result = self.get_incremental_state(incremental_state, "px_state")
          if result is not None:
              return result
          else:
              empty_result: Dict[str, Optional[Tensor]] = {}
              return empty_result
  
      def _set_projected_input_buffer(
          self,
          incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
          buffer: Dict[str, Optional[Tensor]],
      ):
          return self.set_incremental_state(incremental_state, "px_state", buffer)
  
      def output_layer(self, features):
          """Project features to the vocabulary size."""
          if self.adaptive_softmax is None:
              # project back to size of vocabulary
              return self.output_projection(features)
          else:
              return features
  
      def max_positions(self):
          """Maximum output length supported by the decoder."""
          if self.embed_positions is None:
              return self.max_target_positions
          return min(self.max_target_positions, self.embed_positions.max_positions)
  
      def buffered_future_mask(self, tensor):
          dim = tensor.size(0)
          # self._future_mask.device != tensor.device is not working in TorchScript. This is a workaround.
          if (
              self._future_mask.size(0) == 0
              or (not self._future_mask.device == tensor.device)
              or self._future_mask.size(0) < dim
          ):
              self._future_mask = torch.triu(
                  utils.fill_with_neg_inf(torch.zeros([dim, dim])), 1
              )
          self._future_mask = self._future_mask.to(tensor)
          return self._future_mask[:dim, :dim]
  
      def upgrade_state_dict_named(self, state_dict, name):
          """Upgrade a (possibly old) state dict for new versions of fairseq."""
          if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
              weights_key = "{}.embed_positions.weights".format(name)
              if weights_key in state_dict:
                  del state_dict[weights_key]
              state_dict[
                  "{}.embed_positions._float_tensor".format(name)
              ] = torch.FloatTensor(1)
  
          if f"{name}.output_projection.weight" not in state_dict:
              if self.share_input_output_embed:
                  embed_out_key = f"{name}.embed_tokens.weight"
              else:
                  embed_out_key = f"{name}.embed_out"
              if embed_out_key in state_dict:
                  state_dict[f"{name}.output_projection.weight"] = state_dict[
                      embed_out_key
                  ]
                  if not self.share_input_output_embed:
                      del state_dict[embed_out_key]
  
          for i in range(self.num_layers):
              # update layer norms
              layer_norm_map = {
                  "0": "self_attn_layer_norm",
                  "1": "encoder_attn_layer_norm",
                  "2": "final_layer_norm",
              }
              for old, new in layer_norm_map.items():
                  for m in ("weight", "bias"):
                      k = "{}.layers.{}.layer_norms.{}.{}".format(name, i, old, m)
                      if k in state_dict:
                          state_dict[
                              "{}.layers.{}.{}.{}".format(name, i, new, m)
                          ] = state_dict[k]
                          del state_dict[k]
  
          return state_dict
  
  
  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
  
  
  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
  
  
  def get_sinusoidal_positional_embedding(length, embed_dim):
      half_dim = embed_dim // 2
      emb = math.log(10000) / (half_dim - 1)
      emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
      emb = torch.arange(length, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
      emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(length, -1)
      if embed_dim % 2 == 1:
          emb = torch.cat([emb, torch.zeros(length, 1)], dim=1)
      return emb
  
  
  @register_model_architecture("luna", "luna")
  def base_architecture(args):
      args.encoder_embed_path = getattr(args, "encoder_embed_path", None)
      args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
      args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
      args.encoder_layers = getattr(args, "encoder_layers", 6)
      args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
      args.encoder_projected_attention_heads = getattr(args, "encoder_projected_attention_heads", args.encoder_attention_heads)
      args.encoder_learned_pos = getattr(args, "encoder_learned_pos", False)
      args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
  
      args.decoder_embed_path = getattr(args, "decoder_embed_path", None)
      args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
      args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim)
      args.decoder_layers = getattr(args, "decoder_layers", 6)
      args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
      args.decoder_projected_attention_heads = getattr(args, "decoder_projected_attention_heads", args.decoder_attention_heads)
      args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
      args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
  
      args.projection_length = getattr(args, 'projection_length', 32)
      args.fix_projection_length = getattr(args, "fix_projection_length", False)
  
      args.attention_dropout = getattr(args, "attention_dropout", 0.0)
      args.activation_dropout = getattr(args, "activation_dropout", 0.0)
      args.activation_fn = getattr(args, "activation_fn", "relu")
      args.dropout = getattr(args, "dropout", 0.1)
      args.word_dropout = getattr(args, "word_dropout", 0.0)
      args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
      args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
      args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
  
      args.share_decoder_input_output_embed = getattr(args, "share_decoder_input_output_embed", False)
      args.share_all_embeddings = getattr(args, "share_all_embeddings", False)
      args.no_token_positional_embeddings = getattr( args, "no_token_positional_embeddings", False)
      args.adaptive_input = getattr(args, "adaptive_input", False)
  
      args.decoder_output_dim = getattr(args, "decoder_output_dim", args.decoder_embed_dim)
      args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
  
      args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
      args.tie_adaptive_weights = getattr(args, "tie_adaptive_weights", False)
  