





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		:  Dict			       =datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
		import pyspark






@dataclass
class 						UpperCamelCase_  (  datasets.BuilderConfig							):
					_a      :						Optional[datasets.Features]						    =					None


def 		_SCREAMING_SNAKE_CASE     ( lowerCAmelCase__							,lowerCAmelCase__							,):
		import pyspark

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

		return generate_fn






class 						UpperCamelCase_  (  _BaseExamplesIterable							):



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



					def __iter__(		self     :     Optional[int]						):
							yield from self.generate_examples_fn()



					def   __a     (		self     :     int       ,      lowerCamelCase     :     np.random.Generator						):
							lowerCamelCase_	:     Dict 						=      list(range(self.df.rdd.getNumPartitions()						)						)
							generator.shuffle(lowerCamelCase						)
							return SparkExamplesIterable(self.df       ,      partition_order=lowerCamelCase						)



					def   __a     (		self     :     Any       ,      lowerCamelCase     :     int       ,      lowerCamelCase     :     int						):
							lowerCamelCase_	:     str 						=      self.split_shard_indices_by_worker(lowerCamelCase       ,      lowerCamelCase						)
							return SparkExamplesIterable(self.df       ,      partition_order=lowerCamelCase						)





					@property
					def   __a     (		self     :     Any						):
							return len(self.partition_order						)






class 						UpperCamelCase_  (  datasets.DatasetBuilder							):
					_a      :						str						    =					SparkConfig



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

							lowerCamelCase_	:     Any 						=      pyspark.sql.SparkSession.builder.getOrCreate()
							lowerCamelCase_	:     Tuple 						=      df
							lowerCamelCase_	:     List[Any] 						=      working_dir

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



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



					def   __a     (		self     :     Dict       ,      lowerCamelCase     :     datasets.download.download_manager.DownloadManager						):
							return [datasets.SplitGenerator(name=datasets.Split.TRAIN						)]



					def   __a     (		self     :     Tuple       ,      lowerCamelCase     :     Dict						):
							import pyspark

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

							lowerCamelCase_	:     List[Any] 						=      self.df.count()
							lowerCamelCase_	:     Any 						=      df_num_rows if df_num_rows <= 1_00 else 1_00
							# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
							lowerCamelCase_	:     Any 						=      (
							    self.df.limit(lowerCamelCase						)
							    .repartition(1						)
							    .mapInArrow(lowerCamelCase       ,      'batch_bytes: long'						)
							    .agg(pyspark.sql.functions.sum('batch_bytes'						).alias('sample_bytes'						)						)
							    .collect()[0]
							    .sample_bytes
							    / sample_num_rows
							)
							lowerCamelCase_	:     List[Any] 						=      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_	:     Optional[Any] 						=      min(lowerCamelCase       ,      int(approx_total_size / max_shard_size						)						)
									lowerCamelCase_	:     str 						=      self.df.repartition(lowerCamelCase						)



					def   __a     (		self     :     Union[str, Any]       ,      lowerCamelCase     :     str       ,      lowerCamelCase     :     str       ,      lowerCamelCase     :     int       ,      ):
							import pyspark

							lowerCamelCase_	:     Tuple 						=      ParquetWriter if file_format == 'parquet' else ArrowWriter
							lowerCamelCase_	:     Tuple 						=      os.path.join(self._working_dir       ,      os.path.basename(lowerCamelCase						)						) if self._working_dir else fpath
							lowerCamelCase_	:     str 						=      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_	:     Union[str, Any] 						=      self.config.features
							lowerCamelCase_	:     Dict 						=      self._writer_batch_size
							lowerCamelCase_	:     Optional[Any] 						=      self._fs.storage_options

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

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

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

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

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

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

							for task_id, content in self._prepare_split_single(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase						):
									(
									    (
									    lowerCamelCase_
									)						,      (
									    lowerCamelCase_
									)						,      (
									    lowerCamelCase_
									)						,      (
									    lowerCamelCase_
									)						,      
									)	:     Optional[int] 						=      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(lowerCamelCase						)

							lowerCamelCase_	:     Any 						=      total_num_examples
							lowerCamelCase_	:     Any 						=      total_num_bytes

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

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





					def   __a     (		self     :     Optional[int]       ,      lowerCamelCase     :     "datasets.SplitGenerator"       ,      ):
							return SparkExamplesIterable(self.df						)







from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


if TYPE_CHECKING:
		from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType


_lowercase		:  Tuple			       =logging.get_logger(__name__)

_lowercase		:  Any			       ={
    """microsoft/deberta-v2-xlarge""": """https://huggingface.co/microsoft/deberta-v2-xlarge/resolve/main/config.json""",
    """microsoft/deberta-v2-xxlarge""": """https://huggingface.co/microsoft/deberta-v2-xxlarge/resolve/main/config.json""",
    """microsoft/deberta-v2-xlarge-mnli""": (
        """https://huggingface.co/microsoft/deberta-v2-xlarge-mnli/resolve/main/config.json"""
    ),
    """microsoft/deberta-v2-xxlarge-mnli""": (
        """https://huggingface.co/microsoft/deberta-v2-xxlarge-mnli/resolve/main/config.json"""
    ),
}






class 						UpperCamelCase_  (  snake_case__							):
					_a      :						int						    =					'deberta-v2'



					def __init__(		self     :     Any       ,      lowerCamelCase     :     Dict=12_81_00       ,      lowerCamelCase     :     List[Any]=15_36       ,      lowerCamelCase     :     Union[str, Any]=24       ,      lowerCamelCase     :     Union[str, Any]=24       ,      lowerCamelCase     :     Optional[int]=61_44       ,      lowerCamelCase     :     Any="gelu"       ,      lowerCamelCase     :     int=0.1       ,      lowerCamelCase     :     str=0.1       ,      lowerCamelCase     :     Optional[int]=5_12       ,      lowerCamelCase     :     str=0       ,      lowerCamelCase     :     Union[str, Any]=0.02       ,      lowerCamelCase     :     List[str]=1E-7       ,      lowerCamelCase     :     Dict=False       ,      lowerCamelCase     :     List[Any]=-1       ,      lowerCamelCase     :     Union[str, Any]=0       ,      lowerCamelCase     :     Optional[int]=True       ,      lowerCamelCase     :     Optional[int]=None       ,      lowerCamelCase     :     Tuple=0       ,      lowerCamelCase     :     Tuple="gelu"       ,      **lowerCamelCase     :     Optional[int]       ,      ):
							super().__init__(**lowerCamelCase						)

							lowerCamelCase_	:     Any 						=      hidden_size
							lowerCamelCase_	:     Dict 						=      num_hidden_layers
							lowerCamelCase_	:     Optional[int] 						=      num_attention_heads
							lowerCamelCase_	:     Optional[int] 						=      intermediate_size
							lowerCamelCase_	:     str 						=      hidden_act
							lowerCamelCase_	:     int 						=      hidden_dropout_prob
							lowerCamelCase_	:     Tuple 						=      attention_probs_dropout_prob
							lowerCamelCase_	:     Union[str, Any] 						=      max_position_embeddings
							lowerCamelCase_	:     Union[str, Any] 						=      type_vocab_size
							lowerCamelCase_	:     List[Any] 						=      initializer_range
							lowerCamelCase_	:     Optional[Any] 						=      relative_attention
							lowerCamelCase_	:     List[Any] 						=      max_relative_positions
							lowerCamelCase_	:     Optional[int] 						=      pad_token_id
							lowerCamelCase_	:     int 						=      position_biased_input

							# Backwards compatibility
							if type(lowerCamelCase						) == str:
									lowerCamelCase_	:     Dict 						=      [x.strip() for x in pos_att_type.lower().split('|'						)]

							lowerCamelCase_	:     Optional[Any] 						=      pos_att_type
							lowerCamelCase_	:     List[Any] 						=      vocab_size
							lowerCamelCase_	:     Any 						=      layer_norm_eps

							lowerCamelCase_	:     int 						=      kwargs.get('pooler_hidden_size'       ,      lowerCamelCase						)
							lowerCamelCase_	:     Dict 						=      pooler_dropout
							lowerCamelCase_	:     Tuple 						=      pooler_hidden_act






class 						UpperCamelCase_  (  snake_case__							):



					@property
					def   __a     (		self     :     Optional[int]						):
							if self.task == "multiple-choice":
									lowerCamelCase_	:     Optional[Any] 						=      {0: 'batch', 1: 'choice', 2: 'sequence'}
							else:
									lowerCamelCase_	:     Any 						=      {0: 'batch', 1: 'sequence'}
							if self._config.type_vocab_size > 0:
									return OrderedDict(
									    [('input_ids', dynamic_axis), ('attention_mask', dynamic_axis), ('token_type_ids', dynamic_axis)]						)
							else:
									return OrderedDict([('input_ids', dynamic_axis), ('attention_mask', dynamic_axis)]						)



					@property
					def   __a     (		self     :     List[Any]						):
							return 12





					def   __a     (		self     :     Tuple       ,      lowerCamelCase     :     Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]       ,      lowerCamelCase     :     int = -1       ,      lowerCamelCase     :     int = -1       ,      lowerCamelCase     :     int = -1       ,      lowerCamelCase     :     bool = False       ,      lowerCamelCase     :     Optional["TensorType"] = None       ,      lowerCamelCase     :     int = 3       ,      lowerCamelCase     :     int = 40       ,      lowerCamelCase     :     int = 40       ,      lowerCamelCase     :     "PreTrainedTokenizerBase" = None       ,      ):
							lowerCamelCase_	:     Union[str, Any] 						=      super().generate_dummy_inputs(preprocessor=lowerCamelCase       ,      framework=lowerCamelCase						)
							if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
									del dummy_inputs["token_type_ids"]
							return dummy_inputs
