

'''simple docstring'''



import torch
from torch import nn



class 				_snake_case					(   nn.Module			):

  def __init__(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=1      , _SCREAMING_SNAKE_CASE=False	):





      '''simple docstring'''

      super().__init__()

      lowerCAmelCase 						=							n_token
      lowerCAmelCase 						=							d_embed
      lowerCAmelCase 						=							d_proj

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

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

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

      lowerCAmelCase 						=							nn.ModuleList()
      lowerCAmelCase 						=							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	)	):
              lowerCAmelCase 						=							self.cutoff_ends[i], self.cutoff_ends[i + 1]
              lowerCAmelCase 						=							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	)	)

      lowerCAmelCase 						=							keep_order

  def   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      if proj is None:
          lowerCAmelCase 						=							nn.functional.linear(__lowerCAmelCase      , __lowerCAmelCase      , bias=__lowerCAmelCase	)
      else:
          # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
          lowerCAmelCase 						=							nn.functional.linear(__lowerCAmelCase      , proj.t().contiguous()	)
          lowerCAmelCase 						=							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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE      , _SCREAMING_SNAKE_CASE=None      , _SCREAMING_SNAKE_CASE=False	):





      '''simple docstring'''

      if labels is not None:
          # Shift so that tokens < n predict n
          lowerCAmelCase 						=							hidden[..., :-1, :].contiguous()
          lowerCAmelCase 						=							labels[..., 1:].contiguous()
          lowerCAmelCase 						=							hidden.view(-1      , hidden.size(-1	)	)
          lowerCAmelCase 						=							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:
          lowerCAmelCase 						=							hidden.view(-1      , hidden.size(-1	)	)

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

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

              weights.append(__lowerCAmelCase	)
              biases.append(__lowerCAmelCase	)

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

          lowerCAmelCase 						=							self._compute_logit(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase	)
          lowerCAmelCase 						=							nn.functional.log_softmax(__lowerCAmelCase      , dim=1	)

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

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

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

                  if indices_i.numel() == 0:
                      continue

                  lowerCAmelCase 						=							labels.index_select(0      , __lowerCAmelCase	) - l_idx
                  lowerCAmelCase 						=							head_logprob.index_select(0      , __lowerCAmelCase	)
                  lowerCAmelCase 						=							hidden.index_select(0      , __lowerCAmelCase	)
              else:
                  lowerCAmelCase 						=							hidden

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

                  lowerCAmelCase 						=							self._compute_logit(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase	)
                  lowerCAmelCase 						=							nn.functional.log_softmax(__lowerCAmelCase      , dim=1	)
                  lowerCAmelCase 						=							self.cutoffs[0] + i - 1  # No probability for the head cluster
                  if labels is not None:
                      lowerCAmelCase 						=							head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                          1      , target_i[:, None]	).squeeze(1	)
                  else:
                      lowerCAmelCase 						=							head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                      lowerCAmelCase 						=							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   _SCREAMING_SNAKE_CASE				(				self      , _SCREAMING_SNAKE_CASE	):





      '''simple docstring'''

      if self.n_clusters == 0:
          lowerCAmelCase 						=							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
          lowerCAmelCase 						=							[], []
          for i in range(len(self.cutoffs	)	):
              if self.div_val == 1:
                  lowerCAmelCase 						=							self.cutoff_ends[i], self.cutoff_ends[i + 1]
                  lowerCAmelCase 						=							self.out_layers[0].weight[l_idx:r_idx]
                  lowerCAmelCase 						=							self.out_layers[0].bias[l_idx:r_idx]
              else:
                  lowerCAmelCase 						=							self.out_layers[i].weight
                  lowerCAmelCase 						=							self.out_layers[i].bias

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

              weights.append(__lowerCAmelCase	)
              biases.append(__lowerCAmelCase	)

          lowerCAmelCase 						=							weights[0], biases[0], self.out_projs[0]
          lowerCAmelCase 						=							self._compute_logit(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase	)

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

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

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

                  lowerCAmelCase 						=							self._compute_logit(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase	)
                  lowerCAmelCase 						=							nn.functional.log_softmax(__lowerCAmelCase      , dim=1	)

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

          return out













import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE__   :      str				=		logging.get_logger(__name__)


SCREAMING_SNAKE_CASE__   :      Tuple				=		"""▁"""

SCREAMING_SNAKE_CASE__   :      Union[str, Any]				=		{"""vocab_file""": """spiece.model"""}

SCREAMING_SNAKE_CASE__   :      List[Any]				=		{
    """vocab_file""": {
        """google/reformer-crime-and-punishment""": (
            """https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/spiece.model"""
        )
    }
}

SCREAMING_SNAKE_CASE__   :      Optional[int]				=		{
    """google/reformer-crime-and-punishment""": 52_42_88,
}
class     lowerCamelCase_					(   lowerCamelCase			):
				a__       							=	VOCAB_FILES_NAMES
				a__       							=	PRETRAINED_VOCAB_FILES_MAP
				a__       							=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				a__       							=	['''input_ids''', '''attention_mask''']







				def __init__(    self							,       __lowerCAmelCase							,       __lowerCAmelCase="</s>"							,       __lowerCAmelCase="<unk>"							,       __lowerCAmelCase=[]							,       __lowerCAmelCase = None							,       **__lowerCAmelCase							,       ):
								"""simple docstring"""
								__magic_name__		:int   =   {} if sp_model_kwargs is None else sp_model_kwargs

								super().__init__(
								    eos_token=__lowerCAmelCase							,       unk_token=__lowerCAmelCase							,       additional_special_tokens=__lowerCAmelCase							,       sp_model_kwargs=self.sp_model_kwargs							,       **__lowerCAmelCase							,       )

								__magic_name__		:Optional[Any]   =   vocab_file
								__magic_name__		:int   =   spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.Load(__lowerCAmelCase  )







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return self.sp_model.get_piece_size()







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:str   =   {self.convert_ids_to_tokens(__lowerCAmelCase  ): i for i in range(self.vocab_size  )}
								vocab.update(self.added_tokens_encoder  )
								return vocab







				def __getstate__(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   self.__dict__.copy()
								__magic_name__		:Optional[Any]   =   None
								return state







				def __setstate__(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Any   =   d

								# for backward compatibility
								if not hasattr(self							,       '''sp_model_kwargs'''  ):
												__magic_name__		:Optional[int]   =   {}

								__magic_name__		:Union[str, Any]   =   spm.SentencePieceProcessor(**self.sp_model_kwargs  )
								self.sp_model.Load(self.vocab_file  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								return self.sp_model.encode(__lowerCAmelCase							,       out_type=__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								return self.sp_model.piece_to_id(__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								if index < self.sp_model.get_piece_size():
												__magic_name__		:int   =   self.sp_model.IdToPiece(__lowerCAmelCase  )
								return token







				def       A						(    self							,       __lowerCAmelCase  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   []
								__magic_name__		:Tuple   =   ''''''
								for token in tokens:
												# make sure that special tokens are not decoded using sentencepiece model
												if token in self.all_special_tokens:
																out_string += self.sp_model.decode(__lowerCAmelCase  ) + token
																__magic_name__		:Optional[Any]   =   []
												else:
																current_sub_tokens.append(__lowerCAmelCase  )
								out_string += self.sp_model.decode(__lowerCAmelCase  )
								return out_string.strip()







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None  ):
								"""simple docstring"""
								if not os.path.isdir(__lowerCAmelCase  ):
												logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''  )
												return
								__magic_name__		:Optional[int]   =   os.path.join(
								    __lowerCAmelCase							,       (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']  )

								if os.path.abspath(self.vocab_file  ) != os.path.abspath(__lowerCAmelCase  ) and os.path.isfile(self.vocab_file  ):
												copyfile(self.vocab_file							,       __lowerCAmelCase  )
								elif not os.path.isfile(self.vocab_file  ):
												with open(__lowerCAmelCase							,       '''wb'''  ) as fi:
																__magic_name__		:Dict   =   self.sp_model.serialized_model_proto()
																fi.write(__lowerCAmelCase  )

								return (out_vocab_file,)






