






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







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

						_A	:				Tuple   =      rescale_embeddings
						_A	:				Tuple   =      attention_type
						_A	:				Tuple   =      use_bias
						_A	:				Any   =      block_size
						_A	:				Optional[int]   =      num_random_blocks







		def  A__      (					self		):
						_A	:				List[str]   =      ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size		)

						_A	:				Optional[Any]   =      None
						if self.use_attention_mask:
										_A	:				int   =      random_attention_mask([self.batch_size, self.seq_length]		)

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

						_A	:				Tuple   =      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=A__       ,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  A__      (					self		):
						_A	:				Union[str, Any]   =      self.prepare_config_and_inputs()
						_A  ,				_A  ,				_A  ,				_A	:				List[Any]   =      config_and_inputs
						_A	:				Dict   =      {
						    '''input_ids''': input_ids,
						    '''token_type_ids''': token_type_ids,
						    '''attention_mask''': attention_mask,
						}
						return config, inputs_dict






@require_flax
class 						UpperCAmelCase__ (							__snake_case     ,				unittest.TestCase							):
		__snake_case      :       Any     =   (
		    (
		        FlaxBigBirdForCausalLM,
		        FlaxBigBirdModel,
		        FlaxBigBirdForPreTraining,
		        FlaxBigBirdForMaskedLM,
		        FlaxBigBirdForMultipleChoice,
		        FlaxBigBirdForQuestionAnswering,
		        FlaxBigBirdForSequenceClassification,
		        FlaxBigBirdForTokenClassification,
		    )
		    if is_flax_available()
		    else ()
		)

		__snake_case      :       str     =   False
		__snake_case      :       Optional[Any]     =   False







		def  A__      (					self		):
						_A	:				Dict   =      FlaxBigBirdModelTester(self		)







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







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







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







		@slow
		def  A__      (					self		):
						for model_class_name in self.all_model_classes:
										_A	:				str   =      model_class_name.from_pretrained('''google/bigbird-roberta-base'''		)
										self.assertIsNotNone(A__		)







		def  A__      (					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  A__      (					self		):
						_A  ,				_A	:				Any   =      self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__		):
														_A	:				Dict   =      self._prepare_for_class(A__       ,A__		)
														_A	:				Any   =      model_class(A__		)

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

														with self.subTest('''JIT Enabled'''		):
																		_A	:				Dict   =      model_jitted(**A__		).to_tuple()

														with self.subTest('''JIT Disabled'''		):
																		with jax.disable_jit():
																						_A	:				List[str]   =      model_jitted(**A__		).to_tuple()

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







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











from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType






class 						UpperCAmelCase__ (							__snake_case							):
		__snake_case      :       Optional[Any]     =   ["image_processor", "tokenizer"]
		__snake_case      :       Tuple     =   "BridgeTowerImageProcessor"
		__snake_case      :       List[str]     =   ("RobertaTokenizer", "RobertaTokenizerFast")







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







		def __call__(					self       ,A__       ,A__ = None       ,A__ = True       ,A__ = False       ,A__ = None       ,A__ = None       ,A__ = 0       ,A__ = None       ,A__ = None       ,A__ = None       ,A__ = False       ,A__ = False       ,A__ = False       ,A__ = False       ,A__ = True       ,A__ = None       ,**A__       ,):
						_A	:				List[Any]   =      self.tokenizer(
						    text=A__       ,add_special_tokens=A__       ,padding=A__       ,truncation=A__       ,max_length=A__       ,stride=A__       ,pad_to_multiple_of=A__       ,return_token_type_ids=A__       ,return_attention_mask=A__       ,return_overflowing_tokens=A__       ,return_special_tokens_mask=A__       ,return_offsets_mapping=A__       ,return_length=A__       ,verbose=A__       ,return_tensors=A__       ,**A__       ,)
						# add pixel_values + pixel_mask
						_A	:				Optional[Any]   =      self.image_processor(
						    A__       ,return_tensors=A__       ,do_normalize=A__       ,do_center_crop=A__       ,**A__		)
						encoding.update(A__		)

						return encoding







		def  A__      (					self       ,*A__       ,**A__		):
						return self.tokenizer.batch_decode(*A__       ,**A__		)







		def  A__      (					self       ,*A__       ,**A__		):
						return self.tokenizer.decode(*A__       ,**A__		)







		@property
		def  A__      (					self		):
						_A	:				Union[str, Any]   =      self.tokenizer.model_input_names
						_A	:				str   =      self.image_processor.model_input_names
						return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names		)		)



