


'''simple docstring'''
import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
							import jax

							from transformers.models.big_bird.modeling_flax_big_bird import (
							    FlaxBigBirdForCausalLM,
							    FlaxBigBirdForMaskedLM,
							    FlaxBigBirdForMultipleChoice,
							    FlaxBigBirdForPreTraining,
							    FlaxBigBirdForQuestionAnswering,
							    FlaxBigBirdForSequenceClassification,
							    FlaxBigBirdForTokenClassification,
							    FlaxBigBirdModel,
							)






class 					__lowerCAmelCase   (   unittest.TestCase						):

						"""simple docstring"""






						def __init__(  self							:      int ,  lowerCAmelCase							:      Union[str, Any] ,  lowerCAmelCase							:      Dict=2 ,  lowerCAmelCase							:      int=56 ,  lowerCAmelCase							:      List[str]=True ,  lowerCAmelCase							:      Tuple=True ,  lowerCAmelCase							:      Dict=True ,  lowerCAmelCase							:      Optional[Any]=True ,  lowerCAmelCase							:      Any=99 ,  lowerCAmelCase							:      List[Any]=32 ,  lowerCAmelCase							:      Optional[Any]=2 ,  lowerCAmelCase							:      Optional[Any]=2 ,  lowerCAmelCase							:      Any=7 ,  lowerCAmelCase							:      List[Any]="gelu_new" ,  lowerCAmelCase							:      Tuple=0.1 ,  lowerCAmelCase							:      int=0.1 ,  lowerCAmelCase							:      str=5_12 ,  lowerCAmelCase							:      Optional[Any]=16 ,  lowerCAmelCase							:      Dict=2 ,  lowerCAmelCase							:      Optional[Any]=0.0_2 ,  lowerCAmelCase							:      str=4 ,  lowerCAmelCase							:      Union[str, Any]="block_sparse" ,  lowerCAmelCase							:      Any=True ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      List[str]=2 ,  lowerCAmelCase							:      str=3 ,  ):
													A_			    =		parent
													A_			    =		batch_size
													A_			    =		seq_length
													A_			    =		is_training
													A_			    =		use_attention_mask
													A_			    =		use_token_type_ids
													A_			    =		use_labels
													A_			    =		vocab_size
													A_			    =		hidden_size
													A_			    =		num_hidden_layers
													A_			    =		num_attention_heads
													A_			    =		intermediate_size
													A_			    =		hidden_act
													A_			    =		hidden_dropout_prob
													A_			    =		attention_probs_dropout_prob
													A_			    =		max_position_embeddings
													A_			    =		type_vocab_size
													A_			    =		type_sequence_label_size
													A_			    =		initializer_range
													A_			    =		num_choices

													A_			    =		rescale_embeddings
													A_			    =		attention_type
													A_			    =		use_bias
													A_			    =		block_size
													A_			    =		num_random_blocks






						def 				_UpperCAmelCase       (  self							:      Union[str, Any]	):
													A_			    =		ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size	)

													A_			    =		None
													if self.use_attention_mask:
																				A_			    =		random_attention_mask([self.batch_size, self.seq_length]	)

													A_			    =		None
													if self.use_token_type_ids:
																				A_			    =		ids_tensor([self.batch_size, self.seq_length] ,  self.type_vocab_size	)

													A_			    =		BigBirdConfig(
													    vocab_size=self.vocab_size ,  hidden_size=self.hidden_size ,  num_hidden_layers=self.num_hidden_layers ,  num_attention_heads=self.num_attention_heads ,  intermediate_size=self.intermediate_size ,  hidden_act=self.hidden_act ,  hidden_dropout_prob=self.hidden_dropout_prob ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob ,  max_position_embeddings=self.max_position_embeddings ,  type_vocab_size=self.type_vocab_size ,  is_decoder=lowerCAmelCase ,  initializer_range=self.initializer_range ,  attention_type=self.attention_type ,  block_size=self.block_size ,  num_random_blocks=self.num_random_blocks ,  use_bias=self.use_bias ,  rescale_embeddings=self.rescale_embeddings ,  )

													return config, input_ids, token_type_ids, attention_mask






						def 				_UpperCAmelCase       (  self							:      int	):
													A_			    =		self.prepare_config_and_inputs()
													A_    ,					A_    ,					A_    ,					A_			    =		config_and_inputs
													A_			    =		{
													    "input_ids": input_ids,
													    "token_type_ids": token_type_ids,
													    "attention_mask": attention_mask,
													}
													return config, inputs_dict






@require_flax
class 					__lowerCAmelCase   (   lowercase				,			unittest.TestCase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Dict   						=(
						    (
						        FlaxBigBirdForCausalLM,
						        FlaxBigBirdModel,
						        FlaxBigBirdForPreTraining,
						        FlaxBigBirdForMaskedLM,
						        FlaxBigBirdForMultipleChoice,
						        FlaxBigBirdForQuestionAnswering,
						        FlaxBigBirdForSequenceClassification,
						        FlaxBigBirdForTokenClassification,
						    )
						    if is_flax_available()
						    else ()
						)

						_UpperCAmelCase		:       Any   						=False
						_UpperCAmelCase		:       Dict   						=False






						def 				_UpperCAmelCase       (  self							:      Any	):
													A_			    =		FlaxBigBirdModelTester(self	)






						@slow
						# copied from `test_modeling_flax_common` because it takes much longer than other models
						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													super().test_from_pretrained_save_pretrained()






						@slow
						# copied from `test_modeling_flax_common` because it takes much longer than other models
						def 				_UpperCAmelCase       (  self							:      Optional[int]	):
													super().test_from_pretrained_with_no_automatic_init()






						@slow
						# copied from `test_modeling_flax_common` because it takes much longer than other models
						def 				_UpperCAmelCase       (  self							:      Dict	):
													super().test_no_automatic_init()






						@slow
						# copied from `test_modeling_flax_common` because it takes much longer than other models
						def 				_UpperCAmelCase       (  self							:      int	):
													super().test_hidden_states_output()






						@slow
						def 				_UpperCAmelCase       (  self							:      Union[str, Any]	):
													for model_class_name in self.all_model_classes:
																				A_			    =		model_class_name.from_pretrained("google/bigbird-roberta-base"	)
																				self.assertIsNotNone(lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Dict	):
													if self.test_attn_probs:
																				super().test_attention_outputs()






						@slow
						# copied from `test_modeling_flax_common` because it takes much longer than other models
						def 				_UpperCAmelCase       (  self							:      List[str]	):
													A_    ,					A_			    =		self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				with self.subTest(model_class.__name__	):
																											A_			    =		self._prepare_for_class(lowerCAmelCase ,  lowerCAmelCase	)
																											A_			    =		model_class(lowerCAmelCase	)

																											@jax.jit
																											def model_jitted(lowerCAmelCase							:      str ,  lowerCAmelCase							:      Tuple=None ,  **lowerCAmelCase							:      Tuple	):
																																		return model(input_ids=lowerCAmelCase ,  attention_mask=lowerCAmelCase ,  **lowerCAmelCase	)

																											with self.subTest("JIT Enabled"	):
																																		A_			    =		model_jitted(**lowerCAmelCase	).to_tuple()

																											with self.subTest("JIT Disabled"	):
																																		with jax.disable_jit():
																																									A_			    =		model_jitted(**lowerCAmelCase	).to_tuple()

																											self.assertEqual(len(lowerCAmelCase	) ,  len(lowerCAmelCase	)	)
																											for jitted_output, output in zip(lowerCAmelCase ,  lowerCAmelCase	):
																																		self.assertEqual(jitted_output.shape ,  output.shape	)






						def 				_UpperCAmelCase       (  self							:      Any ,  lowerCAmelCase							:      int ,  lowerCAmelCase							:      Optional[Any] ,  lowerCAmelCase							:      Tuple ,  lowerCAmelCase							:      Union[str, Any]=1e-5 ,  lowerCAmelCase							:      List[Any]="outputs" ,  lowerCAmelCase							:      Tuple=None	):
													# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
													# an effort was done to return `attention_probs` (yet to be verified).
													if name.startswith("outputs.attentions"	):
																				return
													else:
																				super().check_pt_flax_outputs(lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase ,  lowerCAmelCase	)









'''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	)





