





from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TFFunnelBaseModel,
           TFFunnelForMaskedLM,
           TFFunnelForMultipleChoice,
           TFFunnelForPreTraining,
           TFFunnelForQuestionAnswering,
           TFFunnelForSequenceClassification,
           TFFunnelForTokenClassification,
           TFFunnelModel,
       )




class     SCREAMING_SNAKE_CASE__ :



 def __init__(							self,__lowerCamelCase,__lowerCamelCase=13,__lowerCamelCase=7,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=99,__lowerCamelCase=[1, 1, 2],__lowerCamelCase=1,__lowerCamelCase=32,__lowerCamelCase=4,__lowerCamelCase=8,__lowerCamelCase=37,__lowerCamelCase="gelu_new",__lowerCamelCase=0.1,__lowerCamelCase=0.1,__lowerCamelCase=0.0,__lowerCamelCase=512,__lowerCamelCase=3,__lowerCamelCase=0.02,__lowerCamelCase=3,__lowerCamelCase=4,__lowerCamelCase=None,__lowerCamelCase=False,):
        A__											=						parent
        A__											=						batch_size
        A__											=						seq_length
        A__											=						is_training
        A__											=						use_input_mask
        A__											=						use_token_type_ids
        A__											=						use_labels
        A__											=						vocab_size
        A__											=						block_sizes
        A__											=						num_decoder_layers
        A__											=						d_model
        A__											=						n_head
        A__											=						d_head
        A__											=						d_inner
        A__											=						hidden_act
        A__											=						hidden_dropout
        A__											=						attention_dropout
        A__											=						activation_dropout
        A__											=						max_position_embeddings
        A__											=						type_vocab_size
        A__											=						2
        A__											=						num_labels
        A__											=						num_choices
        A__											=						scope
        A__											=						initializer_std

        # Used in the tests to check the size of the first attention layer
        A__											=						n_head
        # Used in the tests to check the size of the first hidden state
        A__											=						self.d_model
        # Used in the tests to check the number of output hidden states/attentions
        A__											=						sum(self.block_sizes ) + (0 if base else self.num_decoder_layers)
        # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
        # the last hidden state of the first block (which is the first hidden state of the decoder).
        if not base:
               A__											=						self.num_hidden_layers + 2



 def 				UpperCamelCase			(							self ):
        A__											=						ids_tensor([self.batch_size, self.seq_length],self.vocab_size )

        A__											=						None
        if self.use_input_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__											=						None
        A__											=						None
        A__											=						None
        if self.use_labels:
               A__											=						ids_tensor([self.batch_size],self.type_sequence_label_size )
               A__											=						ids_tensor([self.batch_size, self.seq_length],self.num_labels )
               A__											=						ids_tensor([self.batch_size],self.num_choices )

        A__											=						FunnelConfig(
            vocab_size=self.vocab_size,block_sizes=self.block_sizes,num_decoder_layers=self.num_decoder_layers,d_model=self.d_model,n_head=self.n_head,d_head=self.d_head,d_inner=self.d_inner,hidden_act=self.hidden_act,hidden_dropout=self.hidden_dropout,attention_dropout=self.attention_dropout,activation_dropout=self.activation_dropout,max_position_embeddings=self.max_position_embeddings,type_vocab_size=self.type_vocab_size,initializer_std=self.initializer_std,)

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
        )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						TFFunnelModel(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )

        A__											=						[input_ids, input_mask]
        A__											=						model(__lowerCamelCase )

        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.d_model) )

        A__											=						False
        A__											=						TFFunnelModel(config=__lowerCamelCase )
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.d_model) )

        A__											=						False
        A__											=						TFFunnelModel(config=__lowerCamelCase )
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.d_model) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						TFFunnelBaseModel(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )

        A__											=						[input_ids, input_mask]
        A__											=						model(__lowerCamelCase )

        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, 2, self.d_model) )

        A__											=						False
        A__											=						TFFunnelBaseModel(config=__lowerCamelCase )
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, 3, self.d_model) )

        A__											=						False
        A__											=						TFFunnelBaseModel(config=__lowerCamelCase )
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, 2, self.d_model) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						TFFunnelForPreTraining(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						TFFunnelForMaskedLM(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						self.num_labels
        A__											=						TFFunnelForSequenceClassification(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_labels) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						self.num_choices
        A__											=						TFFunnelForMultipleChoice(config=__lowerCamelCase )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						tf.tile(tf.expand_dims(__lowerCamelCase,1 ),(1, self.num_choices, 1) )
        A__											=						{
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_choices) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						self.num_labels
        A__											=						TFFunnelForTokenClassification(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.num_labels) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__											=						TFFunnelForQuestionAnswering(config=__lowerCamelCase )
        A__											=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        A__											=						model(__lowerCamelCase )
        self.parent.assertEqual(result.start_logits.shape,(self.batch_size, self.seq_length) )
        self.parent.assertEqual(result.end_logits.shape,(self.batch_size, self.seq_length) )




 def 				UpperCamelCase			(							self ):
        A__											=						self.prepare_config_and_inputs()
        (
            (
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	(
            A__
        )					,	
        )											=						config_and_inputs
        A__											=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict




@require_tf
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__				,					UpperCamelCase__				,					unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	(
     (
         TFFunnelModel,
         TFFunnelForMaskedLM,
         TFFunnelForPreTraining,
         TFFunnelForQuestionAnswering,
         TFFunnelForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 __SCREAMING_SNAKE_CASE							       =	(
     {
         '''feature-extraction''': (TFFunnelBaseModel, TFFunnelModel),
         '''fill-mask''': TFFunnelForMaskedLM,
         '''question-answering''': TFFunnelForQuestionAnswering,
         '''text-classification''': TFFunnelForSequenceClassification,
         '''token-classification''': TFFunnelForTokenClassification,
         '''zero-shot''': TFFunnelForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __SCREAMING_SNAKE_CASE							       =	False
 __SCREAMING_SNAKE_CASE							       =	False



 def 				UpperCamelCase			(							self ):
        A__											=						TFFunnelModelTester(self )
        A__											=						ConfigTester(self,config_class=__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        self.config_tester.run_common_tests()



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase )




 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase )




@require_tf
class     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__				,					unittest.TestCase      ):
 __SCREAMING_SNAKE_CASE							       =	(
     (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
 )
 __SCREAMING_SNAKE_CASE							       =	False
 __SCREAMING_SNAKE_CASE							       =	False



 def 				UpperCamelCase			(							self ):
        A__											=						TFFunnelModelTester(self,base=__lowerCamelCase )
        A__											=						ConfigTester(self,config_class=__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        self.config_tester.run_common_tests()



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_base_model(*__lowerCamelCase )



 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase )




 def 				UpperCamelCase			(							self ):
        A__											=						self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowerCamelCase )












class     SCREAMING_SNAKE_CASE__ :



 def __init__(							self,__lowerCamelCase ):
        A__											=						set_counts
        A__											=						max(__lowerCamelCase )
        A__											=						len(__lowerCamelCase )
        A__											=						[1] * num_sets
        A__											=						list(range(__lowerCamelCase ) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase ):
        A__											=						self.get_parent(__lowerCamelCase )
        A__											=						self.get_parent(__lowerCamelCase )

        if src_parent == dst_parent:
               return False

        if self.ranks[dst_parent] >= self.ranks[src_parent]:
               self.set_counts[dst_parent] += self.set_counts[src_parent]
               A__											=						0
               A__											=						dst_parent
               if self.ranks[dst_parent] == self.ranks[src_parent]:
                      self.ranks[dst_parent] += 1
               A__											=						self.set_counts[dst_parent]
        else:
               self.set_counts[src_parent] += self.set_counts[dst_parent]
               A__											=						0
               A__											=						src_parent
               A__											=						self.set_counts[src_parent]

        A__											=						max(self.max_set,__lowerCamelCase )
        return True




 def 				UpperCamelCase			(							self,__lowerCamelCase ):
        if self.parents[disj_set] == disj_set:
               return disj_set
        A__											=						self.get_parent(self.parents[disj_set] )
        return self.parents[disj_set]





