


import torch
from torch import nn






class      UpperCamelCase					(		nn.Module     ):


			'''simple docstring'''







			def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=1   ,							UpperCamelCase_=False		):
								super().__init__()

								lowercase_       :Union[str, Any]     =       n_token
								lowercase_       :Optional[Any]     =       d_embed
								lowercase_       :int     =       d_proj

								lowercase_       :Union[str, Any]     =       cutoffs + [n_token]
								lowercase_       :List[Any]     =       [0] + self.cutoffs
								lowercase_       :int     =       div_val

								lowercase_       :Dict     =       self.cutoffs[0]
								lowercase_       :Optional[int]     =       len(self.cutoffs		) - 1
								lowercase_       :str     =       self.shortlist_size + self.n_clusters

								if self.n_clusters > 0:
													lowercase_       :List[str]     =       nn.Parameter(torch.zeros(self.n_clusters   ,							self.d_embed		)		)
													lowercase_       :Optional[int]     =       nn.Parameter(torch.zeros(self.n_clusters		)		)

								lowercase_       :int     =       nn.ModuleList()
								lowercase_       :List[Any]     =       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(UpperCamelCase_   ,							UpperCamelCase_		)		)		)
																		else:
																							self.out_projs.append(UpperCamelCase_		)

													self.out_layers.append(nn.Linear(UpperCamelCase_   ,							UpperCamelCase_		)		)
								else:
													for i in range(len(self.cutoffs		)		):
																		lowercase_       :str     =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
																		lowercase_       :Union[str, Any]     =       d_embed // (div_val**i)

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

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

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

								return logit
			def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_=None   ,							UpperCamelCase_=False		):

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

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

																		if i == 0:
																							lowercase_       :Union[str, Any]     =       torch.cat([weight_i, self.cluster_weight]   ,							dim=0		)
																							lowercase_       :List[str]     =       torch.cat([bias_i, self.cluster_bias]   ,							dim=0		)

																		weights.append(UpperCamelCase_		)
																		biases.append(UpperCamelCase_		)

													lowercase_       :List[str]     =       weights[0], biases[0], self.out_projs[0]

													lowercase_       :Dict     =       self._compute_logit(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)
													lowercase_       :Any     =       nn.functional.log_softmax(UpperCamelCase_   ,							dim=1		)

													if labels is None:
																		lowercase_       :List[str]     =       hidden.new_empty((head_logit.size(0		), self.n_token)		)
													else:
																		lowercase_       :Dict     =       torch.zeros_like(UpperCamelCase_   ,							dtype=hidden.dtype   ,							device=hidden.device		)

													lowercase_       :str     =       0
													lowercase_       :Tuple     =       [0] + self.cutoffs
													for i in range(len(UpperCamelCase_		) - 1		):
																		lowercase_       :str     =       cutoff_values[i], cutoff_values[i + 1]

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

																							if indices_i.numel() == 0:
																												continue

																							lowercase_       :Any     =       labels.index_select(0   ,							UpperCamelCase_		) - l_idx
																							lowercase_       :int     =       head_logprob.index_select(0   ,							UpperCamelCase_		)
																							lowercase_       :str     =       hidden.index_select(0   ,							UpperCamelCase_		)
																		else:
																							lowercase_       :Tuple     =       hidden

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

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

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

								return out




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

																		if i == 0:
																							lowercase_       :List[str]     =       torch.cat([weight_i, self.cluster_weight]   ,							dim=0		)
																							lowercase_       :Union[str, Any]     =       torch.cat([bias_i, self.cluster_bias]   ,							dim=0		)

																		weights.append(UpperCamelCase_		)
																		biases.append(UpperCamelCase_		)

													lowercase_       :Optional[Any]     =       weights[0], biases[0], self.out_projs[0]
													lowercase_       :List[str]     =       self._compute_logit(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)

													lowercase_       :Tuple     =       hidden.new_empty((head_logit.size(0		), self.n_token)		)
													lowercase_       :Any     =       nn.functional.log_softmax(UpperCamelCase_   ,							dim=1		)

													lowercase_       :Optional[Any]     =       [0] + self.cutoffs
													for i in range(len(UpperCamelCase_		) - 1		):
																		lowercase_       :List[Any]     =       cutoff_values[i], cutoff_values[i + 1]

																		if i == 0:
																							lowercase_       :Dict     =       head_logprob[:, : self.cutoffs[0]]
																		else:
																							lowercase_       :Union[str, Any]     =       weights[i], biases[i], self.out_projs[i]

																							lowercase_       :Dict     =       self._compute_logit(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		)
																							lowercase_       :Union[str, Any]     =       nn.functional.log_softmax(UpperCamelCase_   ,							dim=1		)

																							lowercase_       :List[str]     =       head_logprob[:, -i] + tail_logprob_i
																							lowercase_       :List[Any]     =       logprob_i

													return out





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


SCREAMING_SNAKE_CASE       :		str         =				{
    "configuration_tapas": ["TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP", "TapasConfig"],
    "tokenization_tapas": ["TapasTokenizer"],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   SCREAMING_SNAKE_CASE       :		Dict         =				[
       "TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TapasForMaskedLM",
       "TapasForQuestionAnswering",
       "TapasForSequenceClassification",
       "TapasModel",
       "TapasPreTrainedModel",
       "load_tf_weights_in_tapas",
   ]
try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   SCREAMING_SNAKE_CASE       :		str         =				[
       "TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TFTapasForMaskedLM",
       "TFTapasForQuestionAnswering",
       "TFTapasForSequenceClassification",
       "TFTapasModel",
       "TFTapasPreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
   from .tokenization_tapas import TapasTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tapas import (
          TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
          TapasForMaskedLM,
          TapasForQuestionAnswering,
          TapasForSequenceClassification,
          TapasModel,
          TapasPreTrainedModel,
          load_tf_weights_in_tapas,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_tapas import (
          TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFTapasForMaskedLM,
          TFTapasForQuestionAnswering,
          TFTapasForSequenceClassification,
          TFTapasModel,
          TFTapasPreTrainedModel,
      )


else:
   import sys

   SCREAMING_SNAKE_CASE       :		Union[str, Any]         =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

