


'''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       a_				(	UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		[]
							for part_id in partition_order:
														A_			    =		df.where(f"SPARK_PARTITION_ID() = {part_id}" ).collect()
														for row_idx, row in enumerate(UpperCamelCase_ ):
																					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       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(1_0_0 ).repartition(1 )
							A_			    =		Spark(UpperCamelCase_ )
							# 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=1_6 )
							# Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
							assert spark_builder.df.rdd.getNumPartitions() == 5_0







@require_not_windows
@require_dill_gt_0_3_2
def       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(1_0 ).repartition(2 )
							A_			    =		[1, 0]
							A_			    =		_generate_iterable_examples(UpperCamelCase_						,  UpperCamelCase_ )  # Reverse the partitions.
							A_			    =		_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase_						,  UpperCamelCase_ )

							for i, (row_id, row_dict) in enumerate(generate_fn() ):
														A_    ,					A_			    =		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       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(1_0 ).repartition(1 )
							A_			    =		SparkExamplesIterable(UpperCamelCase_ )
							assert it.n_shards == 1
							for i, (row_id, row_dict) in enumerate(UpperCamelCase_ ):
														assert row_id == f"0_{i}"
														assert row_dict == {"id": i}







@require_not_windows
@require_dill_gt_0_3_2
def       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(3_0 ).repartition(3 )
							# Mock the generator so that shuffle reverses the partition indices.
							with patch("numpy.random.Generator" ) as generator_mock:
														A_			    =		lambda UpperCamelCase_ : x.reverse()
														A_			    =		_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase_						,  [2, 1, 0] )

														A_			    =		SparkExamplesIterable(UpperCamelCase_ ).shuffle_data_sources(UpperCamelCase_ )
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(UpperCamelCase_ ):
																					A_    ,					A_			    =		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       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(2_0 ).repartition(4 )

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

							# Partitions 1 and 3
							A_			    =		SparkExamplesIterable(UpperCamelCase_ ).shard_data_sources(worker_id=1						,  num_workers=2 )
							assert shard_it_a.n_shards == 2
							A_			    =		_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase_						,  [1, 3] )
							for i, (row_id, row_dict) in enumerate(UpperCamelCase_ ):
														A_    ,					A_			    =		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       a_				(	):
							A_			    =		pyspark.sql.SparkSession.builder.master("local[*]" ).appName("pyspark" ).getOrCreate()
							A_			    =		spark.range(1_0_0 ).repartition(1 )
							A_			    =		Spark(UpperCamelCase_ )
							# 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_0_0









'''simple docstring'''
import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE					:				Dict									=		logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					:				int									=		{
    '''facebook/wav2vec2-base-960h''': '''https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json''',
    # See all Wav2Vec2 models at https://huggingface.co/models?filter=wav2vec2
}






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Dict   						="wav2vec2"






						def __init__(  self							:      str ,  lowerCAmelCase							:      Dict=32 ,  lowerCAmelCase							:      List[Any]=7_68 ,  lowerCAmelCase							:      List[str]=12 ,  lowerCAmelCase							:      List[str]=12 ,  lowerCAmelCase							:      List[Any]=30_72 ,  lowerCAmelCase							:      Optional[Any]="gelu" ,  lowerCAmelCase							:      Optional[int]=0.1 ,  lowerCAmelCase							:      Optional[int]=0.1 ,  lowerCAmelCase							:      List[Any]=0.1 ,  lowerCAmelCase							:      Dict=0.0 ,  lowerCAmelCase							:      Union[str, Any]=0.0 ,  lowerCAmelCase							:      str=0.1 ,  lowerCAmelCase							:      List[str]=0.1 ,  lowerCAmelCase							:      Dict=0.0_2 ,  lowerCAmelCase							:      List[str]=1e-5 ,  lowerCAmelCase							:      Optional[Any]="group" ,  lowerCAmelCase							:      List[Any]="gelu" ,  lowerCAmelCase							:      int=(5_12, 5_12, 5_12, 5_12, 5_12, 5_12, 5_12) ,  lowerCAmelCase							:      Union[str, Any]=(5, 2, 2, 2, 2, 2, 2) ,  lowerCAmelCase							:      Any=(10, 3, 3, 3, 3, 2, 2) ,  lowerCAmelCase							:      str=False ,  lowerCAmelCase							:      Tuple=1_28 ,  lowerCAmelCase							:      Union[str, Any]=16 ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      Union[str, Any]=True ,  lowerCAmelCase							:      Any=0.0_5 ,  lowerCAmelCase							:      Any=10 ,  lowerCAmelCase							:      Dict=2 ,  lowerCAmelCase							:      List[Any]=0.0 ,  lowerCAmelCase							:      int=10 ,  lowerCAmelCase							:      List[str]=0 ,  lowerCAmelCase							:      Tuple=3_20 ,  lowerCAmelCase							:      List[str]=2 ,  lowerCAmelCase							:      int=0.1 ,  lowerCAmelCase							:      List[Any]=1_00 ,  lowerCAmelCase							:      List[str]=2_56 ,  lowerCAmelCase							:      int=2_56 ,  lowerCAmelCase							:      Optional[Any]=0.1 ,  lowerCAmelCase							:      str="sum" ,  lowerCAmelCase							:      List[Any]=False ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      Tuple=2_56 ,  lowerCAmelCase							:      str=(5_12, 5_12, 5_12, 5_12, 15_00) ,  lowerCAmelCase							:      List[str]=(5, 3, 3, 1, 1) ,  lowerCAmelCase							:      Tuple=(1, 2, 3, 1, 1) ,  lowerCAmelCase							:      Optional[Any]=5_12 ,  lowerCAmelCase							:      Any=0 ,  lowerCAmelCase							:      str=1 ,  lowerCAmelCase							:      Any=2 ,  lowerCAmelCase							:      Any=False ,  lowerCAmelCase							:      List[Any]=3 ,  lowerCAmelCase							:      Any=2 ,  lowerCAmelCase							:      List[str]=3 ,  lowerCAmelCase							:      int=None ,  lowerCAmelCase							:      Optional[int]=None ,  **lowerCAmelCase							:      Union[str, Any] ,  ):
													super().__init__(**lowerCAmelCase ,  pad_token_id=lowerCAmelCase ,  bos_token_id=lowerCAmelCase ,  eos_token_id=lowerCAmelCase	)
													A_			    =		hidden_size
													A_			    =		feat_extract_norm
													A_			    =		feat_extract_activation
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		conv_bias
													A_			    =		num_conv_pos_embeddings
													A_			    =		num_conv_pos_embedding_groups
													A_			    =		len(self.conv_dim	)
													A_			    =		num_hidden_layers
													A_			    =		intermediate_size
													A_			    =		hidden_act
													A_			    =		num_attention_heads
													A_			    =		hidden_dropout
													A_			    =		attention_dropout
													A_			    =		activation_dropout
													A_			    =		feat_proj_dropout
													A_			    =		final_dropout
													A_			    =		layerdrop
													A_			    =		layer_norm_eps
													A_			    =		initializer_range
													A_			    =		vocab_size
													A_			    =		do_stable_layer_norm
													A_			    =		use_weighted_layer_sum

													if (
													    (len(self.conv_stride	) != self.num_feat_extract_layers)
													    or (len(self.conv_kernel	) != self.num_feat_extract_layers)
													    or (len(self.conv_dim	) != self.num_feat_extract_layers)
													):
																				raise ValueError(
																				    "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
																				    " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
																				    F" {len(self.conv_dim	)}`, `len(config.conv_stride) = {len(self.conv_stride	)}`,"
																				    F" `len(config.conv_kernel) = {len(self.conv_kernel	)}`."	)

													# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
													A_			    =		apply_spec_augment
													A_			    =		mask_time_prob
													A_			    =		mask_time_length
													A_			    =		mask_time_min_masks
													A_			    =		mask_feature_prob
													A_			    =		mask_feature_length
													A_			    =		mask_feature_min_masks

													# parameters for pretraining with codevector quantized representations
													A_			    =		num_codevectors_per_group
													A_			    =		num_codevector_groups
													A_			    =		contrastive_logits_temperature
													A_			    =		feat_quantizer_dropout
													A_			    =		num_negatives
													A_			    =		codevector_dim
													A_			    =		proj_codevector_dim
													A_			    =		diversity_loss_weight

													# ctc loss
													A_			    =		ctc_loss_reduction
													A_			    =		ctc_zero_infinity

													# adapter
													A_			    =		add_adapter
													A_			    =		adapter_kernel_size
													A_			    =		adapter_stride
													A_			    =		num_adapter_layers
													A_			    =		output_hidden_size or hidden_size
													A_			    =		adapter_attn_dim

													# SequenceClassification-specific parameter. Feel free to ignore for other classes.
													A_			    =		classifier_proj_size

													# XVector-specific parameters. Feel free to ignore for other classes.
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		xvector_output_dim






						@property
						def 				_UpperCAmelCase       (  self							:      Tuple	):
													return functools.reduce(operator.mul ,  self.conv_stride ,  1	)





