






import torch
from torch import nn






class 	UpperCAmelCase__		(	nn.Module	):






  """simple docstring"""


  def __init__(							self						:		List[str]     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		Tuple     ,       __lowerCamelCase						:		Union[str, Any]=1     ,       __lowerCamelCase						:		Union[str, Any]=False )		->							str:
       super().__init__()

       SCREAMING_SNAKE_CASE__   =   n_token
       SCREAMING_SNAKE_CASE__   =   d_embed
       SCREAMING_SNAKE_CASE__   =   d_proj

       SCREAMING_SNAKE_CASE__   =   cutoffs + [n_token]
       SCREAMING_SNAKE_CASE__   =   [0] + self.cutoffs
       SCREAMING_SNAKE_CASE__   =   div_val

       SCREAMING_SNAKE_CASE__   =   self.cutoffs[0]
       SCREAMING_SNAKE_CASE__   =   len(self.cutoffs ) - 1
       SCREAMING_SNAKE_CASE__   =   self.shortlist_size + self.n_clusters

       if self.n_clusters > 0:
            SCREAMING_SNAKE_CASE__   =   nn.Parameter(torch.zeros(self.n_clusters     ,       self.d_embed ) )
            SCREAMING_SNAKE_CASE__   =   nn.Parameter(torch.zeros(self.n_clusters ) )

       SCREAMING_SNAKE_CASE__   =   nn.ModuleList()
       SCREAMING_SNAKE_CASE__   =   nn.ParameterList()

       if div_val == 1:
            for i in range(len(self.cutoffs ) ):
                 if d_proj != d_embed:
                      self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCamelCase     ,       __lowerCamelCase ) ) )
                 else:
                      self.out_projs.append(__lowerCamelCase )

            self.out_layers.append(nn.Linear(__lowerCamelCase     ,       __lowerCamelCase ) )
       else:
            for i in range(len(self.cutoffs ) ):
                 SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.cutoff_ends[i], self.cutoff_ends[i + 1]
                 SCREAMING_SNAKE_CASE__   =   d_embed // (div_val**i)

                 self.out_projs.append(nn.Parameter(torch.FloatTensor(__lowerCamelCase     ,       __lowerCamelCase ) ) )

                 self.out_layers.append(nn.Linear(__lowerCamelCase     ,       r_idx - l_idx ) )

       SCREAMING_SNAKE_CASE__   =   keep_order






  def 				lowercase_		(							self						:		int     ,       __lowerCamelCase						:		List[Any]     ,       __lowerCamelCase						:		List[Any]     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		List[str] )		->							List[str]:
       if proj is None:
            SCREAMING_SNAKE_CASE__   =   nn.functional.linear(__lowerCamelCase     ,       __lowerCamelCase     ,       bias=__lowerCamelCase )
       else:
            # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
            SCREAMING_SNAKE_CASE__   =   nn.functional.linear(__lowerCamelCase     ,       proj.t().contiguous() )
            SCREAMING_SNAKE_CASE__   =   nn.functional.linear(__lowerCamelCase     ,       __lowerCamelCase     ,       bias=__lowerCamelCase )
            # else:
            #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
            #     if bias is not None:
            #         logit = logit + bias

       return logit






  def 				lowercase_		(							self						:		Any     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		Union[str, Any]=None     ,       __lowerCamelCase						:		Dict=False )		->							List[Any]:

       if labels is not None:
            # Shift so that tokens < n predict n
            SCREAMING_SNAKE_CASE__   =   hidden[..., :-1, :].contiguous()
            SCREAMING_SNAKE_CASE__   =   labels[..., 1:].contiguous()
            SCREAMING_SNAKE_CASE__   =   hidden.view(-1     ,       hidden.size(-1 ) )
            SCREAMING_SNAKE_CASE__   =   labels.view(-1 )
            if hidden.size(0 ) != labels.size(0 ):
                 raise RuntimeError('''Input and labels should have the same size in the batch dimension.''' )
       else:
            SCREAMING_SNAKE_CASE__   =   hidden.view(-1     ,       hidden.size(-1 ) )

       if self.n_clusters == 0:
            SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       self.out_layers[0].weight     ,       self.out_layers[0].bias     ,       self.out_projs[0] )
            if labels is not None:
                 SCREAMING_SNAKE_CASE__   =   labels != -100
                 SCREAMING_SNAKE_CASE__   =   torch.zeros_like(__lowerCamelCase     ,       dtype=hidden.dtype     ,       device=hidden.device )
                 SCREAMING_SNAKE_CASE__   =   (
                     -nn.functional.log_softmax(__lowerCamelCase     ,       dim=-1 )[mask].gather(1     ,       labels[mask].unsqueeze(1 ) ).squeeze(1 )
                 )
            else:
                 SCREAMING_SNAKE_CASE__   =   nn.functional.log_softmax(__lowerCamelCase     ,       dim=-1 )
       else:
            # construct weights and biases
            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   [], []
            for i in range(len(self.cutoffs ) ):
                 if self.div_val == 1:
                      SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.cutoff_ends[i], self.cutoff_ends[i + 1]
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[0].weight[l_idx:r_idx]
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[0].bias[l_idx:r_idx]
                 else:
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[i].weight
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[i].bias

                 if i == 0:
                      SCREAMING_SNAKE_CASE__   =   torch.cat([weight_i, self.cluster_weight]     ,       dim=0 )
                      SCREAMING_SNAKE_CASE__   =   torch.cat([bias_i, self.cluster_bias]     ,       dim=0 )

                 weights.append(__lowerCamelCase )
                 biases.append(__lowerCamelCase )

            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   weights[0], biases[0], self.out_projs[0]

            SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   nn.functional.log_softmax(__lowerCamelCase     ,       dim=1 )

            if labels is None:
                 SCREAMING_SNAKE_CASE__   =   hidden.new_empty((head_logit.size(0 ), self.n_token) )
            else:
                 SCREAMING_SNAKE_CASE__   =   torch.zeros_like(__lowerCamelCase     ,       dtype=hidden.dtype     ,       device=hidden.device )

            SCREAMING_SNAKE_CASE__   =   0
            SCREAMING_SNAKE_CASE__   =   [0] + self.cutoffs
            for i in range(len(__lowerCamelCase ) - 1 ):
                 SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   cutoff_values[i], cutoff_values[i + 1]

                 if labels is not None:
                      SCREAMING_SNAKE_CASE__   =   (labels >= l_idx) & (labels < r_idx)
                      SCREAMING_SNAKE_CASE__   =   mask_i.nonzero().squeeze()

                      if indices_i.numel() == 0:
                           continue

                      SCREAMING_SNAKE_CASE__   =   labels.index_select(0     ,       __lowerCamelCase ) - l_idx
                      SCREAMING_SNAKE_CASE__   =   head_logprob.index_select(0     ,       __lowerCamelCase )
                      SCREAMING_SNAKE_CASE__   =   hidden.index_select(0     ,       __lowerCamelCase )
                 else:
                      SCREAMING_SNAKE_CASE__   =   hidden

                 if i == 0:
                      if labels is not None:
                           SCREAMING_SNAKE_CASE__   =   head_logprob_i.gather(1     ,       target_i[:, None] ).squeeze(1 )
                      else:
                           SCREAMING_SNAKE_CASE__   =   head_logprob[:, : self.cutoffs[0]]
                 else:
                      SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   weights[i], biases[i], self.out_projs[i]

                      SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )
                      SCREAMING_SNAKE_CASE__   =   nn.functional.log_softmax(__lowerCamelCase     ,       dim=1 )
                      SCREAMING_SNAKE_CASE__   =   self.cutoffs[0] + i - 1  # No probability for the head cluster
                      if labels is not None:
                           SCREAMING_SNAKE_CASE__   =   head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                               1     ,       target_i[:, None] ).squeeze(1 )
                      else:
                           SCREAMING_SNAKE_CASE__   =   head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                           SCREAMING_SNAKE_CASE__   =   logprob_i

                 if labels is not None:
                      if (hasattr(self     ,       '''keep_order''' ) and self.keep_order) or keep_order:
                           out.index_copy_(0     ,       __lowerCamelCase     ,       -logprob_i )
                      else:
                           out[offset : offset + logprob_i.size(0 )].copy_(-logprob_i )
                      offset += logprob_i.size(0 )

       return out







  def 				lowercase_		(							self						:		Optional[int]     ,       __lowerCamelCase						:		Union[str, Any] )		->							Tuple:
       if self.n_clusters == 0:
            SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       self.out_layers[0].weight     ,       self.out_layers[0].bias     ,       self.out_projs[0] )
            return nn.functional.log_softmax(__lowerCamelCase     ,       dim=-1 )
       else:
            # construct weights and biases
            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   [], []
            for i in range(len(self.cutoffs ) ):
                 if self.div_val == 1:
                      SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.cutoff_ends[i], self.cutoff_ends[i + 1]
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[0].weight[l_idx:r_idx]
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[0].bias[l_idx:r_idx]
                 else:
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[i].weight
                      SCREAMING_SNAKE_CASE__   =   self.out_layers[i].bias

                 if i == 0:
                      SCREAMING_SNAKE_CASE__   =   torch.cat([weight_i, self.cluster_weight]     ,       dim=0 )
                      SCREAMING_SNAKE_CASE__   =   torch.cat([bias_i, self.cluster_bias]     ,       dim=0 )

                 weights.append(__lowerCamelCase )
                 biases.append(__lowerCamelCase )

            SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   weights[0], biases[0], self.out_projs[0]
            SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )

            SCREAMING_SNAKE_CASE__   =   hidden.new_empty((head_logit.size(0 ), self.n_token) )
            SCREAMING_SNAKE_CASE__   =   nn.functional.log_softmax(__lowerCamelCase     ,       dim=1 )

            SCREAMING_SNAKE_CASE__   =   [0] + self.cutoffs
            for i in range(len(__lowerCamelCase ) - 1 ):
                 SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   cutoff_values[i], cutoff_values[i + 1]

                 if i == 0:
                      SCREAMING_SNAKE_CASE__   =   head_logprob[:, : self.cutoffs[0]]
                 else:
                      SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   weights[i], biases[i], self.out_projs[i]

                      SCREAMING_SNAKE_CASE__   =   self._compute_logit(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )
                      SCREAMING_SNAKE_CASE__   =   nn.functional.log_softmax(__lowerCamelCase     ,       dim=1 )

                      SCREAMING_SNAKE_CASE__   =   head_logprob[:, -i] + tail_logprob_i
                      SCREAMING_SNAKE_CASE__   =   logprob_i

            return out





from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE					:  List[str]             =      logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE					:  Tuple             =      {
    '''transfo-xl-wt103''': '''https://huggingface.co/transfo-xl-wt103/resolve/main/config.json''',
}






class 	UpperCAmelCase__		(	A__	):






  """simple docstring"""


  a									= "transfo-xl"
  a									= ["mems"]
  a									= {
      "n_token": "vocab_size",
      "hidden_size": "d_model",
      "num_attention_heads": "n_head",
      "num_hidden_layers": "n_layer",
  }






  def __init__(							self						:		Any     ,       __lowerCamelCase						:		int=26_7735     ,       __lowerCamelCase						:		Any=[2_0000, 4_0000, 20_0000]     ,       __lowerCamelCase						:		Dict=1024     ,       __lowerCamelCase						:		Optional[int]=1024     ,       __lowerCamelCase						:		List[Any]=16     ,       __lowerCamelCase						:		Union[str, Any]=64     ,       __lowerCamelCase						:		Dict=4096     ,       __lowerCamelCase						:		int=4     ,       __lowerCamelCase						:		Dict=False     ,       __lowerCamelCase						:		Tuple=18     ,       __lowerCamelCase						:		Optional[int]=1600     ,       __lowerCamelCase						:		str=1000     ,       __lowerCamelCase						:		Any=True     ,       __lowerCamelCase						:		Dict=True     ,       __lowerCamelCase						:		List[Any]=0     ,       __lowerCamelCase						:		int=-1     ,       __lowerCamelCase						:		Tuple=True     ,       __lowerCamelCase						:		List[Any]=0.1     ,       __lowerCamelCase						:		int=0.0     ,       __lowerCamelCase						:		int=True     ,       __lowerCamelCase						:		str="normal"     ,       __lowerCamelCase						:		List[str]=0.01     ,       __lowerCamelCase						:		Any=0.01     ,       __lowerCamelCase						:		List[str]=0.02     ,       __lowerCamelCase						:		List[str]=1e-5     ,       __lowerCamelCase						:		Union[str, Any]=0     ,       **__lowerCamelCase						:		int     ,       )		->							Optional[Any]:
       SCREAMING_SNAKE_CASE__   =   vocab_size
       SCREAMING_SNAKE_CASE__   =   []
       self.cutoffs.extend(__lowerCamelCase )
       if proj_share_all_but_first:
            SCREAMING_SNAKE_CASE__   =   [False] + [True] * len(self.cutoffs )
       else:
            SCREAMING_SNAKE_CASE__   =   [False] + [False] * len(self.cutoffs )
       SCREAMING_SNAKE_CASE__   =   d_model
       SCREAMING_SNAKE_CASE__   =   d_embed
       SCREAMING_SNAKE_CASE__   =   d_head
       SCREAMING_SNAKE_CASE__   =   d_inner
       SCREAMING_SNAKE_CASE__   =   div_val
       SCREAMING_SNAKE_CASE__   =   pre_lnorm
       SCREAMING_SNAKE_CASE__   =   n_layer
       SCREAMING_SNAKE_CASE__   =   n_head
       SCREAMING_SNAKE_CASE__   =   mem_len
       SCREAMING_SNAKE_CASE__   =   same_length
       SCREAMING_SNAKE_CASE__   =   attn_type
       SCREAMING_SNAKE_CASE__   =   clamp_len
       SCREAMING_SNAKE_CASE__   =   sample_softmax
       SCREAMING_SNAKE_CASE__   =   adaptive
       SCREAMING_SNAKE_CASE__   =   dropout
       SCREAMING_SNAKE_CASE__   =   dropatt
       SCREAMING_SNAKE_CASE__   =   untie_r
       SCREAMING_SNAKE_CASE__   =   init
       SCREAMING_SNAKE_CASE__   =   init_range
       SCREAMING_SNAKE_CASE__   =   proj_init_std
       SCREAMING_SNAKE_CASE__   =   init_std
       SCREAMING_SNAKE_CASE__   =   layer_norm_epsilon
       super().__init__(eos_token_id=__lowerCamelCase     ,       **__lowerCamelCase )






  @property
  def 				lowercase_		(							self						:		str )		->							Dict:
       # Message copied from Transformer-XL documentation
       logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )
       return -1







  @max_position_embeddings.setter
  def 				lowercase_		(							self						:		List[str]     ,       __lowerCamelCase						:		Any )		->							List[Any]:
       # Message copied from Transformer-XL documentation
       raise NotImplementedError(
           f'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )


