
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    a_ (			unittest.TestCase    ):






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase=2				,			__UpperCamelCase=56				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=99				,			__UpperCamelCase=32				,			__UpperCamelCase=2				,			__UpperCamelCase=2				,			__UpperCamelCase=7				,			__UpperCamelCase="gelu_new"				,			__UpperCamelCase=0.1				,			__UpperCamelCase=0.1				,			__UpperCamelCase=512				,			__UpperCamelCase=16				,			__UpperCamelCase=2				,			__UpperCamelCase=0.02				,			__UpperCamelCase=4				,			__UpperCamelCase="block_sparse"				,			__UpperCamelCase=True				,			__UpperCamelCase=False				,			__UpperCamelCase=2				,			__UpperCamelCase=3				,			):
									_lowercase					       =					parent
									_lowercase					       =					batch_size
									_lowercase					       =					seq_length
									_lowercase					       =					is_training
									_lowercase					       =					use_attention_mask
									_lowercase					       =					use_token_type_ids
									_lowercase					       =					use_labels
									_lowercase					       =					vocab_size
									_lowercase					       =					hidden_size
									_lowercase					       =					num_hidden_layers
									_lowercase					       =					num_attention_heads
									_lowercase					       =					intermediate_size
									_lowercase					       =					hidden_act
									_lowercase					       =					hidden_dropout_prob
									_lowercase					       =					attention_probs_dropout_prob
									_lowercase					       =					max_position_embeddings
									_lowercase					       =					type_vocab_size
									_lowercase					       =					type_sequence_label_size
									_lowercase					       =					initializer_range
									_lowercase					       =					num_choices

									_lowercase					       =					rescale_embeddings
									_lowercase					       =					attention_type
									_lowercase					       =					use_bias
									_lowercase					       =					block_size
									_lowercase					       =					num_random_blocks






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size  )

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

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

									_lowercase					       =					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=__UpperCamelCase				,			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  ):
									_lowercase					       =					self.prepare_config_and_inputs()
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					config_and_inputs
									_lowercase					       =					{
									    """input_ids""": input_ids,
									    """token_type_ids""": token_type_ids,
									    """attention_mask""": attention_mask,
									}
									return config, inputs_dict







@require_flax
class    a_ (			_a					,			unittest.TestCase    ):
			a	:			Union[str, Any]					   =							(
			    (
			        FlaxBigBirdForCausalLM,
			        FlaxBigBirdModel,
			        FlaxBigBirdForPreTraining,
			        FlaxBigBirdForMaskedLM,
			        FlaxBigBirdForMultipleChoice,
			        FlaxBigBirdForQuestionAnswering,
			        FlaxBigBirdForSequenceClassification,
			        FlaxBigBirdForTokenClassification,
			    )
			    if is_flax_available()
			    else ()
			)

			a	:			Optional[int]					   =							False
			a	:			List[str]					   =							False






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					FlaxBigBirdModelTester(self  )






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






			@slow
			# copied from `test_modeling_flax_common` because it takes much longer than other models
			def    UpperCamelCase_				(      self  ):
									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  ):
									super().test_no_automatic_init()






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






			@slow
			def    UpperCamelCase_				(      self  ):
									for model_class_name in self.all_model_classes:
															_lowercase					       =					model_class_name.from_pretrained("""google/bigbird-roberta-base"""  )
															self.assertIsNotNone(__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									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  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															with self.subTest(model_class.__name__  ):
																					_lowercase					       =					self._prepare_for_class(__UpperCamelCase				,			__UpperCamelCase  )
																					_lowercase					       =					model_class(__UpperCamelCase  )

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

																					with self.subTest("""JIT Enabled"""  ):
																											_lowercase					       =					model_jitted(**__UpperCamelCase  ).to_tuple()

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

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






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase=1E-5				,			__UpperCamelCase="outputs"				,			__UpperCamelCase=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(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )

from ...processing_utils import ProcessorMixin



class    a_ (			_a    ):
			a	:			Optional[int]					   =							'''SpeechT5FeatureExtractor'''
			a	:			List[Any]					   =							'''SpeechT5Tokenizer'''






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase  ):
									super().__init__(__UpperCamelCase				,			__UpperCamelCase  )






			def __call__(      self				,			*__UpperCamelCase				,			**__UpperCamelCase  ):
									_lowercase					       =					kwargs.pop("""audio"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""text"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""text_target"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""audio_target"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""sampling_rate"""				,			__UpperCamelCase  )

									if audio is not None and text is not None:
															raise ValueError(
															    """Cannot process both `audio` and `text` inputs. Did you mean `audio_target` or `text_target`?"""  )
									if audio_target is not None and text_target is not None:
															raise ValueError(
															    """Cannot process both `audio_target` and `text_target` inputs. Did you mean `audio` or `text`?"""  )
									if audio is None and audio_target is None and text is None and text_target is None:
															raise ValueError(
															    """You need to specify either an `audio`, `audio_target`, `text`, or `text_target` input to process."""  )

									if audio is not None:
															_lowercase					       =					self.feature_extractor(__UpperCamelCase				,			*__UpperCamelCase				,			sampling_rate=__UpperCamelCase				,			**__UpperCamelCase  )
									elif text is not None:
															_lowercase					       =					self.tokenizer(__UpperCamelCase				,			**__UpperCamelCase  )
									else:
															_lowercase					       =					None

									if audio_target is not None:
															_lowercase					       =					self.feature_extractor(audio_target=__UpperCamelCase				,			*__UpperCamelCase				,			sampling_rate=__UpperCamelCase				,			**__UpperCamelCase  )
															_lowercase					       =					targets["""input_values"""]
									elif text_target is not None:
															_lowercase					       =					self.tokenizer(__UpperCamelCase				,			**__UpperCamelCase  )
															_lowercase					       =					targets["""input_ids"""]
									else:
															_lowercase					       =					None

									if inputs is None:
															return targets

									if targets is not None:
															_lowercase					       =					labels

															_lowercase					       =					targets.get("""attention_mask"""  )
															if decoder_attention_mask is not None:
																					_lowercase					       =					decoder_attention_mask

									return inputs






			def    UpperCamelCase_				(      self				,			*__UpperCamelCase				,			**__UpperCamelCase  ):
									_lowercase					       =					kwargs.pop("""input_values"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""input_ids"""				,			__UpperCamelCase  )
									_lowercase					       =					kwargs.pop("""labels"""				,			__UpperCamelCase  )

									if input_values is not None and input_ids is not None:
															raise ValueError("""Cannot process both `input_values` and `input_ids` inputs."""  )
									if input_values is None and input_ids is None and labels is None:
															raise ValueError(
															    """You need to specify either an `input_values`, `input_ids`, or `labels` input to be padded."""  )

									if input_values is not None:
															_lowercase					       =					self.feature_extractor.pad(__UpperCamelCase				,			*__UpperCamelCase				,			**__UpperCamelCase  )
									elif input_ids is not None:
															_lowercase					       =					self.tokenizer.pad(__UpperCamelCase				,			**__UpperCamelCase  )
									else:
															_lowercase					       =					None

									if labels is not None:
															if "input_ids" in labels or (isinstance(__UpperCamelCase				,			__UpperCamelCase  ) and "input_ids" in labels[0]):
																					_lowercase					       =					self.tokenizer.pad(__UpperCamelCase				,			**__UpperCamelCase  )
																					_lowercase					       =					targets["""input_ids"""]
															else:
																					_lowercase					       =					self.feature_extractor.feature_size
																					_lowercase					       =					self.feature_extractor.num_mel_bins
																					_lowercase					       =					self.feature_extractor.pad(__UpperCamelCase				,			*__UpperCamelCase				,			**__UpperCamelCase  )
																					_lowercase					       =					feature_size_hack
																					_lowercase					       =					targets["""input_values"""]
									else:
															_lowercase					       =					None

									if inputs is None:
															return targets

									if targets is not None:
															_lowercase					       =					labels

															_lowercase					       =					targets.get("""attention_mask"""  )
															if decoder_attention_mask is not None:
																					_lowercase					       =					decoder_attention_mask

									return inputs






			def    UpperCamelCase_				(      self				,			*__UpperCamelCase				,			**__UpperCamelCase  ):
									return self.tokenizer.batch_decode(*__UpperCamelCase				,			**__UpperCamelCase  )






			def    UpperCamelCase_				(      self				,			*__UpperCamelCase				,			**__UpperCamelCase  ):
									return self.tokenizer.decode(*__UpperCamelCase				,			**__UpperCamelCase  )