





import torch
from torch import nn







class 				__lowerCAmelCase       (    nn.Module):







      def __init__(    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__=1		,      lowerCAmelCase__=False							)   ->						Union[str, Any]:



        '''simple docstring'''




        super().__init__()

        a__	:					Tuple										=n_token
        a__	:					Optional[Any]										=d_embed
        a__	:					Union[str, Any]										=d_proj

        a__	:					int										=cutoffs + [n_token]
        a__	:					Any										=[0] + self.cutoffs
        a__	:					Any										=div_val

        a__	:					Any										=self.cutoffs[0]
        a__	:					str										=len(self.cutoffs							) - 1
        a__	:					Any										=self.shortlist_size + self.n_clusters

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

        a__	:					Dict										=nn.ModuleList()
        a__	:					str										=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							)							):
            a__					,							a__	:					str										=self.cutoff_ends[i], self.cutoff_ends[i + 1]
            a__	:					Optional[int]										=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							)							)

        a__	:					str										=keep_order







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Any:



        '''simple docstring'''




        if proj is None:
          a__	:					List[str]										=nn.functional.linear(lowerCAmelCase__		,      lowerCAmelCase__		,      bias=lowerCAmelCase__							)
        else:
          # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
          a__	:					List[str]										=nn.functional.linear(lowerCAmelCase__		,      proj.t().contiguous()							)
          a__	:					Any										=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		,      lowerCAmelCase__		,      lowerCAmelCase__=None		,      lowerCAmelCase__=False							)   ->						Dict:



        '''simple docstring'''




        if labels is not None:
          # Shift so that tokens < n predict n
          a__	:					List[Any]										=hidden[..., :-1, :].contiguous()
          a__	:					Optional[Any]										=labels[..., 1:].contiguous()
          a__	:					Tuple										=hidden.view(-1		,      hidden.size(-1							)							)
          a__	:					Dict										=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:
          a__	:					Union[str, Any]										=hidden.view(-1		,      hidden.size(-1							)							)

        if self.n_clusters == 0:
          a__	:					List[Any]										=self._compute_logit(lowerCAmelCase__		,      self.out_layers[0].weight		,      self.out_layers[0].bias		,      self.out_projs[0]							)
          if labels is not None:
            a__	:					List[Any]										=labels != -1_0_0
            a__	:					Tuple										=torch.zeros_like(lowerCAmelCase__		,      dtype=hidden.dtype		,      device=hidden.device							)
            a__	:					Optional[int]										=(
                -nn.functional.log_softmax(lowerCAmelCase__		,      dim=-1							)[mask].gather(1		,      labels[mask].unsqueeze(1							)							).squeeze(1							)
            )
          else:
            a__	:					List[Any]										=nn.functional.log_softmax(lowerCAmelCase__		,      dim=-1							)
        else:
          # construct weights and biases
          a__					,							a__	:					Optional[int]										=[], []
          for i in range(len(self.cutoffs							)							):
            if self.div_val == 1:
              a__					,							a__	:					Union[str, Any]										=self.cutoff_ends[i], self.cutoff_ends[i + 1]
              a__	:					Union[str, Any]										=self.out_layers[0].weight[l_idx:r_idx]
              a__	:					Dict										=self.out_layers[0].bias[l_idx:r_idx]
            else:
              a__	:					List[str]										=self.out_layers[i].weight
              a__	:					Optional[int]										=self.out_layers[i].bias

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

            weights.append(lowerCAmelCase__							)
            biases.append(lowerCAmelCase__							)

          a__					,							a__					,							a__	:					int										=weights[0], biases[0], self.out_projs[0]

          a__	:					Optional[int]										=self._compute_logit(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)
          a__	:					Tuple										=nn.functional.log_softmax(lowerCAmelCase__		,      dim=1							)

          if labels is None:
            a__	:					int										=hidden.new_empty((head_logit.size(0							), self.n_token)							)
          else:
            a__	:					Any										=torch.zeros_like(lowerCAmelCase__		,      dtype=hidden.dtype		,      device=hidden.device							)

          a__	:					Optional[Any]										=0
          a__	:					int										=[0] + self.cutoffs
          for i in range(len(lowerCAmelCase__							) - 1							):
            a__					,							a__	:					Dict										=cutoff_values[i], cutoff_values[i + 1]

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

              if indices_i.numel() == 0:
                continue

              a__	:					Any										=labels.index_select(0		,      lowerCAmelCase__							) - l_idx
              a__	:					List[Any]										=head_logprob.index_select(0		,      lowerCAmelCase__							)
              a__	:					Tuple										=hidden.index_select(0		,      lowerCAmelCase__							)
            else:
              a__	:					Optional[Any]										=hidden

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

              a__	:					Dict										=self._compute_logit(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)
              a__	:					List[Any]										=nn.functional.log_softmax(lowerCAmelCase__		,      dim=1							)
              a__	:					Optional[Any]										=self.cutoffs[0] + i - 1  # No probability for the head cluster
              if labels is not None:
                a__	:					Dict										=head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                    1		,      target_i[:, None]							).squeeze(1							)
              else:
                a__	:					Union[str, Any]										=head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                a__	:					Union[str, Any]										=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		,      lowerCAmelCase__							)   ->						int:



        '''simple docstring'''




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

            if i == 0:
              a__	:					Optional[int]										=torch.cat([weight_i, self.cluster_weight]		,      dim=0							)
              a__	:					str										=torch.cat([bias_i, self.cluster_bias]		,      dim=0							)

            weights.append(lowerCAmelCase__							)
            biases.append(lowerCAmelCase__							)

          a__					,							a__					,							a__	:					Optional[Any]										=weights[0], biases[0], self.out_projs[0]
          a__	:					Optional[Any]										=self._compute_logit(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)

          a__	:					Union[str, Any]										=hidden.new_empty((head_logit.size(0							), self.n_token)							)
          a__	:					str										=nn.functional.log_softmax(lowerCAmelCase__		,      dim=1							)

          a__	:					Any										=[0] + self.cutoffs
          for i in range(len(lowerCAmelCase__							) - 1							):
            a__					,							a__	:					Union[str, Any]										=cutoff_values[i], cutoff_values[i + 1]

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

              a__	:					Dict										=self._compute_logit(lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)
              a__	:					Any										=nn.functional.log_softmax(lowerCAmelCase__		,      dim=1							)

              a__	:					List[str]										=head_logprob[:, -i] + tail_logprob_i
              a__	:					Dict										=logprob_i

          return out













import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_funnel import FunnelTokenizer


UpperCAmelCase	:	int		     =       logging.get_logger(__name__)

UpperCAmelCase	:	Optional[Any]		     =       {"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

UpperCAmelCase	:	List[Any]		     =       [
    """small""",
    """small-base""",
    """medium""",
    """medium-base""",
    """intermediate""",
    """intermediate-base""",
    """large""",
    """large-base""",
    """xlarge""",
    """xlarge-base""",
]

UpperCAmelCase	:	Optional[int]		     =       {
    """vocab_file""": {
        """funnel-transformer/small""": """https://huggingface.co/funnel-transformer/small/resolve/main/vocab.txt""",
        """funnel-transformer/small-base""": """https://huggingface.co/funnel-transformer/small-base/resolve/main/vocab.txt""",
        """funnel-transformer/medium""": """https://huggingface.co/funnel-transformer/medium/resolve/main/vocab.txt""",
        """funnel-transformer/medium-base""": (
            """https://huggingface.co/funnel-transformer/medium-base/resolve/main/vocab.txt"""
        ),
        """funnel-transformer/intermediate""": (
            """https://huggingface.co/funnel-transformer/intermediate/resolve/main/vocab.txt"""
        ),
        """funnel-transformer/intermediate-base""": (
            """https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/vocab.txt"""
        ),
        """funnel-transformer/large""": """https://huggingface.co/funnel-transformer/large/resolve/main/vocab.txt""",
        """funnel-transformer/large-base""": """https://huggingface.co/funnel-transformer/large-base/resolve/main/vocab.txt""",
        """funnel-transformer/xlarge""": """https://huggingface.co/funnel-transformer/xlarge/resolve/main/vocab.txt""",
        """funnel-transformer/xlarge-base""": (
            """https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/vocab.txt"""
        ),
    },
    """tokenizer_file""": {
        """funnel-transformer/small""": """https://huggingface.co/funnel-transformer/small/resolve/main/tokenizer.json""",
        """funnel-transformer/small-base""": (
            """https://huggingface.co/funnel-transformer/small-base/resolve/main/tokenizer.json"""
        ),
        """funnel-transformer/medium""": """https://huggingface.co/funnel-transformer/medium/resolve/main/tokenizer.json""",
        """funnel-transformer/medium-base""": (
            """https://huggingface.co/funnel-transformer/medium-base/resolve/main/tokenizer.json"""
        ),
        """funnel-transformer/intermediate""": (
            """https://huggingface.co/funnel-transformer/intermediate/resolve/main/tokenizer.json"""
        ),
        """funnel-transformer/intermediate-base""": (
            """https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/tokenizer.json"""
        ),
        """funnel-transformer/large""": """https://huggingface.co/funnel-transformer/large/resolve/main/tokenizer.json""",
        """funnel-transformer/large-base""": (
            """https://huggingface.co/funnel-transformer/large-base/resolve/main/tokenizer.json"""
        ),
        """funnel-transformer/xlarge""": """https://huggingface.co/funnel-transformer/xlarge/resolve/main/tokenizer.json""",
        """funnel-transformer/xlarge-base""": (
            """https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/tokenizer.json"""
        ),
    },
}
UpperCAmelCase	:	Optional[int]		     =       {F"""funnel-transformer/{name}""": 512 for name in _model_names}
UpperCAmelCase	:	Optional[int]		     =       {F"""funnel-transformer/{name}""": {"""do_lower_case""": True} for name in _model_names}







class 				__lowerCAmelCase       (    UpperCamelCase__):
      _lowercase				:					str						=			VOCAB_FILES_NAMES
      _lowercase				:					List[Any]						=			PRETRAINED_VOCAB_FILES_MAP
      _lowercase				:					Dict						=			PRETRAINED_INIT_CONFIGURATION
      _lowercase				:					Union[str, Any]						=			FunnelTokenizer
      _lowercase				:					Any						=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _lowercase				:					int						=			2







      def __init__(    self		,      lowerCAmelCase__=None		,      lowerCAmelCase__=None		,      lowerCAmelCase__=True		,      lowerCAmelCase__="<unk>"		,      lowerCAmelCase__="<sep>"		,      lowerCAmelCase__="<pad>"		,      lowerCAmelCase__="<cls>"		,      lowerCAmelCase__="<mask>"		,      lowerCAmelCase__="<s>"		,      lowerCAmelCase__="</s>"		,      lowerCAmelCase__=True		,      lowerCAmelCase__=True		,      lowerCAmelCase__=None		,      lowerCAmelCase__="##"		,      **lowerCAmelCase__		,      )   ->						Optional[int]:



        '''simple docstring'''




        super().__init__(
            lowerCAmelCase__		,      tokenizer_file=lowerCAmelCase__		,      do_lower_case=lowerCAmelCase__		,      unk_token=lowerCAmelCase__		,      sep_token=lowerCAmelCase__		,      pad_token=lowerCAmelCase__		,      cls_token=lowerCAmelCase__		,      mask_token=lowerCAmelCase__		,      bos_token=lowerCAmelCase__		,      eos_token=lowerCAmelCase__		,      clean_text=lowerCAmelCase__		,      tokenize_chinese_chars=lowerCAmelCase__		,      strip_accents=lowerCAmelCase__		,      wordpieces_prefix=lowerCAmelCase__		,      **lowerCAmelCase__		,      )

        a__	:					Optional[Any]										=json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
        if (
            normalizer_state.get("lowercase"		,      lowerCAmelCase__							) != do_lower_case
            or normalizer_state.get("strip_accents"		,      lowerCAmelCase__							) != strip_accents
            or normalizer_state.get("handle_chinese_chars"		,      lowerCAmelCase__							) != tokenize_chinese_chars
        ):
          a__	:					List[str]										=getattr(lowerCAmelCase__		,      normalizer_state.pop("type"							)							)
          a__	:					Union[str, Any]										=do_lower_case
          a__	:					Any										=strip_accents
          a__	:					Optional[Any]										=tokenize_chinese_chars
          a__	:					Dict										=normalizer_class(**lowerCAmelCase__							)

        a__	:					Any										=do_lower_case







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__=None							)   ->						str:



        '''simple docstring'''




        a__	:					Dict										=[self.cls_token_id] + token_ids_a + [self.sep_token_id]

        if token_ids_a:
          output += token_ids_a + [self.sep_token_id]

        return output







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = None							)   ->						List[int]:



        '''simple docstring'''




        a__	:					Optional[int]										=[self.sep_token_id]
        a__	:					Union[str, Any]										=[self.cls_token_id]
        if token_ids_a is None:
          return len(cls							) * [self.cls_token_type_id] + len(token_ids_a + sep							) * [0]
        return len(cls							) * [self.cls_token_type_id] + len(token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__ = None							)   ->						Tuple[str]:



        '''simple docstring'''




        a__	:					Tuple										=self._tokenizer.model.save(lowerCAmelCase__		,      name=lowerCAmelCase__							)
        return tuple(lowerCAmelCase__							)






