

import torch
from torch import nn







class 			lowercase    (	nn.Module      ):




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

											snake_case_          =   n_token
											snake_case_          =   d_embed
											snake_case_          =   d_proj

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

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

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

											snake_case_          =   nn.ModuleList()
											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(snake_case				,       snake_case	)	)	)
																									else:
																																self.out_projs.append(snake_case	)

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

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

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

											snake_case_          =   keep_order




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

											return logit




				def    a							( self				,       snake_case				,       snake_case=None				,       snake_case=False	):

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

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

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

																									weights.append(snake_case	)
																									biases.append(snake_case	)

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

																		snake_case_          =   self._compute_logit(snake_case				,       snake_case				,       snake_case				,       snake_case	)
																		snake_case_          =   nn.functional.log_softmax(snake_case				,       dim=1	)

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

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

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

																																if indices_i.numel() == 0:
																																							continue

																																snake_case_          =   labels.index_select(0				,       snake_case	) - l_idx
																																snake_case_          =   head_logprob.index_select(0				,       snake_case	)
																																snake_case_          =   hidden.index_select(0				,       snake_case	)
																									else:
																																snake_case_          =   hidden

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

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

											return out




				def    a							( self				,       snake_case	):
											if self.n_clusters == 0:
																		snake_case_          =   self._compute_logit(snake_case				,       self.out_layers[0].weight				,       self.out_layers[0].bias				,       self.out_projs[0]	)
																		return nn.functional.log_softmax(snake_case				,       dim=-1	)
											else:
																		# construct weights and biases
																		snake_case_    ,   snake_case_          =   [], []
																		for i in range(len(self.cutoffs	)	):
																									if self.div_val == 1:
																																snake_case_    ,   snake_case_          =   self.cutoff_ends[i], self.cutoff_ends[i + 1]
																																snake_case_          =   self.out_layers[0].weight[l_idx:r_idx]
																																snake_case_          =   self.out_layers[0].bias[l_idx:r_idx]
																									else:
																																snake_case_          =   self.out_layers[i].weight
																																snake_case_          =   self.out_layers[i].bias

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

																									weights.append(snake_case	)
																									biases.append(snake_case	)

																		snake_case_    ,   snake_case_    ,   snake_case_          =   weights[0], biases[0], self.out_projs[0]
																		snake_case_          =   self._compute_logit(snake_case				,       snake_case				,       snake_case				,       snake_case	)

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

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

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

																																snake_case_          =   self._compute_logit(snake_case				,       snake_case				,       snake_case				,       snake_case	)
																																snake_case_          =   nn.functional.log_softmax(snake_case				,       dim=1	)

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

																		return out







# Lint as: python3
import itertools
import os
import re


_UpperCAmelCase		:       str						    =	re.compile(R"""([A-Z]+)([A-Z][a-z])""")
_UpperCAmelCase		:       Dict						    =	re.compile(R"""([a-z\d])([A-Z])""")

_UpperCAmelCase		:       Dict						    =	re.compile(R"""(?<!_)_(?!_)""")
_UpperCAmelCase		:       Tuple						    =	re.compile(R"""(_{2,})""")

_UpperCAmelCase		:       Any						    =	R"""^\w+(\.\w+)*$"""

_UpperCAmelCase		:       List[str]						    =	R"""<>:/\|?*"""





def  __lowerCamelCase     (			UpperCamelCase__				):







							'''simple docstring'''






							snake_case_          =   _uppercase_uppercase_re.sub(r'\1_\2'       ,							UpperCamelCase__				)
							snake_case_          =   _lowercase_uppercase_re.sub(r'\1_\2'       ,							UpperCamelCase__				)
							return name.lower()





def  __lowerCamelCase     (			UpperCamelCase__				):







							'''simple docstring'''






							snake_case_          =   _single_underscore_re.split(UpperCamelCase__				)
							snake_case_          =   [_multiple_underscores_re.split(UpperCamelCase__				) for n in name]
							return "".join(n.capitalize() for n in itertools.chain.from_iterable(UpperCamelCase__				) if n != ''				)





def  __lowerCamelCase     (			UpperCamelCase__				):







							'''simple docstring'''






							if os.path.basename(UpperCamelCase__				) != name:
														raise ValueError(F'''Should be a dataset name, not a path: {name}'''				)
							return camelcase_to_snakecase(UpperCamelCase__				)





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__				):







							'''simple docstring'''






							if os.path.basename(UpperCamelCase__				) != name:
														raise ValueError(F'''Should be a dataset name, not a path: {name}'''				)
							if not re.match(_split_re       ,							UpperCamelCase__				):
														raise ValueError(F'''Split name should match \'{_split_re}\'\' but got \'{split}\'.'''				)
							return F'''{filename_prefix_for_name(UpperCamelCase__				)}-{split}'''





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__=None				):







							'''simple docstring'''






							snake_case_          =   filename_prefix_for_split(UpperCamelCase__       ,							UpperCamelCase__				)
							if filetype_suffix:
														prefix += F'''.{filetype_suffix}'''
							snake_case_          =   os.path.join(UpperCamelCase__       ,							UpperCamelCase__				)
							return F'''{filepath}*'''





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__=None       ,							UpperCamelCase__=None				):







							'''simple docstring'''






							snake_case_          =   filename_prefix_for_split(UpperCamelCase__       ,							UpperCamelCase__				)
							snake_case_          =   os.path.join(UpperCamelCase__       ,							UpperCamelCase__				)

							if shard_lengths:
														snake_case_          =   len(UpperCamelCase__				)
														snake_case_          =   [F'''{prefix}-{shard_id:05d}-of-{num_shards:05d}''' for shard_id in range(UpperCamelCase__				)]
														if filetype_suffix:
																					snake_case_          =   [filename + F'''.{filetype_suffix}''' for filename in filenames]
														return filenames
							else:
														snake_case_          =   prefix
														if filetype_suffix:
																					filename += F'''.{filetype_suffix}'''
														return [filename]




