text: from dataclasses import dataclass, field
  import math
  from typing import Any, Dict, List, Optional, Tuple
  from examples.adaptive_span.adaptive_span_model_wrapper import AdaptiveSpanSmallConfig
  
  from fairseq import options, utils
  from fairseq.dataclass import ChoiceEnum, FairseqDataclass
  from fairseq.models import (
      FairseqLanguageModel,
      FairseqIncrementalDecoder,
      register_model,
      register_model_architecture,
  )
  from fairseq.models.transformer import (
      DEFAULT_MIN_PARAMS_TO_WRAP, Embedding, TransformerDecoder
  )
  from fairseq.models.transformer_lm import TransformerLanguageModel, base_lm_architecture
  from fairseq.modules import (
      AdaptiveInput, 
      CharacterTokenEmbedder,
      AdaptiveSoftmax,
      BaseLayer,
      FairseqDropout,
      LayerDropModuleList,
      LayerNorm,
      PositionalEmbedding,
      SinusoidalPositionalEmbedding,
      TransformerDecoderLayer,
      TransformerEncoderLayer
  )
  from omegaconf import II
  import numpy as np
  import torch
  from torch import Tensor
  
  
  DEFAULT_MAX_TARGET_POSITIONS = 1024
  
  @register_model("transformer_numlm")
  class TransformerNumLanguageModel(TransformerLanguageModel):
      def __init__(self, decoder):
          super().__init__(decoder)
  
      @staticmethod
      def add_args(parser):
          TransformerLanguageModel.add_args(parser)
          parser.add_argument('--add-number-token-attention-mask', action="store_true", default=False, help="add self-attention masking to force inner-digits attention")
          parser.add_argument('--add-number-rotation-embedding', action="store_true", default=False, help="add rotation embedding to the <bon> based on number value")
  
      @classmethod
      def build_model(cls, args, task):
          """Build a new model instance."""
  
          if args.decoder_layers_to_keep:
              args.decoder_layers = len(args.decoder_layers_to_keep.split(","))
  
          if getattr(args, "max_target_positions", None) is None:
              args.max_target_positions = getattr(
                  args, "tokens_per_sample", DEFAULT_MAX_TARGET_POSITIONS
              )
  
          if args.character_embeddings:
              embed_tokens = CharacterTokenEmbedder(
                  task.source_dictionary,
                  eval(args.character_filters),
                  args.character_embedding_dim,
                  args.decoder_embed_dim,
                  args.char_embedder_highway_layers,
              )
          elif args.adaptive_input:
              embed_tokens = AdaptiveInput(
                  len(task.source_dictionary),
                  task.source_dictionary.pad(),
                  args.decoder_input_dim,
                  args.adaptive_input_factor,
                  args.decoder_embed_dim,
                  options.eval_str_list(args.adaptive_input_cutoff, type=int),
                  args.quant_noise_pq,
                  args.quant_noise_pq_block_size,
              )
          else:
              embed_tokens = cls.build_embedding(
                  args, task.source_dictionary, args.decoder_input_dim
              )
  
          if args.tie_adaptive_weights:
              assert args.adaptive_input
              assert args.adaptive_input_factor == args.adaptive_softmax_factor
              assert (
                  args.adaptive_softmax_cutoff == args.adaptive_input_cutoff
              ), "{} != {}".format(
                  args.adaptive_softmax_cutoff, args.adaptive_input_cutoff
              )
              assert args.decoder_input_dim == args.decoder_output_dim
  
          decoder = TransformerNumDecoder(
              args, task.target_dictionary, embed_tokens, no_encoder_attn=True
          )
          return cls(decoder)
  
      @classmethod
      def build_embedding(cls, args, dictionary, embed_dim, path=None):
          embed_tokens = Embedding(len(dictionary), embed_dim, dictionary.pad())
          return embed_tokens
  
  
  class TransformerNumDecoder(TransformerDecoder):
      """
      Transformer number 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,
          output_projection=None,
      ):
          super().__init__(args, dictionary, embed_tokens, no_encoder_attn, output_projection)        
          self.add_number_token_attention_mask = getattr(args, "add_number_token_attention_mask", False)
          self.add_number_rotation_embedding = getattr(args, "add_number_rotation_embedding", False)
          self.bon_idx = dictionary.indices['<bon>']
          self.eon_idx = dictionary.indices['<eon>']
  
          if self.add_number_rotation_embedding:
              self.rotation_weight = torch.nn.Parameter(torch.zeros(self.embed_dim, requires_grad=True))
              self.rotation_bias = torch.nn.Parameter(torch.zeros(self.embed_dim, requires_grad=True))
              torch.nn.init.normal_(self.rotation_weight.data, mean=0, std=self.embed_dim ** -0.5)
              torch.nn.init.normal_(self.rotation_bias.data, mean=0, std=self.embed_dim ** -0.5)
  
  
      def forward(
          self,
          prev_output_tokens,
          encoder_out: Optional[Dict[str, List[Tensor]]] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          features_only: bool = False,
          full_context_alignment: bool = False,
          alignment_layer: Optional[int] = None,
          alignment_heads: Optional[int] = None,
          src_lengths: Optional[Any] = None,
          src_number_token_values: Optional[Any] = None,
          src_numlm_mask: 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, should be of size T x B x C
              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).
              full_context_alignment (bool, optional): don't apply
                  auto-regressive mask to self-attention (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,
              full_context_alignment=full_context_alignment,
              alignment_layer=alignment_layer,
              alignment_heads=alignment_heads,
              src_number_token_values=src_number_token_values,
              src_numlm_mask=src_numlm_mask
          )
  
          if not features_only:
              x = self.output_layer(x)
          return x, extra
  
      def extract_features(
          self,
          prev_output_tokens,
          encoder_out: Optional[Dict[str, List[Tensor]]],
          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,
          src_number_token_values: Optional[Any] = None,
          src_numlm_mask: Optional[Any] = None
      ):
          return self.extract_features_scriptable(
              prev_output_tokens,
              encoder_out,
              incremental_state,
              full_context_alignment,
              alignment_layer,
              alignment_heads,
              src_number_token_values=src_number_token_values,
              src_numlm_mask=src_numlm_mask
          )
  
      """
      A scriptable subclass of this class has an extract_features method and calls
      super().extract_features, but super() is not supported in torchscript. A copy of
      this function is made to be used in the subclass instead.
      """
  
      def extract_features_scriptable(
          self,
          prev_output_tokens,
          encoder_out: Optional[Dict[str, List[Tensor]]],
          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,
          src_number_token_values: Optional[Any] = None,
          src_numlm_mask: Optional[Any] = 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
          """
          bs, slen = prev_output_tokens.size()
          if alignment_layer is None:
              alignment_layer = self.num_layers - 1
  
          enc: Optional[Tensor] = None
          padding_mask: Optional[Tensor] = None
          if encoder_out is not None and len(encoder_out["encoder_out"]) > 0:
              enc = encoder_out["encoder_out"][0]
              assert (
                  enc.size()[1] == bs
              ), f"Expected enc.shape == (t, {bs}, c) got {enc.shape}"
          if encoder_out is not None and len(encoder_out["encoder_padding_mask"]) > 0:
              padding_mask = encoder_out["encoder_padding_mask"][0]
  
          # embed positions
          positions = None
          if self.embed_positions is not None:
              positions = self.embed_positions(
                  prev_output_tokens, incremental_state=incremental_state
              )
  
          if incremental_state is not None:
              prev_output_tokens = prev_output_tokens[:, -1:]
              if positions is not None:
                  positions = positions[:, -1:]
  
  
          # combine tokens, numbers and positions embeds
          word_embed = self.embed_scale * self.embed_tokens(prev_output_tokens)
          
          if self.add_number_rotation_embedding:
              number_rotation_embed = self.prepare_number_rotation_embed(src_number_token_values, word_embed.shape, word_embed.dtype, word_embed.device)
              x = word_embed * number_rotation_embed # element-wise multiplication, only <bon>-embeddings have non-1 entries 
          else:
              x = word_embed
  
          if self.quant_noise is not None:
              x = self.quant_noise(x)
  
          if self.project_in_dim is not None:
              x = self.project_in_dim(x)
  
          if positions is not None:
              x += positions
  
          if self.layernorm_embedding is not None:
              x = self.layernorm_embedding(x)
  
          x = self.dropout_module(x)
  
          # B x T x C -> T x B x C
          x = x.transpose(0, 1)
  
          self_attn_padding_mask: Optional[Tensor] = None
          if self.cross_self_attention or prev_output_tokens.eq(self.padding_idx).any():
              self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
  
          if self.add_number_token_attention_mask:
              # add number token inner attention mask to the self_attn_mask
              src_numlm_mask = torch.repeat_interleave(src_numlm_mask, self.args.decoder_attention_heads, dim=0).to(device=word_embed.device)
              assert list(src_numlm_mask.size()) == [bs*self.args.decoder_attention_heads, slen, slen]
  
          # decoder layers
          attn: Optional[Tensor] = None
          inner_states: List[Optional[Tensor]] = [x]
          for idx, layer in enumerate(self.layers):
              if incremental_state is None and not full_context_alignment:
                  self_attn_mask = self.buffered_future_mask(x)
                  if self.add_number_token_attention_mask:
                      self_attn_mask = self_attn_mask + src_numlm_mask # src_attn_mask broadcast to bs*num_heads
              else:
                  self_attn_mask = None
  
              x, layer_attn, _ = layer(
                  x,
                  enc,
                  padding_mask,
                  incremental_state,
                  self_attn_mask=self_attn_mask,
                  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)
  
          # T x B x C -> B x T x C
          x = x.transpose(0, 1)
  
          if self.project_out_dim is not None:
              x = self.project_out_dim(x)
  
          return x, {"attn": [attn], "inner_states": inner_states}
  
  
      def get_rotation_embed(self, value):
          # right now -1 and 1 have the same rotation embedding for now.
          # TODO: differentiate a number and its negative
          assert math.isfinite(value), "[transformer_numlm] value is not finite!"
          # rotation_embed = self.rotation_weight * math.log(abs(value)) + self.rotation_bias
          # value has been math.log(abs(number.value)) in collate_fn to prevent --fp16 cast number too large to inf
          rotation_embed = self.rotation_weight * value + self.rotation_bias
          rotation_embed_ = torch.zeros_like(rotation_embed)
          assert rotation_embed_.shape[0] % 2 == 0, "[tranformer_numlm] Embedding dimension size is odd."
          cos_mask = torch.LongTensor([1, 0]).repeat(int(rotation_embed_.shape[0] // 2)).to(device=rotation_embed.device)
          sin_mask = torch.LongTensor([0, 1]).repeat(int(rotation_embed_.shape[0] // 2)).to(device=rotation_embed.device)
          rotation_embed_ = torch.sin(rotation_embed) * sin_mask + torch.cos(rotation_embed) * cos_mask
          return rotation_embed_
  
      
      def prepare_number_rotation_embed(self, src_number_token_values, _shape, _dtype, _device):
          '''
              in:
                  prev_output_tokens: dim - [NUM_SENT, SEQ_LEN]
              out:
                  number_embed: dim - same as word_embed
          '''
          # only <bon> tokens are embeded
          number_rotation_embed = torch.ones(_shape, dtype=_dtype, device=_device)
          for sent_i, bon_pos, value in src_number_token_values:
              number_rotation_embed[int(sent_i)][int(bon_pos)] = self.get_rotation_embed(value)
          return number_rotation_embed
      
  
  
  @register_model_architecture("transformer_numlm", "transformer_numlm_base")
  def transformer_numlm_base(args):
      args.add_number_token_attention_mask = getattr(args, "add_number_token_attention_mask", False)
      base_lm_architecture(args)
  