





"""simple docstring"""






import tensorflow as tf

from ...tf_utils import shape_list




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







					'''simple docstring'''




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

									__a			:    Dict      							=	vocab_size
									__a			:    Optional[int]      							=	d_embed
									__a			:    Union[str, Any]      							=	d_proj

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

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

									__a			:    str      							=	[]
									__a			:    Optional[int]      							=	[]







					def 			_lowerCamelCase		(    self						,      _UpperCAmelCase  ):
									if self.n_clusters > 0:
													__a			:    Union[str, Any]      							=	self.add_weight(
													    shape=(self.n_clusters, self.d_embed)						,      initializer='''zeros'''						,      trainable=_UpperCAmelCase						,      name='''cluster_weight'''  )
													__a			:    List[str]      							=	self.add_weight(
													    shape=(self.n_clusters,)						,      initializer='''zeros'''						,      trainable=_UpperCAmelCase						,      name='''cluster_bias'''  )

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

																	__a			:    List[Any]      							=	self.add_weight(
																	    shape=(d_emb_i, self.d_proj)						,      initializer='''zeros'''						,      trainable=_UpperCAmelCase						,      name=f"""out_projs_._{i}"""  )
																	self.out_projs.append(_UpperCAmelCase  )
																	__a			:    List[Any]      							=	self.add_weight(
																	    shape=(r_idx - l_idx, d_emb_i)						,      initializer='''zeros'''						,      trainable=_UpperCAmelCase						,      name=f"""out_layers_._{i}_._weight"""						,      )
																	__a			:    List[Any]      							=	self.add_weight(
																	    shape=(r_idx - l_idx,)						,      initializer='''zeros'''						,      trainable=_UpperCAmelCase						,      name=f"""out_layers_._{i}_._bias"""						,      )
																	self.out_layers.append((weight, bias)  )
									super().build(_UpperCAmelCase  )







					@staticmethod
					def 			_lowerCamelCase		(    _UpperCAmelCase						,      _UpperCAmelCase						,      _UpperCAmelCase						,      _UpperCAmelCase=None  ):
									__a			:    Any      							=	x
									if proj is not None:
													__a			:    Optional[int]      							=	tf.einsum('''ibd,ed->ibe'''						,      _UpperCAmelCase						,      _UpperCAmelCase  )
									return tf.einsum('''ibd,nd->ibn'''						,      _UpperCAmelCase						,      _UpperCAmelCase  ) + b







					@staticmethod
					def 			_lowerCamelCase		(    _UpperCAmelCase						,      _UpperCAmelCase  ):
									__a			:    Optional[Any]      							=	shape_list(_UpperCAmelCase  )
									__a			:    List[str]      							=	tf.range(lp_size[0]						,      dtype=target.dtype  )
									__a			:    Optional[int]      							=	tf.stack([r, target]						,      1  )
									return tf.gather_nd(_UpperCAmelCase						,      _UpperCAmelCase  )







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

																	if self.div_val == 1:
																					__a			:    Tuple      							=	self.out_layers[0][0][l_idx:r_idx]
																					__a			:    Optional[Any]      							=	self.out_layers[0][1][l_idx:r_idx]
																	else:
																					__a			:    Union[str, Any]      							=	self.out_layers[i][0]
																					__a			:    Optional[int]      							=	self.out_layers[i][1]

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

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

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

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

									return out


import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


A_			:					Any    			=     datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
	import pyspark



@dataclass
class        _a		(datasets.BuilderConfig							):


					'''simple docstring'''






					UpperCAmelCase__:  Optional[datasets.Features]    				=			None



def 		UpperCamelCase   (lowercase_:					"pyspark.sql.DataFrame"       , lowercase_:					List[int]       , )	-> Dict:
	import pyspark

	def generate_fn():
		A__				:       List[str]	  =						df.select("""*"""       , pyspark.sql.functions.spark_partition_id().alias("""part_id"""   )   )
		for partition_id in partition_order:
			A__				:       List[Any]	  =						df_with_partition_id.select("""*"""   ).where(f"""part_id = {partition_id}"""   ).drop("""part_id"""   )
			A__				:       Optional[int]	  =						partition_df.collect()
			A__				:       Any	  =						0
			for row in rows:
				yield f"""{partition_id}_{row_id}""", row.asDict()
				row_id += 1

	return generate_fn



class        _a		(_BaseExamplesIterable							):


					'''simple docstring'''






					def __init__(					self       ,			A__       ,			A__=None       ,			):
						A__				:       List[str]	  =						df
						A__				:       Optional[int]	  =						partition_order or range(self.df.rdd.getNumPartitions()     )
						A__				:       Tuple	  =						_generate_iterable_examples(self.df       ,			self.partition_order     )






					def __iter__(					self     ):
						yield from self.generate_examples_fn()






					def     __A     (					self       ,			A__     ):
						A__				:       str	  =						list(range(self.df.rdd.getNumPartitions()     )     )
						generator.shuffle(A__     )
						return SparkExamplesIterable(self.df       ,			partition_order=A__     )






					def     __A     (					self       ,			A__       ,			A__     ):
						A__				:       Optional[int]	  =						self.split_shard_indices_by_worker(A__       ,			A__     )
						return SparkExamplesIterable(self.df       ,			partition_order=A__     )






					@property
					def     __A     (					self     ):
						return len(self.partition_order     )



class        _a		(datasets.DatasetBuilder							):


					'''simple docstring'''






					UpperCAmelCase__:  Union[str, Any]    				=			SparkConfig






					def __init__(					self       ,			A__       ,			A__ = None       ,			A__ = None       ,			**A__       ,			):
						import pyspark

						A__				:       Dict	  =						pyspark.sql.SparkSession.builder.getOrCreate()
						A__				:       int	  =						df
						A__				:       Any	  =						working_dir

						super().__init__(
						    cache_dir=A__       ,			config_name=str(self.df.semanticHash()     )       ,			**A__       ,			)






					def     __A     (					self     ):
						# Returns the path of the created file.
						def create_cache_and_write_probe(A__     ):
							# makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
							# already exist.
							os.makedirs(self._cache_dir       ,			exist_ok=A__     )
							A__				:       List[str]	  =						os.path.join(self._cache_dir       ,			"""fs_test""" + uuid.uuida().hex     )
							# Opening the file in append mode will create a new file unless it already exists, in which case it will not
							# change the file contents.
							open(A__       ,			"""a"""     )
							return [probe_file]

						if self._spark.conf.get("""spark.master"""       ,			""""""     ).startswith("""local"""     ):
							return

						# If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
						# accessible to the driver.
						# TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
						if self._cache_dir:
							A__				:       int	  =						(
							    self._spark.sparkContext.parallelize(range(1     )       ,			1     ).mapPartitions(A__     ).collect()
							)
							if os.path.isfile(probe[0]     ):
								return

						raise ValueError(
						    """When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir"""     )






					def     __A     (					self     ):
						return datasets.DatasetInfo(features=self.config.features     )






					def     __A     (					self       ,			A__     ):
						return [datasets.SplitGenerator(name=datasets.Split.TRAIN     )]






					def     __A     (					self       ,			A__     ):
						import pyspark

						def get_arrow_batch_size(A__     ):
							for batch in it:
								yield pa.RecordBatch.from_pydict({"""batch_bytes""": [batch.nbytes]}     )

						A__				:       Dict	  =						self.df.count()
						A__				:       List[Any]	  =						df_num_rows if df_num_rows <= 100 else 100
						# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
						A__				:       Union[str, Any]	  =						(
						    self.df.limit(A__     )
						    .repartition(1     )
						    .mapInArrow(A__       ,			"""batch_bytes: long"""     )
						    .agg(pyspark.sql.functions.sum("""batch_bytes"""     ).alias("""sample_bytes"""     )     )
						    .collect()[0]
						    .sample_bytes
						    / sample_num_rows
						)
						A__				:       str	  =						approx_bytes_per_row * df_num_rows
						if approx_total_size > max_shard_size:
							# Make sure there is at least one row per partition.
							A__				:       Dict	  =						min(A__       ,			int(approx_total_size / max_shard_size     )     )
							A__				:       int	  =						self.df.repartition(A__     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			):
						import pyspark

						A__				:       Optional[int]	  =						ParquetWriter if file_format == """parquet""" else ArrowWriter
						A__				:       Any	  =						os.path.join(self._working_dir       ,			os.path.basename(A__     )     ) if self._working_dir else fpath
						A__				:       Union[str, Any]	  =						file_format == """parquet"""

						# Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
						# pickling the SparkContext.
						A__				:       str	  =						self.config.features
						A__				:       Union[str, Any]	  =						self._writer_batch_size
						A__				:       Tuple	  =						self._fs.storage_options

						def write_arrow(A__     ):
							# Within the same SparkContext, no two task attempts will share the same attempt ID.
							A__				:       str	  =						pyspark.TaskContext().taskAttemptId()
							A__				:       str	  =						next(A__       ,			A__     )
							if first_batch is None:
								# Some partitions might not receive any data.
								return pa.RecordBatch.from_arrays(
								    [[task_id], [0], [0]]       ,			names=["""task_id""", """num_examples""", """num_bytes"""]       ,			)
							A__				:       List[Any]	  =						0
							A__				:       Optional[Any]	  =						writer_class(
							    features=A__       ,			path=working_fpath.replace("""SSSSS"""       ,			F"""{shard_id:05d}"""     ).replace("""TTTTT"""       ,			F"""{task_id:05d}"""     )       ,			writer_batch_size=A__       ,			storage_options=A__       ,			embed_local_files=A__       ,			)
							A__				:       Tuple	  =						pa.Table.from_batches([first_batch]     )
							writer.write_table(A__     )
							for batch in it:
								if max_shard_size is not None and writer._num_bytes >= max_shard_size:
									A__  ,  A__				:       Union[str, Any]	  =						writer.finalize()
									writer.close()
									yield pa.RecordBatch.from_arrays(
									    [[task_id], [num_examples], [num_bytes]]       ,			names=["""task_id""", """num_examples""", """num_bytes"""]       ,			)
									shard_id += 1
									A__				:       Tuple	  =						writer_class(
									    features=writer._features       ,			path=working_fpath.replace("""SSSSS"""       ,			F"""{shard_id:05d}"""     ).replace("""TTTTT"""       ,			F"""{task_id:05d}"""     )       ,			writer_batch_size=A__       ,			storage_options=A__       ,			embed_local_files=A__       ,			)
								A__				:       Optional[int]	  =						pa.Table.from_batches([batch]     )
								writer.write_table(A__     )

							if writer._num_bytes > 0:
								A__  ,  A__				:       Any	  =						writer.finalize()
								writer.close()
								yield pa.RecordBatch.from_arrays(
								    [[task_id], [num_examples], [num_bytes]]       ,			names=["""task_id""", """num_examples""", """num_bytes"""]       ,			)

							if working_fpath != fpath:
								for file in os.listdir(os.path.dirname(A__     )     ):
									A__				:       Optional[Any]	  =						os.path.join(os.path.dirname(A__     )       ,			os.path.basename(A__     )     )
									shutil.move(A__       ,			A__     )

						A__				:       Tuple	  =						(
						    self.df.mapInArrow(A__       ,			"""task_id: long, num_examples: long, num_bytes: long"""     )
						    .groupBy("""task_id"""     )
						    .agg(
						        pyspark.sql.functions.sum("""num_examples"""     ).alias("""total_num_examples"""     )       ,			pyspark.sql.functions.sum("""num_bytes"""     ).alias("""total_num_bytes"""     )       ,			pyspark.sql.functions.count("""num_bytes"""     ).alias("""num_shards"""     )       ,			pyspark.sql.functions.collect_list("""num_examples"""     ).alias("""shard_lengths"""     )       ,			)
						    .collect()
						)
						for row in stats:
							yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)






					def     __A     (					self       ,			A__       ,			A__ = "arrow"       ,			A__ = None       ,			A__ = None       ,			**A__       ,			):
						self._validate_cache_dir()

						A__				:       Union[str, Any]	  =						convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE     )
						self._repartition_df_if_needed(A__     )
						A__				:       Any	  =						not is_remote_filesystem(self._fs     )
						A__				:       Optional[int]	  =						os.path.join if is_local else posixpath.join

						A__				:       Dict	  =						"""-TTTTT-SSSSS-of-NNNNN"""
						A__				:       Any	  =						F"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
						A__				:       Any	  =						path_join(self._output_dir       ,			A__     )

						A__				:       Tuple	  =						0
						A__				:       str	  =						0
						A__				:       List[Any]	  =						0
						A__				:       List[Any]	  =						[]
						A__				:       Optional[Any]	  =						[]

						for task_id, content in self._prepare_split_single(A__       ,			A__       ,			A__     ):
							(
							    (
							    A__
							)  ,  (
							    A__
							)  ,  (
							    A__
							)  ,  (
							    A__
							)  ,  
							)				:       List[Any]	  =						content
							if num_bytes > 0:
								total_num_examples += num_examples
								total_num_bytes += num_bytes
								total_shards += num_shards
								task_id_and_num_shards.append((task_id, num_shards)     )
								all_shard_lengths.extend(A__     )

						A__				:       Optional[int]	  =						total_num_examples
						A__				:       Union[str, Any]	  =						total_num_bytes

						# should rename everything at the end
						logger.debug(F"""Renaming {total_shards} shards."""     )
						if total_shards > 1:
							A__				:       int	  =						all_shard_lengths

							# Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
							# pickling error due to pickling the SparkContext.
							A__				:       Dict	  =						self._fs

							# use the -SSSSS-of-NNNNN pattern
							def _rename_shard(
							    A__       ,			A__       ,			A__       ,			):
								rename(
								    A__       ,			fpath.replace("""SSSSS"""       ,			F"""{shard_id:05d}"""     ).replace("""TTTTT"""       ,			F"""{task_id:05d}"""     )       ,			fpath.replace("""TTTTT-SSSSS"""       ,			F"""{global_shard_id:05d}"""     ).replace("""NNNNN"""       ,			F"""{total_shards:05d}"""     )       ,			)

							A__				:       List[Any]	  =						[]
							A__				:       Union[str, Any]	  =						0
							for i in range(len(A__     )     ):
								A__  ,  A__				:       Optional[int]	  =						task_id_and_num_shards[i]
								for shard_id in range(A__     ):
									args.append([task_id, shard_id, global_shard_id]     )
									global_shard_id += 1
							self._spark.sparkContext.parallelize(A__       ,			len(A__     )     ).map(lambda A__     : _rename_shard(*A__     )     ).collect()
						else:
							# don't use any pattern
							A__				:       List[Any]	  =						0
							A__				:       List[str]	  =						task_id_and_num_shards[0][0]
							self._rename(
							    fpath.replace("""SSSSS"""       ,			F"""{shard_id:05d}"""     ).replace("""TTTTT"""       ,			F"""{task_id:05d}"""     )       ,			fpath.replace(A__       ,			""""""     )       ,			)






					def     __A     (					self       ,			A__       ,			):
						return SparkExamplesIterable(self.df     )



