



'''simple docstring'''

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


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

if TYPE_CHECKING:
							import pyspark

@dataclass
class      __magic_name__						(   datasets.BuilderConfig       ):
							lowerCAmelCase	:   Optional[datasets.Features]           = None



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_						,      )      ->				Optional[int]:
		import pyspark

		def generate_fn():
				_a					:				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      __magic_name__						(   _BaseExamplesIterable       ):







							def __init__(							self :	Dict     ,_UpperCAmelCase :	"pyspark.sql.DataFrame"     ,_UpperCAmelCase :	Any=None     ,):
									_a					:				List[str]				=							df
									_a					:				str				=							partition_order or range(self.df.rdd.getNumPartitions()		)
									_a					:				Union[str, Any]				=							_generate_iterable_examples(self.df     ,self.partition_order		)







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







							def 							__lowercase      (							self :	str     ,_UpperCAmelCase :	np.random.Generator		):
									_a					:				int				=							list(range(self.df.rdd.getNumPartitions()		)		)
									generator.shuffle(_UpperCAmelCase		)
									return SparkExamplesIterable(self.df     ,partition_order=_UpperCAmelCase		)







							def 							__lowercase      (							self :	Optional[int]     ,_UpperCAmelCase :	int     ,_UpperCAmelCase :	int		):
									_a					:				Union[str, Any]				=							self.split_shard_indices_by_worker(_UpperCAmelCase     ,_UpperCAmelCase		)
									return SparkExamplesIterable(self.df     ,partition_order=_UpperCAmelCase		)







							@property
							def 							__lowercase      (							self :	Union[str, Any]		):
									return len(self.partition_order		)

class      __magic_name__						(   datasets.DatasetBuilder       ):
							lowerCAmelCase	:   Optional[Any]           = SparkConfig







							def __init__(							self :	List[Any]     ,_UpperCAmelCase :	"pyspark.sql.DataFrame"     ,_UpperCAmelCase :	str = None     ,_UpperCAmelCase :	str = None     ,**_UpperCAmelCase :	int     ,):
									import pyspark

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

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







							def 							__lowercase      (							self :	int		):
									# Returns the path of the created file.
									def create_cache_and_write_probe(_UpperCAmelCase :	int		):
											# 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=_UpperCAmelCase		)
											_a					:				List[Any]				=							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(_UpperCAmelCase     ,'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					:				List[str]				=							(
											    self._spark.sparkContext.parallelize(range(1		)     ,1		).mapPartitions(_UpperCAmelCase		).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 							__lowercase      (							self :	List[Any]		):
									return datasets.DatasetInfo(features=self.config.features		)







							def 							__lowercase      (							self :	Any     ,_UpperCAmelCase :	datasets.download.download_manager.DownloadManager		):
									return [datasets.SplitGenerator(name=datasets.Split.TRAIN		)]







							def 							__lowercase      (							self :	Optional[Any]     ,_UpperCAmelCase :	List[str]		):
									import pyspark

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

									_a					:				int				=							self.df.count()
									_a					:				Dict				=							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					:				Optional[Any]				=							(
									    self.df.limit(_UpperCAmelCase		)
									    .repartition(1		)
									    .mapInArrow(_UpperCAmelCase     ,'batch_bytes: long'		)
									    .agg(pyspark.sql.functions.sum('batch_bytes'		).alias('sample_bytes'		)		)
									    .collect()[0]
									    .sample_bytes
									    / sample_num_rows
									)
									_a					:				Optional[int]				=							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					:				List[Any]				=							min(_UpperCAmelCase     ,int(approx_total_size / max_shard_size		)		)
											_a					:				Union[str, Any]				=							self.df.repartition(_UpperCAmelCase		)







							def 							__lowercase      (							self :	Optional[Any]     ,_UpperCAmelCase :	str     ,_UpperCAmelCase :	str     ,_UpperCAmelCase :	int     ,):
									import pyspark

									_a					:				int				=							ParquetWriter if file_format == 'parquet' else ArrowWriter
									_a					:				Any				=							os.path.join(self._working_dir     ,os.path.basename(_UpperCAmelCase		)		) if self._working_dir else fpath
									_a					:				Optional[int]				=							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					:				List[Any]				=							self.config.features
									_a					:				Optional[Any]				=							self._writer_batch_size
									_a					:				str				=							self._fs.storage_options

									def write_arrow(_UpperCAmelCase :	str		):
											# Within the same SparkContext, no two task attempts will share the same attempt ID.
											_a					:				Optional[int]				=							pyspark.TaskContext().taskAttemptId()
											_a					:				Optional[Any]				=							next(_UpperCAmelCase     ,_UpperCAmelCase		)
											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					:				Union[str, Any]				=							0
											_a					:				Any				=							writer_class(
											    features=_UpperCAmelCase     ,path=working_fpath.replace('SSSSS'     ,F"""{shard_id:05d}"""		).replace('TTTTT'     ,F"""{task_id:05d}"""		)     ,writer_batch_size=_UpperCAmelCase     ,storage_options=_UpperCAmelCase     ,embed_local_files=_UpperCAmelCase     ,)
											_a					:				Any				=							pa.Table.from_batches([first_batch]		)
											writer.write_table(_UpperCAmelCase		)
											for batch in it:
													if max_shard_size is not None and writer._num_bytes >= max_shard_size:
															_a					:				str				=							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					:				List[str]				=							writer_class(
															    features=writer._features     ,path=working_fpath.replace('SSSSS'     ,F"""{shard_id:05d}"""		).replace('TTTTT'     ,F"""{task_id:05d}"""		)     ,writer_batch_size=_UpperCAmelCase     ,storage_options=_UpperCAmelCase     ,embed_local_files=_UpperCAmelCase     ,)
													_a					:				Optional[Any]				=							pa.Table.from_batches([batch]		)
													writer.write_table(_UpperCAmelCase		)

											if writer._num_bytes > 0:
													_a					:				List[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(_UpperCAmelCase		)		):
															_a					:				Tuple				=							os.path.join(os.path.dirname(_UpperCAmelCase		)     ,os.path.basename(_UpperCAmelCase		)		)
															shutil.move(_UpperCAmelCase     ,_UpperCAmelCase		)

									_a					:				List[Any]				=							(
									    self.df.mapInArrow(_UpperCAmelCase     ,'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 							__lowercase      (							self :	int     ,_UpperCAmelCase :	"datasets.SplitGenerator"     ,_UpperCAmelCase :	str = "arrow"     ,_UpperCAmelCase :	Optional[Union[str, int]] = None     ,_UpperCAmelCase :	Optional[int] = None     ,**_UpperCAmelCase :	int     ,):
									self._validate_cache_dir()

									_a					:				List[Any]				=							convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE		)
									self._repartition_df_if_needed(_UpperCAmelCase		)
									_a					:				int				=							not is_remote_filesystem(self._fs		)
									_a					:				List[Any]				=							os.path.join if is_local else posixpath.join

									_a					:				List[str]				=							'-TTTTT-SSSSS-of-NNNNN'
									_a					:				Optional[int]				=							F"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
									_a					:				Union[str, Any]				=							path_join(self._output_dir     ,_UpperCAmelCase		)

									_a					:				Dict				=							0
									_a					:				List[Any]				=							0
									_a					:				Tuple				=							0
									_a					:				Dict				=							[]
									_a					:				List[Any]				=							[]

									for task_id, content in self._prepare_split_single(_UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase		):
											(
											    _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(_UpperCAmelCase		)

									_a					:				Any				=							total_num_examples
									_a					:				Optional[int]				=							total_num_bytes

									# should rename everything at the end
									logger.debug(F"""Renaming {total_shards} shards."""		)
									if total_shards > 1:
											_a					:				List[Any]				=							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					:				Union[str, Any]				=							self._fs

											# use the -SSSSS-of-NNNNN pattern
											def _rename_shard(
											    _UpperCAmelCase :	int     ,_UpperCAmelCase :	int     ,_UpperCAmelCase :	int     ,):
													rename(
													    _UpperCAmelCase     ,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					:				Optional[Any]				=							[]
											_a					:				Optional[Any]				=							0
											for i in range(len(_UpperCAmelCase		)		):
													_a					:				Optional[int]				=							task_id_and_num_shards[i]
													for shard_id in range(_UpperCAmelCase		):
															args.append([task_id, shard_id, global_shard_id]		)
															global_shard_id += 1
											self._spark.sparkContext.parallelize(_UpperCAmelCase     ,len(_UpperCAmelCase		)		).map(lambda _UpperCAmelCase		: _rename_shard(*_UpperCAmelCase		)		).collect()
									else:
											# don't use any pattern
											_a					:				str				=							0
											_a					:				Union[str, Any]				=							task_id_and_num_shards[0][0]
											self._rename(
											    fpath.replace('SSSSS'     ,F"""{shard_id:05d}"""		).replace('TTTTT'     ,F"""{task_id:05d}"""		)     ,fpath.replace(_UpperCAmelCase     ,''		)     ,)







							def 							__lowercase      (							self :	List[str]     ,_UpperCAmelCase :	"datasets.SplitGenerator"     ,):
									return SparkExamplesIterable(self.df		)





'''simple docstring'''

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.eager import context
							from tensorflow.python.framework import ops

							from transformers import GradientAccumulator, create_optimizer

@require_tf
class      __magic_name__						(   unittest.TestCase       ):







							def 							__lowercase      (							self :	int     ,_UpperCAmelCase :	int     ,_UpperCAmelCase :	str     ,_UpperCAmelCase :	Optional[Any]		):
									self.assertEqual(len(_UpperCAmelCase		)     ,len(_UpperCAmelCase		)		)
									for a, b in zip(_UpperCAmelCase     ,_UpperCAmelCase		):
											self.assertAlmostEqual(_UpperCAmelCase     ,_UpperCAmelCase     ,delta=_UpperCAmelCase		)







							def 							__lowercase      (							self :	int		):
									_a					:				int				=							GradientAccumulator()
									accumulator([tf.constant([1.0, 2.0]		)]		)
									accumulator([tf.constant([-2.0, 1.0]		)]		)
									accumulator([tf.constant([-1.0, 2.0]		)]		)
									with self.assertRaises(_UpperCAmelCase		):
											accumulator([tf.constant([1.0, 1.0]		), tf.constant([2.0, 2.0]		)]		)
									self.assertEqual(accumulator.step     ,3		)
									self.assertEqual(len(accumulator.gradients		)     ,1		)
									self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()     ,[-2.0, 5.0]     ,tol=1E-2		)
									accumulator.reset()
									self.assertEqual(accumulator.step     ,0		)
									self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()     ,[0.0, 0.0]     ,tol=1E-2		)







							def 							__lowercase      (							self :	Any		):
									_a					:				int				=							None
									ops.enable_eager_execution_internal()
									_a					:				Optional[int]				=							tf.config.list_physical_devices('CPU'		)
									if len(_UpperCAmelCase		) == 1:
											tf.config.set_logical_device_configuration(
											    physical_devices[0]     ,[tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]		)
									_a					:				Tuple				=							tf.config.list_logical_devices(device_type='CPU'		)
									_a					:				List[str]				=							tf.distribute.MirroredStrategy(devices=devices[:2]		)

									with strategy.scope():
											_a					:				Tuple				=							GradientAccumulator()
											_a					:				List[Any]				=							tf.Variable([4.0, 3.0]		)
											_a	, _a					:				Dict				=							create_optimizer(5E-5     ,10     ,5		)
											_a					:				Tuple				=							tf.Variable([0.0, 0.0]     ,trainable=_UpperCAmelCase		)

									def accumulate_on_replica(_UpperCAmelCase :	str		):
											accumulator([gradient]		)

									def apply_on_replica():
											optimizer.apply_gradients(list(zip(accumulator.gradients     ,[variable]		)		)		)

									@tf.function
									def accumulate(_UpperCAmelCase :	Optional[int]     ,_UpperCAmelCase :	Optional[int]		):
											with strategy.scope():
													_a					:				Union[str, Any]				=							strategy.experimental_local_results(_UpperCAmelCase		)
													local_variables[0].assign(_UpperCAmelCase		)
													local_variables[1].assign(_UpperCAmelCase		)
													strategy.run(_UpperCAmelCase     ,args=(gradient_placeholder,)		)

									@tf.function
									def apply_grad():
											with strategy.scope():
													strategy.run(_UpperCAmelCase		)

									def _check_local_values(_UpperCAmelCase :	Tuple     ,_UpperCAmelCase :	List[str]		):
											_a					:				List[Any]				=							strategy.experimental_local_results(accumulator._gradients[0]		)
											self.assertListAlmostEqual(values[0].value()     ,_UpperCAmelCase     ,tol=1E-2		)
											self.assertListAlmostEqual(values[1].value()     ,_UpperCAmelCase     ,tol=1E-2		)

									accumulate([1.0, 2.0]     ,[-1.0, 1.0]		)
									accumulate([3.0, -1.0]     ,[-1.0, -1.0]		)
									accumulate([-2.0, 2.0]     ,[3.0, -2.0]		)
									self.assertEqual(accumulator.step     ,3		)
									_check_local_values([2.0, 3.0]     ,[1.0, -2.0]		)
									apply_grad()
									self.assertListAlmostEqual(variable.value()     ,[4.0, 3.0]     ,tol=1E-2		)
									accumulator.reset()
									self.assertEqual(accumulator.step     ,0		)
									_check_local_values([0.0, 0.0]     ,[0.0, 0.0]		)




