


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_							, lowerCamelCase_ )    ->						List[Any]:
    _lowercase       :							List[Any]      		=		[]
    for part_id in partition_order:
        _lowercase       :							List[str]      		=		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_(      )    ->						Union[str, Any]:
    _lowercase       :							Optional[int]      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							Any      		=		spark.range(100 ).repartition(1 )
    _lowercase       :							int      		=		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_(      )    ->						str:
    _lowercase       :							Any      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							Union[str, Any]      		=		spark.range(10 ).repartition(2 )
    _lowercase       :							Optional[Any]      		=		[1, 0]
    _lowercase       :							int      		=		_generate_iterable_examples(lowerCamelCase_							, lowerCamelCase_ )  # Reverse the partitions.
    _lowercase       :							Tuple      		=		_get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase_							, lowerCamelCase_ )

    for i, (row_id, row_dict) in enumerate(generate_fn() ):
        _lowercase		,				_lowercase       :							Tuple      		=		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:
    _lowercase       :							Optional[Any]      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							Dict      		=		spark.range(10 ).repartition(1 )
    _lowercase       :							Any      		=		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_(      )    ->						str:
    _lowercase       :							Optional[int]      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							str      		=		spark.range(30 ).repartition(3 )
    # Mock the generator so that shuffle reverses the partition indices.
    with patch('numpy.random.Generator' ) as generator_mock:
        _lowercase       :							Dict      		=		lambda lowerCamelCase_ : x.reverse()
        _lowercase       :							Dict      		=		_get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase_							, [2, 1, 0] )

        _lowercase       :							Dict      		=		SparkExamplesIterable(lowerCamelCase_ ).shuffle_data_sources(lowerCamelCase_ )
        assert shuffled_it.n_shards == 3
        for i, (row_id, row_dict) in enumerate(lowerCamelCase_ ):
            _lowercase		,				_lowercase       :							Any      		=		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_(      )    ->						Union[str, Any]:
    _lowercase       :							List[str]      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							int      		=		spark.range(20 ).repartition(4 )

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

    # Partitions 1 and 3
    _lowercase       :							Union[str, Any]      		=		SparkExamplesIterable(lowerCamelCase_ ).shard_data_sources(worker_id=1							, num_workers=2 )
    assert shard_it_a.n_shards == 2
    _lowercase       :							str      		=		_get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase_							, [1, 3] )
    for i, (row_id, row_dict) in enumerate(lowerCamelCase_ ):
        _lowercase		,				_lowercase       :							Union[str, Any]      		=		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]:
    _lowercase       :							Union[str, Any]      		=		pyspark.sql.SparkSession.builder.master('local[*]' ).appName('pyspark' ).getOrCreate()
    _lowercase       :							List[str]      		=		spark.range(100 ).repartition(1 )
    _lowercase       :							Dict      		=		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() == 100






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


SCREAMING_SNAKE_CASE       :      int												= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE       :      List[str]												= {
    "SenseTime/deformable-detr": "https://huggingface.co/sensetime/deformable-detr/resolve/main/config.json",
    # See all Deformable DETR models at https://huggingface.co/models?filter=deformable-detr
}






class _lowerCamelCase(						_a   ):
       lowercase_ :	Dict     =			"""deformable_detr"""
       lowercase_ :	int     =			{
           """hidden_size""": """d_model""",
           """num_attention_heads""": """encoder_attention_heads""",
       }
       def __init__(					self,      lowerCamelCase=True,      lowerCamelCase=None,      lowerCamelCase=3,      lowerCamelCase=3_00,      lowerCamelCase=10_24,      lowerCamelCase=6,      lowerCamelCase=10_24,      lowerCamelCase=8,      lowerCamelCase=6,      lowerCamelCase=10_24,      lowerCamelCase=8,      lowerCamelCase=0.0,      lowerCamelCase=True,      lowerCamelCase="relu",      lowerCamelCase=2_56,      lowerCamelCase=0.1,      lowerCamelCase=0.0,      lowerCamelCase=0.0,      lowerCamelCase=0.0_2,      lowerCamelCase=1.0,      lowerCamelCase=True,      lowerCamelCase=False,      lowerCamelCase="sine",      lowerCamelCase="resnet50",      lowerCamelCase=True,      lowerCamelCase=False,      lowerCamelCase=4,      lowerCamelCase=4,      lowerCamelCase=4,      lowerCamelCase=False,      lowerCamelCase=3_00,      lowerCamelCase=False,      lowerCamelCase=1,      lowerCamelCase=5,      lowerCamelCase=2,      lowerCamelCase=1,      lowerCamelCase=1,      lowerCamelCase=5,      lowerCamelCase=2,      lowerCamelCase=0.1,      lowerCamelCase=0.2_5,      lowerCamelCase=False,      **lowerCamelCase,      )       ->   Optional[int]:




           """simple docstring"""

           if backbone_config is not None and use_timm_backbone:
               raise ValueError('You can\'t specify both `backbone_config` and `use_timm_backbone`.')

           if not use_timm_backbone:
               if backbone_config is None:
                   logger.info('`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.')
                   _lowercase       :							List[str]      		=		CONFIG_MAPPING['resnet'](out_features=['stage4'])
               elif isinstance(lowerCamelCase,      lowerCamelCase):
                   _lowercase       :							List[str]      		=		backbone_config.get('model_type')
                   _lowercase       :							str      		=		CONFIG_MAPPING[backbone_model_type]
                   _lowercase       :							Optional[int]      		=		config_class.from_dict(lowerCamelCase)
           _lowercase       :							Tuple      		=		use_timm_backbone
           _lowercase       :							List[str]      		=		backbone_config
           _lowercase       :							Tuple      		=		num_channels
           _lowercase       :							Optional[Any]      		=		num_queries
           _lowercase       :							Optional[Any]      		=		max_position_embeddings
           _lowercase       :							Optional[int]      		=		d_model
           _lowercase       :							int      		=		encoder_ffn_dim
           _lowercase       :							List[Any]      		=		encoder_layers
           _lowercase       :							str      		=		encoder_attention_heads
           _lowercase       :							str      		=		decoder_ffn_dim
           _lowercase       :							Optional[Any]      		=		decoder_layers
           _lowercase       :							List[str]      		=		decoder_attention_heads
           _lowercase       :							Optional[int]      		=		dropout
           _lowercase       :							Optional[Any]      		=		attention_dropout
           _lowercase       :							int      		=		activation_dropout
           _lowercase       :							Any      		=		activation_function
           _lowercase       :							Optional[int]      		=		init_std
           _lowercase       :							int      		=		init_xavier_std
           _lowercase       :							Union[str, Any]      		=		encoder_layerdrop
           _lowercase       :							Tuple      		=		auxiliary_loss
           _lowercase       :							Union[str, Any]      		=		position_embedding_type
           _lowercase       :							str      		=		backbone
           _lowercase       :							List[Any]      		=		use_pretrained_backbone
           _lowercase       :							Any      		=		dilation
           # deformable attributes
           _lowercase       :							Any      		=		num_feature_levels
           _lowercase       :							Dict      		=		encoder_n_points
           _lowercase       :							Dict      		=		decoder_n_points
           _lowercase       :							Dict      		=		two_stage
           _lowercase       :							Union[str, Any]      		=		two_stage_num_proposals
           _lowercase       :							str      		=		with_box_refine
           if two_stage is True and with_box_refine is False:
               raise ValueError('If two_stage is True, with_box_refine must be True.')
           # Hungarian matcher
           _lowercase       :							Tuple      		=		class_cost
           _lowercase       :							int      		=		bbox_cost
           _lowercase       :							Optional[int]      		=		giou_cost
           # Loss coefficients
           _lowercase       :							Optional[Any]      		=		mask_loss_coefficient
           _lowercase       :							Dict      		=		dice_loss_coefficient
           _lowercase       :							Tuple      		=		bbox_loss_coefficient
           _lowercase       :							Optional[int]      		=		giou_loss_coefficient
           _lowercase       :							Union[str, Any]      		=		eos_coefficient
           _lowercase       :							Union[str, Any]      		=		focal_alpha
           _lowercase       :							Dict      		=		disable_custom_kernels
           super().__init__(is_encoder_decoder=lowerCamelCase,      **lowerCamelCase)
       @property
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           return self.encoder_attention_heads
       @property
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           return self.d_model




       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		copy.deepcopy(self.__dict__)
           if self.backbone_config is not None:
               _lowercase       :							Union[str, Any]      		=		self.backbone_config.to_dict()
           _lowercase       :							Tuple      		=		self.__class__.model_type
           return output


