






'''simple docstring'''







import tensorflow as tf

from ...tf_utils import shape_list

class _lowercase			(       tf.keras.layers.Layer	):




							def __init__(							self:       Union[str, Any]    ,							UpperCamelCase__:       Optional[int]    ,							UpperCamelCase__:       Union[str, Any]    ,							UpperCamelCase__:       Tuple    ,							UpperCamelCase__:       Optional[int]    ,							UpperCamelCase__:       Any=1    ,							UpperCamelCase__:       Optional[Any]=False    ,							**UpperCamelCase__:       List[Any]  ):
														super().__init__(**_snake_case  )

														lowerCamelCase__      :			Dict 						= vocab_size
														lowerCamelCase__      :			List[str] 						= d_embed
														lowerCamelCase__      :			Tuple 						= d_proj

														lowerCamelCase__      :			Optional[Any] 						= cutoffs + [vocab_size]
														lowerCamelCase__      :			Union[str, Any] 						= [0] + self.cutoffs
														lowerCamelCase__      :			Union[str, Any] 						= div_val

														lowerCamelCase__      :			Union[str, Any] 						= self.cutoffs[0]
														lowerCamelCase__      :			List[str] 						= len(self.cutoffs  ) - 1
														lowerCamelCase__      :			List[Any] 						= self.shortlist_size + self.n_clusters
														lowerCamelCase__      :			Tuple 						= keep_order

														lowerCamelCase__      :			Any 						= []
														lowerCamelCase__      :			int 						= []




							def 	lowerCamelCase_       (							self:       Optional[int]    ,							UpperCamelCase__:       Any  ):


														if self.n_clusters > 0:
																					lowerCamelCase__      :			List[str] 						= self.add_weight(
																					    shape=(self.n_clusters, self.d_embed)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name="""cluster_weight"""  )
																					lowerCamelCase__      :			int 						= self.add_weight(
																					    shape=(self.n_clusters,)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name="""cluster_bias"""  )

														if self.div_val == 1:
																					for i in range(len(self.cutoffs  )  ):
																												if self.d_proj != self.d_embed:
																																			lowerCamelCase__      :			List[str] 						= self.add_weight(
																																			    shape=(self.d_embed, self.d_proj)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_projs_._{i}'''    ,							)
																																			self.out_projs.append(_snake_case  )
																												else:
																																			self.out_projs.append(_snake_case  )
																												lowerCamelCase__      :			Optional[int] 						= self.add_weight(
																												    shape=(self.vocab_size, self.d_embed)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_layers_._{i}_._weight'''    ,							)
																												lowerCamelCase__      :			Any 						= self.add_weight(
																												    shape=(self.vocab_size,)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_layers_._{i}_._bias'''    ,							)
																												self.out_layers.append((weight, bias)  )
														else:
																					for i in range(len(self.cutoffs  )  ):
																												lowerCamelCase__						,	lowerCamelCase__      :			int 						= self.cutoff_ends[i], self.cutoff_ends[i + 1]
																												lowerCamelCase__      :			List[Any] 						= self.d_embed // (self.div_val**i)

																												lowerCamelCase__      :			Optional[int] 						= self.add_weight(
																												    shape=(d_emb_i, self.d_proj)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_projs_._{i}'''  )
																												self.out_projs.append(_snake_case  )
																												lowerCamelCase__      :			Tuple 						= self.add_weight(
																												    shape=(r_idx - l_idx, d_emb_i)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_layers_._{i}_._weight'''    ,							)
																												lowerCamelCase__      :			int 						= self.add_weight(
																												    shape=(r_idx - l_idx,)    ,							initializer="""zeros"""    ,							trainable=_snake_case    ,							name=F'''out_layers_._{i}_._bias'''    ,							)
																												self.out_layers.append((weight, bias)  )
														super().build(_snake_case  )




							@staticmethod
							def 	lowerCamelCase_       (							UpperCamelCase__:       List[Any]    ,							UpperCamelCase__:       str    ,							UpperCamelCase__:       Any    ,							UpperCamelCase__:       Union[str, Any]=None  ):
														lowerCamelCase__      :			List[Any] 						= x
														if proj is not None:
																					lowerCamelCase__      :			str 						= tf.einsum("""ibd,ed->ibe"""    ,							_snake_case    ,							_snake_case  )
														return tf.einsum("""ibd,nd->ibn"""    ,							_snake_case    ,							_snake_case  ) + b




							@staticmethod
							def 	lowerCamelCase_       (							UpperCamelCase__:       Optional[int]    ,							UpperCamelCase__:       Optional[Any]  ):
														lowerCamelCase__      :			Union[str, Any] 						= shape_list(_snake_case  )
														lowerCamelCase__      :			Tuple 						= tf.range(lp_size[0]    ,							dtype=target.dtype  )
														lowerCamelCase__      :			List[Any] 						= tf.stack([r, target]    ,							1  )
														return tf.gather_nd(_snake_case    ,							_snake_case  )




							def 	lowerCamelCase_       (							self:       Optional[int]    ,							UpperCamelCase__:       Any    ,							UpperCamelCase__:       Optional[Any]    ,							UpperCamelCase__:       Any=True    ,							UpperCamelCase__:       Union[str, Any]=False  ):
														lowerCamelCase__      :			Optional[Any] 						= 0
														if self.n_clusters == 0:
																					lowerCamelCase__      :			int 						= self._logit(_snake_case    ,							self.out_layers[0][0]    ,							self.out_layers[0][1]    ,							self.out_projs[0]  )
																					if target is not None:
																												lowerCamelCase__      :			str 						= tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_snake_case    ,							logits=_snake_case  )
																					lowerCamelCase__      :			List[Any] 						= tf.nn.log_softmax(_snake_case    ,							axis=-1  )
														else:
																					lowerCamelCase__      :			str 						= shape_list(_snake_case  )
																					lowerCamelCase__      :			Tuple 						= []
																					lowerCamelCase__      :			Dict 						= tf.zeros(hidden_sizes[:2]  )
																					for i in range(len(self.cutoffs  )  ):
																												lowerCamelCase__						,	lowerCamelCase__      :			Tuple 						= self.cutoff_ends[i], self.cutoff_ends[i + 1]
																												if target is not None:
																																			lowerCamelCase__      :			Any 						= (target >= l_idx) & (target < r_idx)
																																			lowerCamelCase__      :			Optional[Any] 						= tf.where(_snake_case  )
																																			lowerCamelCase__      :			Union[str, Any] 						= tf.boolean_mask(_snake_case    ,							_snake_case  ) - l_idx

																												if self.div_val == 1:
																																			lowerCamelCase__      :			Tuple 						= self.out_layers[0][0][l_idx:r_idx]
																																			lowerCamelCase__      :			Any 						= self.out_layers[0][1][l_idx:r_idx]
																												else:
																																			lowerCamelCase__      :			Dict 						= self.out_layers[i][0]
																																			lowerCamelCase__      :			Optional[Any] 						= self.out_layers[i][1]

																												if i == 0:
																																			lowerCamelCase__      :			List[str] 						= tf.concat([cur_W, self.cluster_weight]    ,							0  )
																																			lowerCamelCase__      :			List[Any] 						= tf.concat([cur_b, self.cluster_bias]    ,							0  )

																																			lowerCamelCase__      :			List[str] 						= self._logit(_snake_case    ,							_snake_case    ,							_snake_case    ,							self.out_projs[0]  )
																																			lowerCamelCase__      :			Dict 						= tf.nn.log_softmax(_snake_case  )
																																			out.append(head_logprob[..., : self.cutoffs[0]]  )
																																			if target is not None:
																																										lowerCamelCase__      :			Tuple 						= tf.boolean_mask(_snake_case    ,							_snake_case  )
																																										lowerCamelCase__      :			Tuple 						= self._gather_logprob(_snake_case    ,							_snake_case  )
																												else:
																																			lowerCamelCase__      :			Union[str, Any] 						= self._logit(_snake_case    ,							_snake_case    ,							_snake_case    ,							self.out_projs[i]  )
																																			lowerCamelCase__      :			List[Any] 						= tf.nn.log_softmax(_snake_case  )
																																			lowerCamelCase__      :			Dict 						= self.cutoffs[0] + i - 1  # No probability for the head cluster
																																			lowerCamelCase__      :			Tuple 						= head_logprob[..., cluster_prob_idx, None] + tail_logprob
																																			out.append(_snake_case  )
																																			if target is not None:
																																										lowerCamelCase__      :			Optional[int] 						= tf.boolean_mask(_snake_case    ,							_snake_case  )
																																										lowerCamelCase__      :			List[str] 						= tf.boolean_mask(_snake_case    ,							_snake_case  )
																																										lowerCamelCase__      :			Optional[int] 						= self._gather_logprob(_snake_case    ,							_snake_case  )
																																										cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																												if target is not None:
																																			loss += tf.scatter_nd(_snake_case    ,							-cur_logprob    ,							shape_list(_snake_case  )  )
																					lowerCamelCase__      :			Tuple 						= tf.concat(_snake_case    ,							axis=-1  )

														if target is not None:
																					if return_mean:
																												lowerCamelCase__      :			List[str] 						= tf.reduce_mean(_snake_case  )
																					# Add the training-time loss value to the layer using `self.add_loss()`.
																					self.add_loss(_snake_case  )

																					# Log the loss as a metric (we could log arbitrary metrics,
																					# including different metrics for training and inference.
																					self.add_metric(_snake_case    ,							name=self.name    ,							aggregation="""mean""" if return_mean else """"""  )

														return out









'''simple docstring'''







import os
from typing import Optional

import fsspec
from fsspec.archive import AbstractArchiveFileSystem
from fsspec.utils import DEFAULT_BLOCK_SIZE

class _lowercase			(       _lowercase	):
							a          = """"""
							a          = (
							    None  # protocol passed in prefix to the url. ex: "gzip", for gzip://file.txt::http://foo.bar/file.txt.gz
							)
							a          = None  # compression type in fsspec. ex: "gzip"
							a          = None  # extension of the filename to strip. ex: "".gz" to get file.txt from file.txt.gz




							def __init__(							self:       str    ,							UpperCamelCase__:       str = ""    ,							UpperCamelCase__:       Optional[str] = None    ,							UpperCamelCase__:       Optional[dict] = None    ,							**UpperCamelCase__:       List[Any]  ):
														super().__init__(self    ,							**UpperCamelCase__  )
														# always open as "rb" since fsspec can then use the TextIOWrapper to make it work for "r" mode
														lowerCamelCase__      :			List[Any] 						= fsspec.open(
														    UpperCamelCase__    ,							mode="""rb"""    ,							protocol=UpperCamelCase__    ,							compression=self.compression    ,							client_kwargs={
														        """requote_redirect_url""": False,  # see https://github.com/huggingface/datasets/pull/5459
														        """trust_env""": True,  # Enable reading proxy env variables.
														        **(target_options or {}).pop("""client_kwargs"""    ,							{}  ),  # To avoid issues if it was already passed.
														    }    ,							**(target_options or {})    ,							)
														lowerCamelCase__      :			str 						= os.path.basename(self.file.path.split("""::"""  )[0]  )
														lowerCamelCase__      :			Union[str, Any] 						= (
														    self.compressed_name[: self.compressed_name.rindex("""."""  )]
														    if """.""" in self.compressed_name
														    else self.compressed_name
														)
														lowerCamelCase__      :			Tuple 						= None




							@classmethod
							def 	lowerCamelCase_       (							cls:       Optional[int]    ,							UpperCamelCase__:       Optional[int]  ):
														# compressed file paths are always relative to the archive root
														return super()._strip_protocol(UpperCamelCase__  ).lstrip("""/"""  )




							def 	lowerCamelCase_       (							self:       Tuple  ):
														if self.dir_cache is None:
																					lowerCamelCase__      :			Dict 						= {**self.file.fs.info(self.file.path  ), """name""": self.uncompressed_name}
																					lowerCamelCase__      :			int 						= {f["""name"""]: f}




							def 	lowerCamelCase_       (							self:       List[str]    ,							UpperCamelCase__:       str  ):
														return self.file.open().read()




							def 	lowerCamelCase_       (							self:       Optional[int]    ,							UpperCamelCase__:       str    ,							UpperCamelCase__:       str = "rb"    ,							UpperCamelCase__:       Optional[int]=None    ,							UpperCamelCase__:       Tuple=True    ,							UpperCamelCase__:       Tuple=None    ,							**UpperCamelCase__:       Optional[Any]    ,							):
														lowerCamelCase__      :			Union[str, Any] 						= self._strip_protocol(UpperCamelCase__  )
														if mode != "rb":
																					raise ValueError(F'''Tried to read with mode {mode} on file {self.file.path} opened with mode \'rb\''''  )
														return self.file.open()



class _lowercase			(       _lowercase	):
							a          = """bz2"""
							a          = """bz2"""
							a          = """.bz2"""



class _lowercase			(       _lowercase	):
							a          = """gzip"""
							a          = """gzip"""
							a          = """.gz"""



class _lowercase			(       _lowercase	):
							a          = """lz4"""
							a          = """lz4"""
							a          = """.lz4"""



class _lowercase			(       _lowercase	):
							a          = """xz"""
							a          = """xz"""
							a          = """.xz"""



class _lowercase			(       _lowercase	):
							a          = """zstd"""
							a          = """zstd"""
							a          = """.zst"""




							def __init__(							self:       int    ,							UpperCamelCase__:       str    ,							UpperCamelCase__:       str = "rb"    ,							UpperCamelCase__:       Optional[str] = None    ,							UpperCamelCase__:       Optional[dict] = None    ,							UpperCamelCase__:       int = DEFAULT_BLOCK_SIZE    ,							**UpperCamelCase__:       Dict    ,							):
														super().__init__(
														    fo=UpperCamelCase__    ,							mode=UpperCamelCase__    ,							target_protocol=UpperCamelCase__    ,							target_options=UpperCamelCase__    ,							block_size=UpperCamelCase__    ,							**UpperCamelCase__    ,							)
														# We need to wrap the zstd decompressor to avoid this error in fsspec==2021.7.0 and zstandard==0.15.2:
														#
														# File "/Users/user/.virtualenvs/hf-datasets/lib/python3.7/site-packages/fsspec/core.py", line 145, in open
														#     out.close = close
														# AttributeError: 'zstd.ZstdDecompressionReader' object attribute 'close' is read-only
														#
														# see https://github.com/intake/filesystem_spec/issues/725
														lowerCamelCase__      :			Tuple 						= self.file.__enter__

														class _lowercase			:




																					def __init__(							self:       Optional[int]    ,							UpperCamelCase__:       Any  ):
																												lowerCamelCase__      :			Optional[int] 						= file_




																					def __enter__(							self:       List[Any]  ):
																												self._file.__enter__()
																												return self




																					def __exit__(							self:       Any    ,							*UpperCamelCase__:       str    ,							**UpperCamelCase__:       Any  ):
																												self._file.__exit__(*UpperCamelCase__    ,							**UpperCamelCase__  )




																					def __iter__(							self:       Any  ):
																												return iter(self._file  )




																					def 	lowerCamelCase_       (							self:       List[Any]  ):
																												return next(self._file  )




																					def __getattr__(							self:       List[str]    ,							UpperCamelCase__:       Dict  ):
																												return getattr(self._file    ,							UpperCamelCase__  )

														def fixed_enter(*UpperCamelCase__:       Union[str, Any]    ,							**UpperCamelCase__:       List[str]  ):
																					return WrappedFile(_enter(*UpperCamelCase__    ,							**UpperCamelCase__  )  )

														lowerCamelCase__      :			Optional[Any] 						= fixed_enter

