




import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


UpperCAmelCase__				     = datasets.utils.logging.get_logger(__name__)







@dataclass
class a (     datasets.BuilderConfig  ):
						"""simple docstring"""

						UpperCamelCase_	:       Optional[datasets.Features]       	=		None
						UpperCamelCase_	:       str       	=		"utf-8"
						UpperCamelCase_	:       Optional[str]       	=		None
						UpperCamelCase_	:       Optional[str]       	=		None
						UpperCamelCase_	:       bool       	=		True  # deprecated
						UpperCamelCase_	:       Optional[int]       	=		None  # deprecated
						UpperCamelCase_	:       int       	=		10 << 20  # 10MB
						UpperCamelCase_	:       Optional[bool]       	=		None







class a (     datasets.ArrowBasedBuilder  ):
						"""simple docstring"""

						UpperCamelCase_	:       str       	=		JsonConfig






						def     UpperCAmelCase_      ( self :     Optional[int]				)    ->						Tuple:
													"""simple docstring"""


													if self.config.block_size is not None:
																				logger.warning('''The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead'''				)
																				__lowercase      =    self.config.block_size
													if self.config.use_threads is not True:
																				logger.warning(
																				    '''The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.'''				)
													if self.config.newlines_in_values is not None:
																				raise ValueError('''The JSON loader parameter `newlines_in_values` is no longer supported'''				)
													return datasets.DatasetInfo(features=self.config.features				)






						def     UpperCAmelCase_      ( self :     Union[str, Any]				,							lowerCamelCase__ :     Union[str, Any]				)    ->						List[Any]:
													"""simple docstring"""


													if not self.config.data_files:
																				raise ValueError(f'At least one data file must be specified, but got data_files={self.config.data_files}'				)
													__lowercase      =    dl_manager.download_and_extract(self.config.data_files				)
													if isinstance(UpperCAmelCase_				,							(str, list, tuple)				):
																				__lowercase      =    data_files
																				if isinstance(UpperCAmelCase_				,							UpperCAmelCase_				):
																											__lowercase      =    [files]
																				__lowercase      =    [dl_manager.iter_files(UpperCAmelCase_				) for file in files]
																				return [datasets.SplitGenerator(name=datasets.Split.TRAIN				,							gen_kwargs={'''files''': files}				)]
													__lowercase      =    []
													for split_name, files in data_files.items():
																				if isinstance(UpperCAmelCase_				,							UpperCAmelCase_				):
																											__lowercase      =    [files]
																				__lowercase      =    [dl_manager.iter_files(UpperCAmelCase_				) for file in files]
																				splits.append(datasets.SplitGenerator(name=UpperCAmelCase_				,							gen_kwargs={'''files''': files}				)				)
													return splits






						def     UpperCAmelCase_      ( self :     Tuple				,							lowerCamelCase__ :     pa.Table				)    ->						pa.Table:
													"""simple docstring"""


													if self.config.features is not None:
																				# adding missing columns
																				for column_name in set(self.config.features				) - set(pa_table.column_names				):
																											__lowercase      =    self.config.features.arrow_schema.field(UpperCAmelCase_				).type
																											__lowercase      =    pa_table.append_column(UpperCAmelCase_				,							pa.array([None] * len(UpperCAmelCase_				)				,							type=UpperCAmelCase_				)				)
																				# more expensive cast to support nested structures with keys in a different order
																				# allows str <-> int/float or str to Audio for example
																				__lowercase      =    table_cast(UpperCAmelCase_				,							self.config.features.arrow_schema				)
													return pa_table






						def     UpperCAmelCase_      ( self :     Optional[Any]				,							lowerCamelCase__ :     str				)    ->						List[str]:
													"""simple docstring"""


													for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase_				)				):
																				# If the file is one json object and if we need to look at the list of items in one specific field
																				if self.config.field is not None:
																											with open(UpperCAmelCase_				,							encoding=self.config.encoding				,							errors=self.config.encoding_errors				) as f:
																																		__lowercase      =    json.load(UpperCAmelCase_				)

																											# We keep only the field we are interested in
																											__lowercase      =    dataset[self.config.field]

																											# We accept two format: a list of dicts or a dict of lists
																											if isinstance(UpperCAmelCase_				,							(list, tuple)				):
																																		__lowercase      =    set().union(*[row.keys() for row in dataset]				)
																																		__lowercase      =    {col: [row.get(UpperCAmelCase_				) for row in dataset] for col in keys}
																											else:
																																		__lowercase      =    dataset
																											__lowercase      =    pa.Table.from_pydict(UpperCAmelCase_				)
																											yield file_idx, self._cast_table(UpperCAmelCase_				)

																				# If the file has one json object per line
																				else:
																											with open(UpperCAmelCase_				,							'''rb'''				) as f:
																																		__lowercase      =    0
																																		# Use block_size equal to the chunk size divided by 32 to leverage multithreading
																																		# Set a default minimum value of 16kB if the chunk size is really small
																																		__lowercase      =    max(self.config.chunksize // 32				,							16 << 10				)
																																		__lowercase      =    (
																																		    self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
																																		)
																																		while True:
																																									__lowercase      =    f.read(self.config.chunksize				)
																																									if not batch:
																																																break
																																									# Finish current line
																																									try:
																																																batch += f.readline()
																																									except (AttributeError, io.UnsupportedOperation):
																																																batch += readline(UpperCAmelCase_				)
																																									# PyArrow only accepts utf-8 encoded bytes
																																									if self.config.encoding != "utf-8":
																																																__lowercase      =    batch.decode(self.config.encoding				,							errors=UpperCAmelCase_				).encode('''utf-8'''				)
																																									try:
																																																while True:
																																																							try:
																																																														__lowercase      =    paj.read_json(
																																																														    io.BytesIO(UpperCAmelCase_				)				,							read_options=paj.ReadOptions(block_size=UpperCAmelCase_				)				)
																																																														break
																																																							except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
																																																														if (
																																																														    isinstance(UpperCAmelCase_				,							pa.ArrowInvalid				)
																																																														    and "straddling" not in str(UpperCAmelCase_				)
																																																														    or block_size > len(UpperCAmelCase_				)
																																																														):
																																																																					raise
																																																														else:
																																																																					# Increase the block size in case it was too small.
																																																																					# The block size will be reset for the next file.
																																																																					logger.debug(
																																																																					    f'Batch of {len(UpperCAmelCase_				)} bytes couldn\'t be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.'				)
																																																																					block_size *= 2
																																									except pa.ArrowInvalid as e:
																																																try:
																																																							with open(
																																																							    UpperCAmelCase_				,							encoding=self.config.encoding				,							errors=self.config.encoding_errors				) as f:
																																																														__lowercase      =    json.load(UpperCAmelCase_				)
																																																except json.JSONDecodeError:
																																																							logger.error(f'Failed to read file \'{file}\' with error {type(UpperCAmelCase_				)}: {e}'				)
																																																							raise e
																																																# If possible, parse the file as a list of json objects and exit the loop
																																																if isinstance(UpperCAmelCase_				,							UpperCAmelCase_				):  # list is the only sequence type supported in JSON
																																																							try:
																																																														__lowercase      =    set().union(*[row.keys() for row in dataset]				)
																																																														__lowercase      =    {col: [row.get(UpperCAmelCase_				) for row in dataset] for col in keys}
																																																														__lowercase      =    pa.Table.from_pydict(UpperCAmelCase_				)
																																																							except (pa.ArrowInvalid, AttributeError) as e:
																																																														logger.error(f'Failed to read file \'{file}\' with error {type(UpperCAmelCase_				)}: {e}'				)
																																																														raise ValueError(f'Not able to read records in the JSON file at {file}.'				) from None
																																																							yield file_idx, self._cast_table(UpperCAmelCase_				)
																																																							break
																																																else:
																																																							logger.error(f'Failed to read file \'{file}\' with error {type(UpperCAmelCase_				)}: {e}'				)
																																																							raise ValueError(
																																																							    f'Not able to read records in the JSON file at {file}. '
																																																							    f'You should probably indicate the field of the JSON file containing your records. '
																																																							    f'This JSON file contain the following fields: {str(list(dataset.keys()				)				)}. '
																																																							    f'Select the correct one and provide it as `field=\'XXX\'` to the dataset loading method. '				) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																																									yield (file_idx, batch_idx), self._cast_table(UpperCAmelCase_				)
																																									batch_idx += 1









import tensorflow as tf

from ...tf_utils import shape_list







class a (     tf.keras.layers.Layer  ):
						"""simple docstring"""

						def __init__( self :     int				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Tuple=1				,							lowerCamelCase__ :     Optional[Any]=False				,							**lowerCamelCase__ :     Optional[int]				)    ->						str:
													"""simple docstring"""


													super().__init__(**lowerCamelCase__				)

													__lowercase      =    vocab_size
													__lowercase      =    d_embed
													__lowercase      =    d_proj

													__lowercase      =    cutoffs + [vocab_size]
													__lowercase      =    [0] + self.cutoffs
													__lowercase      =    div_val

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

													__lowercase      =    []
													__lowercase      =    []






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				)    ->						Tuple:
													"""simple docstring"""


													if self.n_clusters > 0:
																				__lowercase      =    self.add_weight(
																				    shape=(self.n_clusters, self.d_embed)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name='''cluster_weight'''				)
																				__lowercase      =    self.add_weight(
																				    shape=(self.n_clusters,)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name='''cluster_bias'''				)

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

																											__lowercase      =    self.add_weight(
																											    shape=(d_emb_i, self.d_proj)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_projs_._{i}'				)
																											self.out_projs.append(lowerCamelCase__				)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx, d_emb_i)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx,)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._bias'				,							)
																											self.out_layers.append((weight, bias)				)
													super().build(lowerCamelCase__				)






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     int				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     List[str]=None				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    x
													if proj is not None:
																				__lowercase      =    tf.einsum('''ibd,ed->ibe'''				,							lowerCamelCase__				,							lowerCamelCase__				)
													return tf.einsum('''ibd,nd->ibn'''				,							lowerCamelCase__				,							lowerCamelCase__				) + b






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[str]				)    ->						str:
													"""simple docstring"""


													__lowercase      =    shape_list(lowerCamelCase__				)
													__lowercase      =    tf.range(lp_size[0]				,							dtype=target.dtype				)
													__lowercase      =    tf.stack([r, target]				,							1				)
													return tf.gather_nd(lowerCamelCase__				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Optional[int]=True				,							lowerCamelCase__ :     Optional[int]=False				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    0
													if self.n_clusters == 0:
																				__lowercase      =    self._logit(lowerCamelCase__				,							self.out_layers[0][0]				,							self.out_layers[0][1]				,							self.out_projs[0]				)
																				if target is not None:
																											__lowercase      =    tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__				,							logits=lowerCamelCase__				)
																				__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				,							axis=-1				)
													else:
																				__lowercase      =    shape_list(lowerCamelCase__				)
																				__lowercase      =    []
																				__lowercase      =    tf.zeros(hidden_sizes[:2]				)
																				for i in range(len(self.cutoffs				)				):
																											__lowercase     ,     __lowercase      =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
																											if target is not None:
																																		__lowercase      =    (target >= l_idx) & (target < r_idx)
																																		__lowercase      =    tf.where(lowerCamelCase__				)
																																		__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				) - l_idx

																											if self.div_val == 1:
																																		__lowercase      =    self.out_layers[0][0][l_idx:r_idx]
																																		__lowercase      =    self.out_layers[0][1][l_idx:r_idx]
																											else:
																																		__lowercase      =    self.out_layers[i][0]
																																		__lowercase      =    self.out_layers[i][1]

																											if i == 0:
																																		__lowercase      =    tf.concat([cur_W, self.cluster_weight]				,							0				)
																																		__lowercase      =    tf.concat([cur_b, self.cluster_bias]				,							0				)

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

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

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

													return out






