text: import math
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from fairseq import options, utils
  from fairseq.models import (
      FairseqEncoder,
      FairseqIncrementalDecoder,
      FairseqEncoderDecoderModel,
      register_model,
      register_model_architecture,
  )
  from fairseq.modules import (
      MultiheadAttention,
      LayerNorm,
  )
  from fairseq.modules.transformer_sentence_encoder import init_bert_params
  from .learned_positional_embedding import LearnedPositionalEmbedding
  
  DEFAULT_MAX_SOURCE_POSITIONS = 512
  DEFAULT_MAX_TARGET_POSITIONS = 512
  
  
  @register_model('transformer_mass')
  class TransformerMASSModel(FairseqEncoderDecoderModel):
      """
      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, encoder, decoder):
          super().__init__(encoder, decoder)
  
      @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('--attention-dropout', type=float, metavar='D',
                              help='dropout probability for attention weights')
          parser.add_argument('--activation-dropout', type=float, metavar='D',
                              help='dropout probability after activation in FFN.')
  
          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('--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('--share-all-embeddings', action='store_true',
                              help='share encoder, decoder and output embeddings'
                                   ' (requires shared dictionary and embed dim)')
          parser.add_argument('--load-from-pretrained-model', type=str, default=None,
                              help='Load from pretrained model')
          # fmt: on
  
      @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 not hasattr(args, 'max_source_positions'):
              args.max_source_positions = DEFAULT_MAX_SOURCE_POSITIONS
          if not hasattr(args, 'max_target_positions'):
              args.max_target_positions = DEFAULT_MAX_TARGET_POSITIONS
  
          src_dict, tgt_dict = task.source_dictionary, task.target_dictionary
  
  
          def build_embedding(dictionary, embed_dim):
              num_embeddings = len(dictionary)
              padding_idx = dictionary.pad()
              emb = Embedding(num_embeddings, embed_dim, padding_idx)
              return emb
  
          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')
              encoder_embed_tokens = build_embedding(
                  src_dict, args.encoder_embed_dim
              )
              decoder_embed_tokens = encoder_embed_tokens
              args.share_decoder_input_output_embed = True
          else:
              encoder_embed_tokens = build_embedding(
                  src_dict, args.encoder_embed_dim
              )
              decoder_embed_tokens = build_embedding(
                  tgt_dict, args.decoder_embed_dim
              )
  
          encoder = TransformerEncoder(args, src_dict, encoder_embed_tokens)
          decoder = TransformerDecoder(args, tgt_dict, decoder_embed_tokens)
  
          model = TransformerMASSModel(encoder, decoder)
  
          if args.load_from_pretrained_model is not None:
              states = torch.load(args.load_from_pretrained_model, map_location='cpu')
              model.load_state_dict(states)
              args.load_from_pretrained_model = None # Clear this param
  
          return TransformerMASSModel(encoder, decoder)
  
      def max_positions(self):
          return min(self.encoder.max_positions(), self.decoder.max_positions())
  
      def forward(self, src_tokens=None, src_lengths=None, prev_output_tokens=None, **kwargs):
          """
          Run the forward pass for an encoder-decoder model.
  
          First feed a batch of source tokens through the encoder. Then, feed the
          encoder output and previous decoder outputs (i.e., teacher forcing) to
          the decoder to produce the next outputs::
  
              encoder_out = self.encoder(src_tokens, src_lengths)
              return self.decoder(prev_output_tokens, encoder_out)
  
          Args:
              src_tokens (LongTensor): tokens in the source language of shape
                  `(batch, src_len)`
              src_lengths (LongTensor): source sentence lengths of shape `(batch)`
              prev_output_tokens (LongTensor): previous decoder outputs of shape
                  `(batch, tgt_len)`, for teacher forcing
  
          Returns:
              tuple:
                  - the decoder's output of shape `(batch, tgt_len, vocab)`
                  - a dictionary with any model-specific outputs
          """
          encoder_out = self.encoder(src_tokens, src_lengths=src_lengths, **kwargs)
          decoder_out = self.decoder(prev_output_tokens, encoder_out=encoder_out, **kwargs)
          return decoder_out
  
  
  class TransformerEncoderLayer(nn.Module):
      """
      Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
      models.
      """
  
      def __init__(
          self,
          embedding_dim: float = 768,
          ffn_embedding_dim: float = 3072,
          num_attention_heads: float = 8,
          dropout: float = 0.1,
          attention_dropout: float = 0.1,
          activation_dropout: float = 0.1,
          activation_fn: str = 'relu',
          add_bias_kv: bool = False,
          add_zero_attn: bool = False,
          export: bool = False,
      ) -> None:
  
          super().__init__()
          # Initialize parameters
          self.embedding_dim = embedding_dim
          self.dropout = dropout
          self.activation_dropout = activation_dropout
  
          # Initialize blocks
          self.activation_fn = utils.get_activation_fn(activation_fn)
          self.self_attn = MultiheadAttention(
              self.embedding_dim,
              num_attention_heads,
              dropout=attention_dropout,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
              self_attention=True
          )
  
          # layer norm associated with the self attention layer
          self.self_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)
          self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
          self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
  
          # layer norm associated with the position wise feed-forward NN
          self.final_layer_norm = LayerNorm(self.embedding_dim, export=export)
  
      def forward(
          self,
          x: torch.Tensor,
          self_attn_mask: torch.Tensor = None,
          self_attn_padding_mask: torch.Tensor = None,
      ):
          """
          LayerNorm is applied either before or after the self-attention/ffn
          modules similar to the original Transformer imlementation.
          """
          residual = x
          x, attn = self.self_attn(
              query=x,
              key=x,
              value=x,
              key_padding_mask=self_attn_padding_mask,
              need_weights=False,
              attn_mask=self_attn_mask,
          )
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          x = self.self_attn_layer_norm(x)
  
          residual = x
          x = self.activation_fn(self.fc1(x))
          x = F.dropout(x, p=self.activation_dropout, training=self.training)
          x = self.fc2(x)
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          x = self.final_layer_norm(x)
          return x, attn
  
  
  class TransformerDecoderLayer(nn.Module):
      def __init__(
          self,
          embedding_dim: float = 768,
          ffn_embedding_dim: float = 3072,
          num_attention_heads: float = 8,
          dropout: float = 0.1,
          attention_dropout: float = 0.1,
          activation_dropout: float = 0.1,
          activation_fn: str = 'relu',
          add_bias_kv: bool = False,
          add_zero_attn: bool = False,
          export: bool = False,
      ):
          super().__init__()
  
          self.embedding_dim = embedding_dim
          self.dropout = dropout
          self.activation_dropout = activation_dropout
  
          # Initialize blocks
          self.activation_fn = utils.get_activation_fn(activation_fn)
          self.self_attn = MultiheadAttention(
              self.embedding_dim,
              num_attention_heads,
              dropout=attention_dropout,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
              self_attention=True
          )
  
          # layer norm associated with the self attention layer
          self.self_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)
  
          self.encoder_attn = MultiheadAttention(
              self.embedding_dim,
              num_attention_heads,
              kdim=embedding_dim,
              vdim=embedding_dim,
              dropout=attention_dropout,
              encoder_decoder_attention=True,
          )
          self.encoder_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)
  
          self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
          self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
  
          # layer norm associated with the position wise feed-forward NN
          self.final_layer_norm = LayerNorm(self.embedding_dim, export=export)
          self.need_attn = False
          
      def forward(
          self,
          x,
          encoder_out=None,
          encoder_mask=None,
          incremental_state=None,
          prev_self_attn_state=None,
          prev_attn_state=None,
          self_attn_mask=None,
          self_attn_padding_mask=None,
      ):
          residual = x
          if prev_self_attn_state is not None:
              if incremental_state is None:
                  incremental_state = {}
              prev_key, prev_value = prev_self_attn_state
              saved_state = {"prev_key": prev_key, "prev_value": prev_value}
              self.self_attn._set_input_buffer(incremental_state, saved_state)
  
          x, attn = self.self_attn(
              query=x,
              key=x,
              value=x,
              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
          x = self.self_attn_layer_norm(x)
  
          residual = x
          if prev_attn_state is not None:
              if incremental_state is None:
                  incremental_state = {}
              prev_key, prev_value = prev_attn_state
              saved_state = {"prev_key": prev_key, "prev_value": prev_value}
              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_mask,
              incremental_state=incremental_state,
              static_kv=True,
              need_weights=(not self.training and self.need_attn),
          )
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          x = self.encoder_attn_layer_norm(x)
  
          residual = x
          x = self.activation_fn(self.fc1(x))
          x = F.dropout(x, p=self.activation_dropout, training=self.training)
          x = self.fc2(x)
          x = F.dropout(x, p=self.dropout, training=self.training)
          x = residual + x
          x = self.final_layer_norm(x)
          return x, attn
      
      def make_generation_fast_(self, need_attn=False, **kwargs):
          self.need_attn = need_attn
  
  
  class TransformerEncoder(FairseqEncoder):
      """
      Transformer encoder consisting of *args.encoder_layers* layers. Each layer
      is a :class:`TransformerEncoderLayer`.
  
      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 = args.dropout
  
          embed_dim = embed_tokens.embedding_dim
          self.padding_idx = embed_tokens.padding_idx
          self.max_source_positions = args.max_source_positions
  
          self.embed_tokens = embed_tokens
          self.embed_scale = math.sqrt(embed_dim)
          self.embed_positions = LearnedPositionalEmbedding(
              args.max_source_positions + 1 + self.padding_idx, embed_dim, self.padding_idx,
          )
  
          self.layers = nn.ModuleList([])
  
          self.layers.extend([
              TransformerEncoderLayer(
                  args.encoder_embed_dim,
                  args.encoder_ffn_embed_dim,
                  args.encoder_attention_heads,
                  args.dropout,
                  args.attention_dropout,
                  args.activation_dropout,
                  args.activation_fn,
              )
              for i in range(args.encoder_layers)
          ])
  
          self.emb_layer_norm = LayerNorm(embed_dim)
  
          self.apply(init_bert_params)
  
      def forward(self, src_tokens, src_lengths, **unused):
          """
          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)`
  
          Returns:
              dict:
                  - **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)`
          """
          # compute padding mask
          encoder_padding_mask = src_tokens.eq(self.padding_idx)
          if not encoder_padding_mask.any():
              encoder_padding_mask = None
  
          # embed tokens and positions
          x = self.embed_scale * self.embed_tokens(src_tokens)
          if self.embed_positions is not None:
              x += self.embed_positions(src_tokens)
  
          if self.emb_layer_norm:
              x = self.emb_layer_norm(x)
  
          x = F.dropout(x, p=self.dropout, training=self.training)
          
          if encoder_padding_mask is not None:
              x *= 1 - encoder_padding_mask.unsqueeze(-1).type_as(x)
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
  
          # encoder layers
          for layer in self.layers:
              x, _ = layer(x, self_attn_padding_mask=encoder_padding_mask)
  
          return {
              'encoder_out': x,  # T x B x C
              'encoder_padding_mask': encoder_padding_mask,  # B x T
          }
  
      def reorder_encoder_out(self, encoder_out, 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*
          """
          if encoder_out['encoder_out'] is not None:
              encoder_out['encoder_out'] = \
                  encoder_out['encoder_out'].index_select(1, new_order)
          if encoder_out['encoder_padding_mask'] is not None:
              encoder_out['encoder_padding_mask'] = \
                  encoder_out['encoder_padding_mask'].index_select(0, new_order)
          return encoder_out
  
      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())
  
  
  class TransformerDecoder(FairseqIncrementalDecoder):
      """
      Transformer decoder consisting of *args.decoder_layers* layers. Each layer
      is a :class:`TransformerDecoderLayer`.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
          dictionary (~fairseq.data.Dictionary): decoding dictionary
          embed_tokens (torch.nn.Embedding): output embedding
          no_encoder_attn (bool, optional): whether to attend to encoder outputs
              (default: False).
      """
  
      def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
          super().__init__(dictionary)
          self.register_buffer('version', torch.Tensor([3]))
  
          self.dropout = args.dropout
          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.padding_idx = embed_tokens.padding_idx
          self.max_target_positions = args.max_target_positions
          self.embed_dim = embed_dim
          self.embed_tokens = embed_tokens
          self.embed_scale = math.sqrt(embed_dim)  # todo: try with input_embed_dim
          
          self.embed_positions = LearnedPositionalEmbedding(
              args.max_target_positions + 1 + self.padding_idx, embed_dim, self.padding_idx,
          )
  
          self.layers = nn.ModuleList([])
          self.layers.extend([
              TransformerDecoderLayer(
                  args.decoder_embed_dim,
                  args.decoder_ffn_embed_dim,
                  args.decoder_attention_heads,
                  args.dropout,
                  args.attention_dropout,
                  args.activation_dropout,
                  args.activation_fn,
              )
              for _ in range(args.decoder_layers)    
          ])
  
          if not self.share_input_output_embed:
              self.embed_out = nn.Parameter(torch.Tensor(len(dictionary), self.embed_dim))
              nn.init.normal_(self.embed_out, mean=0, std=self.embed_dim ** -0.5)
  
          self.emb_layer_norm = LayerNorm(embed_dim)
          self.apply(init_bert_params)
  
      def forward(self,
                  prev_output_tokens, 
                  encoder_out=None, 
                  incremental_state=None, 
                  **unused):
          x, extra = self.extract_features(prev_output_tokens, encoder_out, incremental_state, **unused)
          x = self.output_layer(x)
          return x, extra
  
      def extract_features(self, prev_output_tokens, encoder_out=None, incremental_state=None, **unused):
          # embed positions
          if 'positions' in unused:
              positions = self.embed_positions._forward(unused['positions'])
          else:
              positions = self.embed_positions(
                  prev_output_tokens,
                  incremental_state=incremental_state,
              ) if self.embed_positions is not None else None
  
          if incremental_state is not None:
              prev_output_tokens = prev_output_tokens[:, -1:]
              if positions is not None:
                  positions = positions[:, -1:]
  
          # embed tokens and positions
          x = self.embed_scale * self.embed_tokens(prev_output_tokens)
  
          if positions is not None:
              x += positions
  
          if self.emb_layer_norm:
              x = self.emb_layer_norm(x)
  
          x = F.dropout(x, p=self.dropout, training=self.training)
  
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
          attn = None
  
          inner_states = [x]
  
          # decoder layers
          for layer in self.layers:
              x, attn = layer(
                  x,
                  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,
                  incremental_state,
                  self_attn_mask=self.buffered_future_mask(x) if incremental_state is None else None,
              )
              inner_states.append(x)
  
          # T x B x C -> B x T x C
          x = x.transpose(0, 1)
  
          return x, {'attn': attn, 'inner_states': inner_states}
  
      def output_layer(self, features, **kwargs):
          """Project features to the vocabulary size."""
          # project back to size of vocabulary
          if self.share_input_output_embed:
              return F.linear(features, self.embed_tokens.weight)
          else:
              return F.linear(features, self.embed_out)
  
      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)
          if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device or self._future_mask.size(0) < dim:
              self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
          return self._future_mask[:dim, :dim]
  
  
  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.)
      return m
  
  
  @register_model_architecture('transformer_mass', 'transformer_mass')
  def base_architecture(args):
      args.activation_fn = getattr(args, 'activation_fn', 'relu')
      args.dropout = getattr(args, 'dropout', 0.1)
      args.attention_dropout = getattr(args, 'attention_dropout', 0.)
      args.activation_dropout = getattr(args, 'activation_dropout', 0.)
  
      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.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 512)
      args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 2048)
      args.decoder_layers = getattr(args, 'decoder_layers', 6)
      args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 8)
  
      args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', False)
      args.share_all_embeddings = getattr(args, 'share_all_embeddings', False)
  
  @register_model_architecture('transformer_mass', 'transformer_mass_base')
  def transformer_base(args):
      args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768)
      args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 3072)
      args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 12)
      args.encoder_layers = getattr(args, 'encoder_layers', 6)
  
      args.dropout = getattr(args, 'dropout', 0.1)
      args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
      args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
      args.activation_fn = getattr(args, 'activation_fn', 'gelu')
  
      args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 768)
      args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 3072)
      args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 12)
      args.decoder_layers = getattr(args, 'decoder_layers', 6)
  
      args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
      args.share_all_embeddings = getattr(args, 'share_all_embeddings', True)
      base_architecture(args)
  
  
  @register_model_architecture('transformer_mass', 'transformer_mass_big')
  def transformer_big(args):
      args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
      args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4096)
      args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
      args.encoder_layers = getattr(args, 'encoder_layers', 12)
  
      args.dropout = getattr(args, 'dropout', 0.1)
      args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
      args.activation_dropout = getattr(args, 'activation_dropout', 0.1)
      args.activation_fn = getattr(args, 'activation_fn', 'gelu')
  
      args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', 1024)
      args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', 4096)
      args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
      args.decoder_layers = getattr(args, 'decoder_layers', 12)
  
      args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
      args.share_all_embeddings = getattr(args, 'share_all_embeddings', True)
      base_architecture(args)
  