
'''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


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

if TYPE_CHECKING:
						import pyspark




@dataclass
class   UpperCAmelCase_					(		datasets.BuilderConfig			):






						'''simple docstring'''






						_lowercase		:				Optional[datasets.Features]									=       None






def A				(__lowerCamelCase				:"pyspark.sql.DataFrame"			,     __lowerCamelCase				:List[int]			,     ):
						import pyspark

						def generate_fn():
												_lowerCAmelCase  						=			df.select("""*"""			,     pyspark.sql.functions.spark_partition_id().alias("""part_id"""		)		)
												for partition_id in partition_order:
																		_lowerCAmelCase  						=			df_with_partition_id.select("""*"""		).where(f'part_id = {partition_id}'		).drop("""part_id"""		)
																		_lowerCAmelCase  						=			partition_df.collect()
																		_lowerCAmelCase  						=			0
																		for row in rows:
																								yield f'{partition_id}_{row_id}', row.asDict()
																								row_id += 1

						return generate_fn




class   UpperCAmelCase_					(		_BaseExamplesIterable			):






						'''simple docstring'''






						def __init__( self   ,       _lowercase   ,       _lowercase=None   ,       ):

												"""simple docstring"""



												_lowerCAmelCase  						=			df
												_lowerCAmelCase  						=			partition_order or range(self.df.rdd.getNumPartitions()				)
												_lowerCAmelCase  						=			_generate_iterable_examples(self.df   ,       self.partition_order				)




						def __iter__( self				):

												"""simple docstring"""



												yield from self.generate_examples_fn()




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			list(range(self.df.rdd.getNumPartitions()				)				)
												generator.shuffle(_lowercase				)
												return SparkExamplesIterable(self.df   ,       partition_order=_lowercase				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.split_shard_indices_by_worker(_lowercase   ,       _lowercase				)
												return SparkExamplesIterable(self.df   ,       partition_order=_lowercase				)




						@property
						def  _lowercase     ( self				):

												"""simple docstring"""



												return len(self.partition_order				)






class   UpperCAmelCase_					(		datasets.DatasetBuilder			):






						'''simple docstring'''






						_lowercase		:				str									=       SparkConfig




						def __init__( self   ,       _lowercase   ,       _lowercase = None   ,       _lowercase = None   ,       **_lowercase   ,       ):

												"""simple docstring"""



												import pyspark

												_lowerCAmelCase  						=			pyspark.sql.SparkSession.builder.getOrCreate()
												_lowerCAmelCase  						=			df
												_lowerCAmelCase  						=			working_dir

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




						def  _lowercase     ( self				):

												"""simple docstring"""



												def create_cache_and_write_probe(_lowercase				):
																		# 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=_lowercase				)
																		_lowerCAmelCase  						=			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(_lowercase   ,       """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:
																		_lowerCAmelCase  						=			(
																		    self._spark.sparkContext.parallelize(range(1				)   ,       1				).mapPartitions(_lowercase				).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				):

												"""simple docstring"""



												return datasets.DatasetInfo(features=self.config.features				)




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												return [datasets.SplitGenerator(name=datasets.Split.TRAIN				)]




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												import pyspark

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

												_lowerCAmelCase  						=			self.df.count()
												_lowerCAmelCase  						=			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.
												_lowerCAmelCase  						=			(
												    self.df.limit(_lowercase				)
												    .repartition(1				)
												    .mapInArrow(_lowercase   ,       """batch_bytes: long"""				)
												    .agg(pyspark.sql.functions.sum("""batch_bytes"""				).alias("""sample_bytes"""				)				)
												    .collect()[0]
												    .sample_bytes
												    / sample_num_rows
												)
												_lowerCAmelCase  						=			approx_bytes_per_row * df_num_rows
												if approx_total_size > max_shard_size:
																		# Make sure there is at least one row per partition.
																		_lowerCAmelCase  						=			min(_lowercase   ,       int(approx_total_size / max_shard_size				)				)
																		_lowerCAmelCase  						=			self.df.repartition(_lowercase				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       ):

												"""simple docstring"""



												import pyspark

												_lowerCAmelCase  						=			ParquetWriter if file_format == """parquet""" else ArrowWriter
												_lowerCAmelCase  						=			os.path.join(self._working_dir   ,       os.path.basename(_lowercase				)				) if self._working_dir else fpath
												_lowerCAmelCase  						=			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.
												_lowerCAmelCase  						=			self.config.features
												_lowerCAmelCase  						=			self._writer_batch_size
												_lowerCAmelCase  						=			self._fs.storage_options

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

																		if writer._num_bytes > 0:
																								_lowerCAmelCase						,   _lowerCAmelCase  						=			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(_lowercase				)				):
																														_lowerCAmelCase  						=			os.path.join(os.path.dirname(_lowercase				)   ,       os.path.basename(_lowercase				)				)
																														shutil.move(_lowercase   ,       _lowercase				)

												_lowerCAmelCase  						=			(
												    self.df.mapInArrow(_lowercase   ,       """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   ,       _lowercase   ,       _lowercase = "arrow"   ,       _lowercase = None   ,       _lowercase = None   ,       **_lowercase   ,       ):

												"""simple docstring"""



												self._validate_cache_dir()

												_lowerCAmelCase  						=			convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE				)
												self._repartition_df_if_needed(_lowercase				)
												_lowerCAmelCase  						=			not is_remote_filesystem(self._fs				)
												_lowerCAmelCase  						=			os.path.join if is_local else posixpath.join

												_lowerCAmelCase  						=			"""-TTTTT-SSSSS-of-NNNNN"""
												_lowerCAmelCase  						=			F'{self.name}-{split_generator.name}{SUFFIX}.{file_format}'
												_lowerCAmelCase  						=			path_join(self._output_dir   ,       _lowercase				)

												_lowerCAmelCase  						=			0
												_lowerCAmelCase  						=			0
												_lowerCAmelCase  						=			0
												_lowerCAmelCase  						=			[]
												_lowerCAmelCase  						=			[]

												for task_id, content in self._prepare_split_single(_lowercase   ,       _lowercase   ,       _lowercase				):
																		(
																		    (
																		    _lowerCAmelCase
																		)						,   (
																		    _lowerCAmelCase
																		)						,   (
																		    _lowerCAmelCase
																		)						,   (
																		    _lowerCAmelCase
																		)						,   
																		)  						=			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(_lowercase				)

												_lowerCAmelCase  						=			total_num_examples
												_lowerCAmelCase  						=			total_num_bytes

												# should rename everything at the end
												logger.debug(F'Renaming {total_shards} shards.'				)
												if total_shards > 1:
																		_lowerCAmelCase  						=			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.
																		_lowerCAmelCase  						=			self._fs

																		# use the -SSSSS-of-NNNNN pattern
																		def _rename_shard(
																		    _lowercase   ,       _lowercase   ,       _lowercase   ,       ):
																								rename(
																								    _lowercase   ,       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}'				)   ,       )

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




						def  _lowercase     ( self   ,       _lowercase   ,       ):

												"""simple docstring"""



												return SparkExamplesIterable(self.df				)


'''simple docstring'''



import inspect
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
_lowercase    =		"""src/transformers"""


# This is to make sure the transformers module imported is the one in the repo.
_lowercase    =		direct_transformers_import(PATH_TO_TRANSFORMERS)

_lowercase    =		transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
_lowercase    =		re.compile(R"""\[(.+?)\]\((https://huggingface\.co/.+?)\)""")


_lowercase    =		{
    """DecisionTransformerConfig""",
    """EncoderDecoderConfig""",
    """MusicgenConfig""",
    """RagConfig""",
    """SpeechEncoderDecoderConfig""",
    """TimmBackboneConfig""",
    """VisionEncoderDecoderConfig""",
    """VisionTextDualEncoderConfig""",
    """LlamaConfig""",
}


def A				(__lowerCamelCase				:str		):
						_lowerCAmelCase  						=			None

						# source code of `config_class`
						_lowerCAmelCase  						=			inspect.getsource(__lowerCamelCase		)
						_lowerCAmelCase  						=			_re_checkpoint.findall(__lowerCamelCase		)

						# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
						# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
						for ckpt_name, ckpt_link in checkpoints:
												# allow the link to end with `/`
												if ckpt_link.endswith("""/"""		):
																		_lowerCAmelCase  						=			ckpt_link[:-1]

												# verify the checkpoint name corresponds to the checkpoint link
												_lowerCAmelCase  						=			f'https://huggingface.co/{ckpt_name}'
												if ckpt_link == ckpt_link_from_name:
																		_lowerCAmelCase  						=			ckpt_name
																		break

						return checkpoint


def A				():
						_lowerCAmelCase  						=			[]

						for config_class in list(CONFIG_MAPPING.values()		):
												# Skip deprecated models
												if "models.deprecated" in config_class.__module__:
																		continue
												_lowerCAmelCase  						=			get_checkpoint_from_config_class(__lowerCamelCase		)

												_lowerCAmelCase  						=			config_class.__name__
												if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
																		configs_without_checkpoint.append(__lowerCamelCase		)

						if len(__lowerCamelCase		) > 0:
												_lowerCAmelCase  						=			"""\n""".join(sorted(__lowerCamelCase		)		)
												raise ValueError(f'The following configurations don\'t contain any valid checkpoint:\n{message}'		)


if __name__ == "__main__":
						check_config_docstrings_have_checkpoints()
