





'''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 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__       ):







  '''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(SCREAMING_SNAKE_CASE__       ):
      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 				snake_case_      (		):







  '''simple docstring'''



  _snake_case     		=				pyspark.sql.SparkSession.builder.master("local[*]"       ).appName("pyspark"       ).getOrCreate()
  _snake_case     		=				spark.range(1_00       ).repartition(1       )
  _snake_case     		=				Spark(SCREAMING_SNAKE_CASE__       )
  # 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 				snake_case_      (		):







  '''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(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__       )  # Reverse the partitions.
  _snake_case     		=				_get_expected_row_ids_and_row_dicts_for_partition_order(SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__       )

  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 				snake_case_      (		):







  '''simple docstring'''



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




@require_not_windows
@require_dill_gt_0_3_2
def 				snake_case_      (		):







  '''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 SCREAMING_SNAKE_CASE__       : x.reverse()
    _snake_case     		=				_get_expected_row_ids_and_row_dicts_for_partition_order(SCREAMING_SNAKE_CASE__					,	[2, 1, 0]       )

    _snake_case     		=				SparkExamplesIterable(SCREAMING_SNAKE_CASE__       ).shuffle_data_sources(SCREAMING_SNAKE_CASE__       )
    assert shuffled_it.n_shards == 3
    for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__       ):
      _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 				snake_case_      (		):







  '''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(SCREAMING_SNAKE_CASE__       ).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(SCREAMING_SNAKE_CASE__					,	[0, 2]       )
  for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__       ):
    _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(SCREAMING_SNAKE_CASE__       ).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(SCREAMING_SNAKE_CASE__					,	[1, 3]       )
  for i, (row_id, row_dict) in enumerate(SCREAMING_SNAKE_CASE__       ):
    _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 				snake_case_      (		):







  '''simple docstring'''



  _snake_case     		=				pyspark.sql.SparkSession.builder.master("local[*]"       ).appName("pyspark"       ).getOrCreate()
  _snake_case     		=				spark.range(1_00       ).repartition(1       )
  _snake_case     		=				Spark(SCREAMING_SNAKE_CASE__       )
  # 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









'''simple docstring'''



from collections import OrderedDict
from typing import Mapping

from packaging import version

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


__magic_name__   :						Dict    =			logging.get_logger(__name__)

__magic_name__   :						Dict    =			{
    """google/mobilenet_v1_1.0_224""": """https://huggingface.co/google/mobilenet_v1_1.0_224/resolve/main/config.json""",
    """google/mobilenet_v1_0.75_192""": """https://huggingface.co/google/mobilenet_v1_0.75_192/resolve/main/config.json""",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
}


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						List[Any]      =   '''mobilenet_v1'''

     def __init__(					self						,   lowerCamelCase=3						,   lowerCamelCase=224						,   lowerCamelCase=1.0						,   lowerCamelCase=8						,   lowerCamelCase="relu6"						,   lowerCamelCase=True						,   lowerCamelCase=0.999						,   lowerCamelCase=0.02						,   lowerCamelCase=0.001						,   **lowerCamelCase						,   ):
       super().__init__(**lowerCamelCase	)

       if depth_multiplier <= 0:
         raise ValueError("depth_multiplier must be greater than zero."	)

       _snake_case     		=				num_channels
       _snake_case     		=				image_size
       _snake_case     		=				depth_multiplier
       _snake_case     		=				min_depth
       _snake_case     		=				hidden_act
       _snake_case     		=				tf_padding
       _snake_case     		=				classifier_dropout_prob
       _snake_case     		=				initializer_range
       _snake_case     		=				layer_norm_eps







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						int      =   version.parse('''1.11'''			)

     @property
     def 	UpperCamelCase(					self	):
       return OrderedDict([("pixel_values", {0: "batch"})]	)

     @property
     def 	UpperCamelCase(					self	):
       if self.task == "image-classification":
         return OrderedDict([("logits", {0: "batch"})]	)
       else:
         return OrderedDict([("last_hidden_state", {0: "batch"}), ("pooler_output", {0: "batch"})]	)






     @property
     def 	UpperCamelCase(					self	):
       return 1e-4


