




'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase_							:	List[str]	     =		{
    '''configuration_roberta''': ['''ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RobertaConfig''', '''RobertaOnnxConfig'''],
    '''tokenization_roberta''': ['''RobertaTokenizer'''],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase_							:	Union[str, Any]	     =		['''RobertaTokenizerFast''']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase_							:	Any	     =		[
         '''ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''RobertaForCausalLM''',
         '''RobertaForMaskedLM''',
         '''RobertaForMultipleChoice''',
         '''RobertaForQuestionAnswering''',
         '''RobertaForSequenceClassification''',
         '''RobertaForTokenClassification''',
         '''RobertaModel''',
         '''RobertaPreTrainedModel''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase_							:	Any	     =		[
         '''TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFRobertaForCausalLM''',
         '''TFRobertaForMaskedLM''',
         '''TFRobertaForMultipleChoice''',
         '''TFRobertaForQuestionAnswering''',
         '''TFRobertaForSequenceClassification''',
         '''TFRobertaForTokenClassification''',
         '''TFRobertaMainLayer''',
         '''TFRobertaModel''',
         '''TFRobertaPreTrainedModel''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase_							:	Optional[int]	     =		[
         '''FlaxRobertaForCausalLM''',
         '''FlaxRobertaForMaskedLM''',
         '''FlaxRobertaForMultipleChoice''',
         '''FlaxRobertaForQuestionAnswering''',
         '''FlaxRobertaForSequenceClassification''',
         '''FlaxRobertaForTokenClassification''',
         '''FlaxRobertaModel''',
         '''FlaxRobertaPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaOnnxConfig
     from .tokenization_roberta import RobertaTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_roberta_fast import RobertaTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_roberta import (
              ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
              RobertaForCausalLM,
              RobertaForMaskedLM,
              RobertaForMultipleChoice,
              RobertaForQuestionAnswering,
              RobertaForSequenceClassification,
              RobertaForTokenClassification,
              RobertaModel,
              RobertaPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_roberta import (
              TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFRobertaForCausalLM,
              TFRobertaForMaskedLM,
              TFRobertaForMultipleChoice,
              TFRobertaForQuestionAnswering,
              TFRobertaForSequenceClassification,
              TFRobertaForTokenClassification,
              TFRobertaMainLayer,
              TFRobertaModel,
              TFRobertaPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_roberta import (
              FlaxRobertaForCausalLM,
              FlaxRobertaForMaskedLM,
              FlaxRobertaForMultipleChoice,
              FlaxRobertaForQuestionAnswering,
              FlaxRobertaForSequenceClassification,
              FlaxRobertaForTokenClassification,
              FlaxRobertaModel,
              FlaxRobertaPreTrainedModel,
          )

else:
     import sys

     UpperCAmelCase_							:	str	     =		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)












'''simple docstring'''



from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)







def       _UpperCamelCase			(_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  List[Any]						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	[]
     for part_id in partition_order:
          __snake_case            =	df.where(f'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
          for row_idx, row in enumerate(_lowerCamelCase						):
               expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())						)
     return expected_row_ids_and_row_dicts







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Any:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16						)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(2						)
     __snake_case            =	[1, 0]
     __snake_case            =	_generate_iterable_examples(_lowerCamelCase ,						_lowerCamelCase						)  # Reverse the partitions.
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						_lowerCamelCase						)

     for i, (row_id, row_dict) in enumerate(generate_fn()						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> int:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(10						).repartition(1						)
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          assert row_id == f'''0_{i}'''
          assert row_dict == {"id": i}







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Union[str, Any]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(30						).repartition(3						)
     # Mock the generator so that shuffle reverses the partition indices.
     with patch('''numpy.random.Generator'''						) as generator_mock:
          __snake_case            =	lambda _lowerCamelCase						: x.reverse()
          __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[2, 1, 0]						)

          __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shuffle_data_sources(_lowerCamelCase						)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
               __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Tuple:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(20						).repartition(4						)

     # Partitions 0 and 2
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=0 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[0, 2]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     __snake_case            =	SparkExamplesIterable(_lowerCamelCase						).shard_data_sources(worker_id=1 ,						num_workers=2						)
     assert shard_it_a.n_shards == 2
     __snake_case            =	_get_expected_row_ids_and_row_dicts_for_partition_order(_lowerCamelCase ,						[1, 3]						)
     for i, (row_id, row_dict) in enumerate(_lowerCamelCase						):
          __snake_case       ,     __snake_case            =	expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Optional[int]:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1						)
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 1_00






