
"""simple docstring"""


import torch
from torch import nn




class  __lowerCamelCase  ( nn.Module			):




				def __init__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase=1							,						UpperCAmelCase=False   ):
							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(UpperCAmelCase							,						UpperCAmelCase   )   )   )
													else:
																self.out_projs.append(UpperCAmelCase   )

										self.out_layers.append(nn.Linear(UpperCAmelCase							,						UpperCAmelCase   )   )
							else:
										for i in range(len(self.cutoffs   )   ):
													lowerCamelCase_		, lowerCamelCase_					=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
													lowerCamelCase_					=  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   )   )

							lowerCamelCase_					=  keep_order




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   ):
							if proj is None:
										lowerCamelCase_					=  nn.functional.linear(UpperCAmelCase							,						UpperCAmelCase							,						bias=UpperCAmelCase   )
							else:
										# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
										lowerCamelCase_					=  nn.functional.linear(UpperCAmelCase							,						proj.t().contiguous()   )
										lowerCamelCase_					=  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
										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(UpperCAmelCase							,						self.out_layers[0].weight							,						self.out_layers[0].bias							,						self.out_projs[0]   )
										if labels is not None:
													lowerCamelCase_					=  labels != -100
													lowerCamelCase_					=  torch.zeros_like(UpperCAmelCase							,						dtype=hidden.dtype							,						device=hidden.device   )
													lowerCamelCase_					=  (
													    -nn.functional.log_softmax(UpperCAmelCase							,						dim=-1   )[mask].gather(1							,						labels[mask].unsqueeze(1   )   ).squeeze(1   )
													)
										else:
													lowerCamelCase_					=  nn.functional.log_softmax(UpperCAmelCase							,						dim=-1   )
							else:
										# construct weights and biases
										lowerCamelCase_		, lowerCamelCase_					=  [], []
										for i in range(len(self.cutoffs   )   ):
													if self.div_val == 1:
																lowerCamelCase_		, 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(UpperCAmelCase   )
													biases.append(UpperCAmelCase   )

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

										lowerCamelCase_					=  self._compute_logit(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
										lowerCamelCase_					=  nn.functional.log_softmax(UpperCAmelCase							,						dim=1   )

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

										lowerCamelCase_					=  0
										lowerCamelCase_					=  [0] + self.cutoffs
										for i in range(len(UpperCAmelCase   ) - 1   ):
													lowerCamelCase_		, 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							,						UpperCAmelCase   ) - l_idx
																lowerCamelCase_					=  head_logprob.index_select(0							,						UpperCAmelCase   )
																lowerCamelCase_					=  hidden.index_select(0							,						UpperCAmelCase   )
													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_		, lowerCamelCase_		, lowerCamelCase_					=  weights[i], biases[i], self.out_projs[i]

																lowerCamelCase_					=  self._compute_logit(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
																lowerCamelCase_					=  nn.functional.log_softmax(UpperCAmelCase							,						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							,						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:
										lowerCamelCase_					=  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
										lowerCamelCase_		, lowerCamelCase_					=  [], []
										for i in range(len(self.cutoffs   )   ):
													if self.div_val == 1:
																lowerCamelCase_		, 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(UpperCAmelCase   )
													biases.append(UpperCAmelCase   )

										lowerCamelCase_		, lowerCamelCase_		, lowerCamelCase_					=  weights[0], biases[0], self.out_projs[0]
										lowerCamelCase_					=  self._compute_logit(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )

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

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

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

																lowerCamelCase_					=  self._compute_logit(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
																lowerCamelCase_					=  nn.functional.log_softmax(UpperCAmelCase							,						dim=1   )

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

										return out





"""simple docstring"""


import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD






def  lowercase						(	lowerCAmelCase__						):
			def wrapper(*lowerCAmelCase__	,**lowerCAmelCase__						):
						lowerCamelCase_					=  timeit.default_timer()
						lowerCamelCase_					=  func(*lowerCAmelCase__	,**lowerCAmelCase__						)
						lowerCamelCase_					=  timeit.default_timer() - starttime
						return delta

			lowerCamelCase_					=  func.__name__

			return wrapper






def  lowercase						(	lowerCAmelCase__	,lowerCAmelCase__=100	,lowerCAmelCase__=None						):
			lowerCamelCase_					=  []
			lowerCamelCase_					=  seq_shapes or {}
			for i in range(lowerCAmelCase__						):
						lowerCamelCase_					=  {}
						for col_id, (k, v) in enumerate(features.items()						):
									if isinstance(lowerCAmelCase__	,_ArrayXD						):
												lowerCamelCase_					=  np.random.rand(*v.shape						).astype(v.dtype						)
									elif isinstance(lowerCAmelCase__	,datasets.Value						):
												if v.dtype == "string":
															lowerCamelCase_					=  '''The small grey turtle was surprisingly fast when challenged.'''
												else:
															lowerCamelCase_					=  np.random.randint(10	,size=1						).astype(v.dtype						).item()
									elif isinstance(lowerCAmelCase__	,datasets.Sequence						):
												while isinstance(lowerCAmelCase__	,datasets.Sequence						):
															lowerCamelCase_					=  v.feature
												lowerCamelCase_					=  seq_shapes[k]
												lowerCamelCase_					=  np.random.rand(*lowerCAmelCase__						).astype(v.dtype						)
									lowerCamelCase_					=  data

						dummy_data.append((i, example)						)

			return dummy_data






def  lowercase						(	lowerCAmelCase__	,lowerCAmelCase__	,lowerCAmelCase__=100	,lowerCAmelCase__=None						):
			lowerCamelCase_					=  generate_examples(lowerCAmelCase__	,num_examples=lowerCAmelCase__	,seq_shapes=lowerCAmelCase__						)

			with ArrowWriter(features=lowerCAmelCase__	,path=lowerCAmelCase__						) as writer:
						for key, record in dummy_data:
									lowerCamelCase_					=  features.encode_example(lowerCAmelCase__						)
									writer.write(lowerCAmelCase__						)

						lowerCamelCase_		, lowerCamelCase_					=  writer.finalize()

			if not num_final_examples == num_examples:
						raise ValueError(
						    f"Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}."						)

			lowerCamelCase_					=  datasets.Dataset.from_file(filename=lowerCAmelCase__	,info=datasets.DatasetInfo(features=lowerCAmelCase__						)						)

			return dataset



