text: import torch
  from fairseq import utils
  from fairseq.models.transformer import TransformerModel, TransformerEncoder, TransformerDecoder, base_architecture
  from fairseq.models import (
      register_model,
      register_model_architecture,
  )
  
  DEFAULT_MAX_SOURCE_POSITIONS = 1024
  DEFAULT_MAX_TARGET_POSITIONS = 1024
  
  
  @register_model("our_transformer")
  class OurTransformerModel(TransformerModel):
      """
      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, args, encoder, decoder):
          super().__init__(args, encoder, decoder)
  
      @staticmethod
      def add_args(parser):
          """Add model-specific arguments to the parser."""
          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', '--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-normalize-before', action='store_true',
                              help='apply layernorm before each encoder block')
          parser.add_argument('--encoder-learned-pos', action='store_true',
                              help='use learned positional embeddings in the encoder')
          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-learned-pos', action='store_true',
                              help='use learned positional embeddings in the decoder')
          parser.add_argument('--decoder-normalize-before', action='store_true',
                              help='apply layernorm before each decoder block')
          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('--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('--layernorm-embedding', action='store_true',
                              help='add layernorm to embedding')
          parser.add_argument('--no-scale-embedding', action='store_true',
                              help='if True, dont scale embeddings')
          parser.add_argument('--checkpoint-activations', action='store_true',
                              help='checkpoint activations at each layer, which saves GPU '
                                   'memory usage at the cost of some additional compute')
          # args for "Cross+Self-Attention for Transformer Models" (Peitz et al., 2019)
          parser.add_argument('--no-cross-attention', default=False, action='store_true',
                              help='do not perform cross-attention')
          parser.add_argument('--cross-self-attention', default=False, action='store_true',
                              help='perform cross+self-attention')
          # 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')
          # args for printing score
          parser.add_argument('--print-attn-score', default=False, action="store_true",
                              help="Whether print attention score ")
          # fmt: on
  
      @classmethod
      def build_model(cls, args, task):
          """Build a new model instance."""
  
          # make sure all arguments are present in older models
          transformer_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_decoder(cls, args, tgt_dict, embed_tokens):
          return OurTransformerDecoder(
              args,
              tgt_dict,
              embed_tokens,
              no_encoder_attn=getattr(args, "no_cross_attention", False),
          )
  
  
  class OurTransformerDecoder(TransformerDecoder):
      """
      Transformer decoder consisting of *args.decoder_layers* layers. Each layer
      is a :class:`SparseTransformerDecoderLayer`.
      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__(args, dictionary=dictionary, embed_tokens=embed_tokens, no_encoder_attn=no_encoder_attn)
          self.print_attn_score = args.print_attn_score
  
      def forward(
              self,
              prev_output_tokens,
              encoder_out=None,
              incremental_state=None,
              features_only: bool = False,
              full_context_alignment: bool = False,
              alignment_layer=None,
              alignment_heads=None,
              src_lengths=None,
              return_all_hiddens=False,
      ):
          x, extra = self.extract_features(
              prev_output_tokens,
              encoder_out=encoder_out,
              incremental_state=incremental_state,
              full_context_alignment=full_context_alignment,
              alignment_layer=alignment_layer,
              alignment_heads=alignment_heads,
          )
          if self.print_attn_score:
              print(extra['attn'])
          if not features_only:
              x = self.output_layer(x)
          return x, extra
  
  
  @register_model_architecture("our_transformer", "our_transformer")
  def transformer_base_architecture(args):
      args.print_attn_score = getattr(args, "print_attn_score", False)
      base_architecture(args)
  
  
  @register_model_architecture('our_transformer', 'transformer_base')
  def transformer_base(args):
      args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
      args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 1024)
      args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 4)
      args.encoder_layers = getattr(args, "encoder_layers", 4)
      args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 512)
      args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 1024)
      args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 4)
      args.decoder_layers = getattr(args, "decoder_layers", 4)
      transformer_base_architecture(args)
  
  
  @register_model_architecture('our_transformer', 'transformer_large')
  def transformer_large(args):
      args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
      args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 2048)
      args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 8)
      args.encoder_layers = getattr(args, "encoder_layers", 6)
      args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 512)
      args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 2048)
      args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
      args.decoder_layers = getattr(args, "decoder_layers", 6)
      transformer_base_architecture(args)
  