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






 def __init__(	self     :Optional[int]		, _A     :Any		, _A     :int=2		, _A     :Optional[Any]=56		, _A     :Tuple=True		, _A     :Union[str, Any]=True		, _A     :Optional[Any]=True		, _A     :str=True		, _A     :str=99		, _A     :Tuple=32		, _A     :int=2		, _A     :Optional[int]=2		, _A     :Tuple=7		, _A     :Optional[int]="gelu_new"		, _A     :Any=0.1		, _A     :List[str]=0.1		, _A     :List[Any]=512		, _A     :Union[str, Any]=16		, _A     :int=2		, _A     :Dict=0.02		, _A     :Tuple=4		, _A     :Union[str, Any]="block_sparse"		, _A     :Optional[Any]=True		, _A     :Any=False		, _A     :Any=2		, _A     :int=3		, ) ->	Optional[int]:







       '''simple docstring'''



       __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 		lowercase_  (	self     :int	) ->	Dict:







       '''simple docstring'''



       __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=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 		lowercase_  (	self     :Dict	) ->	Dict:







       '''simple docstring'''



       __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        UpperCamelCase__     (			__lowerCamelCase					,   unittest.TestCase):
 UpperCAmelCase__			:			Optional[int]   =    (
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 UpperCAmelCase__			:			Optional[int]   =    False
 UpperCAmelCase__			:			Optional[int]   =    False






 def 		lowercase_  (	self     :List[str]	) ->	Optional[Any]:







       '''simple docstring'''



       __A						=	FlaxBigBirdModelTester(self	)






 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		lowercase_  (	self     :Optional[Any]	) ->	List[Any]:







       '''simple docstring'''



       super().test_from_pretrained_save_pretrained()






 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		lowercase_  (	self     :List[Any]	) ->	List[str]:







       '''simple docstring'''



       super().test_from_pretrained_with_no_automatic_init()






 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		lowercase_  (	self     :List[Any]	) ->	int:







       '''simple docstring'''



       super().test_no_automatic_init()






 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		lowercase_  (	self     :List[str]	) ->	Optional[Any]:







       '''simple docstring'''



       super().test_hidden_states_output()






 @slow
 def 		lowercase_  (	self     :Optional[Any]	) ->	Dict:







       '''simple docstring'''



       for model_class_name in self.all_model_classes:
             __A						=	model_class_name.from_pretrained('google/bigbird-roberta-base'	)
             self.assertIsNotNone(UpperCamelCase_	)






 def 		lowercase_  (	self     :Dict	) ->	Optional[int]:







       '''simple docstring'''



       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 		lowercase_  (	self     :List[Any]	) ->	List[Any]:







       '''simple docstring'''



       __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(UpperCamelCase_		, UpperCamelCase_	)
                   __A						=	model_class(UpperCamelCase_	)

                   @jax.jit
                   def model_jitted(_A     :Tuple		, _A     :List[Any]=None		, **_A     :Union[str, Any]	):
                         return model(input_ids=UpperCamelCase_		, attention_mask=UpperCamelCase_		, **UpperCamelCase_	)

                   with self.subTest('JIT Enabled'	):
                         __A						=	model_jitted(**UpperCamelCase_	).to_tuple()

                   with self.subTest('JIT Disabled'	):
                         with jax.disable_jit():
                               __A						=	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 		lowercase_  (	self     :List[Any]		, _A     :int		, _A     :List[str]		, _A     :Any		, _A     :Dict=1E-5		, _A     :List[str]="outputs"		, _A     :List[str]=None	) ->	str:







       '''simple docstring'''



       if name.startswith('outputs.attentions'	):
             return
       else:
             super().check_pt_flax_outputs(UpperCamelCase_		, UpperCamelCase_		, UpperCamelCase_		, UpperCamelCase_		, UpperCamelCase_		, UpperCamelCase_	)












from transformers import DistilBertTokenizer, DistilBertTokenizerFast
from transformers.testing_utils import require_tokenizers, slow

from ..bert.test_tokenization_bert import BertTokenizationTest





@require_tokenizers
class     lowerCamelCase__(					__lowerCamelCase):
      UpperCAmelCase__      :	Dict    		=   DistilBertTokenizer
      UpperCAmelCase__      :	Dict    		=   DistilBertTokenizerFast
      UpperCAmelCase__      :	Tuple    		=   True
      @slow
      def 			lowerCAmelCase__						(				self: Tuple   ):
             __lowerCamelCase   =						DistilBertTokenizer.from_pretrained("""distilbert-base-uncased"""   )

             __lowerCamelCase   =						tokenizer.encode("""sequence builders"""    ,       add_special_tokens=UpperCamelCase_   )
             __lowerCamelCase   =						tokenizer.encode("""multi-sequence build"""    ,       add_special_tokens=UpperCamelCase_   )

             __lowerCamelCase   =						tokenizer.build_inputs_with_special_tokens(UpperCamelCase_   )
             __lowerCamelCase   =						tokenizer.build_inputs_with_special_tokens(UpperCamelCase_    ,       UpperCamelCase_   )

             assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
             assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + text_a + [
                 tokenizer.sep_token_id
             ]

