





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 					A__   :





      def __init__( self       ,							A_       ,							A_=13       ,							A_=7       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=True       ,							A_=99       ,							A_=[1, 1, 2]       ,							A_=1       ,							A_=32       ,							A_=4       ,							A_=8       ,							A_=37       ,							A_="gelu_new"       ,							A_=0.1       ,							A_=0.1       ,							A_=0.0       ,							A_=512       ,							A_=3       ,							A_=0.02       ,							A_=3       ,							A_=4       ,							A_=None       ,							A_=False       ,							):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   parent
       UpperCamelCase	:					Optional[int]       	=   batch_size
       UpperCamelCase	:					List[Any]       	=   seq_length
       UpperCamelCase	:					Tuple       	=   is_training
       UpperCamelCase	:					Optional[Any]       	=   use_input_mask
       UpperCamelCase	:					Any       	=   use_token_type_ids
       UpperCamelCase	:					Optional[Any]       	=   use_labels
       UpperCamelCase	:					List[str]       	=   vocab_size
       UpperCamelCase	:					Tuple       	=   block_sizes
       UpperCamelCase	:					List[Any]       	=   num_decoder_layers
       UpperCamelCase	:					Optional[Any]       	=   d_model
       UpperCamelCase	:					List[Any]       	=   n_head
       UpperCamelCase	:					Union[str, Any]       	=   d_head
       UpperCamelCase	:					Tuple       	=   d_inner
       UpperCamelCase	:					str       	=   hidden_act
       UpperCamelCase	:					List[Any]       	=   hidden_dropout
       UpperCamelCase	:					Any       	=   attention_dropout
       UpperCamelCase	:					Tuple       	=   activation_dropout
       UpperCamelCase	:					Optional[int]       	=   max_position_embeddings
       UpperCamelCase	:					Tuple       	=   type_vocab_size
       UpperCamelCase	:					List[str]       	=   2
       UpperCamelCase	:					Optional[Any]       	=   num_labels
       UpperCamelCase	:					Union[str, Any]       	=   num_choices
       UpperCamelCase	:					Tuple       	=   scope
       UpperCamelCase	:					Dict       	=   initializer_std

       # Used in the tests to check the size of the first attention layer
       UpperCamelCase	:					str       	=   n_head
       # Used in the tests to check the size of the first hidden state
       UpperCamelCase	:					Optional[Any]       	=   self.d_model
       # Used in the tests to check the number of output hidden states/attentions
       UpperCamelCase	:					Optional[Any]       	=   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:
        UpperCamelCase	:					Union[str, Any]       	=   self.num_hidden_layers + 2





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




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

       UpperCamelCase	:					int       	=   None
       if self.use_input_mask:
        UpperCamelCase	:					List[str]       	=   random_attention_mask([self.batch_size, self.seq_length]   )

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

       UpperCamelCase	:					List[Any]       	=   None
       UpperCamelCase	:					Union[str, Any]       	=   None
       UpperCamelCase	:					Optional[int]       	=   None
       if self.use_labels:
        UpperCamelCase	:					Any       	=   ids_tensor([self.batch_size]       ,							self.type_sequence_label_size   )
        UpperCamelCase	:					Optional[int]       	=   ids_tensor([self.batch_size, self.seq_length]       ,							self.num_labels   )
        UpperCamelCase	:					Dict       	=   ids_tensor([self.batch_size]       ,							self.num_choices   )

       UpperCamelCase	:					Dict       	=   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       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   TFFunnelModel(config=A_   )
       UpperCamelCase	:					int       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					List[str]       	=   model(A_   )

       UpperCamelCase	:					Union[str, Any]       	=   [input_ids, input_mask]
       UpperCamelCase	:					Tuple       	=   model(A_   )

       UpperCamelCase	:					Dict       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.d_model)   )

       UpperCamelCase	:					int       	=   False
       UpperCamelCase	:					Optional[int]       	=   TFFunnelModel(config=A_   )
       UpperCamelCase	:					str       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.d_model)   )

       UpperCamelCase	:					Union[str, Any]       	=   False
       UpperCamelCase	:					str       	=   TFFunnelModel(config=A_   )
       UpperCamelCase	:					List[Any]       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.d_model)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					List[str]       	=   TFFunnelBaseModel(config=A_   )
       UpperCamelCase	:					Tuple       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					Optional[int]       	=   model(A_   )

       UpperCamelCase	:					List[Any]       	=   [input_ids, input_mask]
       UpperCamelCase	:					int       	=   model(A_   )

       UpperCamelCase	:					str       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, 2, self.d_model)   )

       UpperCamelCase	:					str       	=   False
       UpperCamelCase	:					Union[str, Any]       	=   TFFunnelBaseModel(config=A_   )
       UpperCamelCase	:					Optional[int]       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, 3, self.d_model)   )

       UpperCamelCase	:					Union[str, Any]       	=   False
       UpperCamelCase	:					List[Any]       	=   TFFunnelBaseModel(config=A_   )
       UpperCamelCase	:					Dict       	=   model(A_   )
       self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, 2, self.d_model)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   TFFunnelForPreTraining(config=A_   )
       UpperCamelCase	:					List[str]       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					Union[str, Any]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   TFFunnelForMaskedLM(config=A_   )
       UpperCamelCase	:					List[Any]       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					List[Any]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   self.num_labels
       UpperCamelCase	:					Dict       	=   TFFunnelForSequenceClassification(config=A_   )
       UpperCamelCase	:					Optional[int]       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					Any       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   self.num_choices
       UpperCamelCase	:					str       	=   TFFunnelForMultipleChoice(config=A_   )
       UpperCamelCase	:					int       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					Optional[Any]       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					str       	=   tf.tile(tf.expand_dims(A_       ,							1   )       ,							(1, self.num_choices, 1)   )
       UpperCamelCase	:					Any       	=   {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
           "token_type_ids": multiple_choice_token_type_ids,
       }
       UpperCamelCase	:					List[str]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_choices)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					List[Any]       	=   self.num_labels
       UpperCamelCase	:					Optional[Any]       	=   TFFunnelForTokenClassification(config=A_   )
       UpperCamelCase	:					int       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					List[str]       	=   model(A_   )
       self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)   )





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							A_       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   TFFunnelForQuestionAnswering(config=A_   )
       UpperCamelCase	:					Optional[Any]       	=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
       UpperCamelCase	:					Union[str, Any]       	=   model(A_   )
       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   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    (
           UpperCamelCase
       )	,    
       )	:					Dict       	=   config_and_inputs
       UpperCamelCase	:					Any       	=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict



@require_tf
class 					A__   (       __snake_case       ,   __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:List[str]						    =   (
          (
              TFFunnelModel,
              TFFunnelForMaskedLM,
              TFFunnelForPreTraining,
              TFFunnelForQuestionAnswering,
              TFFunnelForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase	:Dict						    =   (
          {
              'feature-extraction': (TFFunnelBaseModel, TFFunnelModel),
              'fill-mask': TFFunnelForMaskedLM,
              'question-answering': TFFunnelForQuestionAnswering,
              'text-classification': TFFunnelForSequenceClassification,
              'token-classification': TFFunnelForTokenClassification,
              'zero-shot': TFFunnelForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase	:Union[str, Any]						    =   False
      _UpperCAmelCase	:Tuple						    =   False





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   TFFunnelModelTester(self   )
       UpperCamelCase	:					List[str]       	=   ConfigTester(self       ,							config_class=A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					str       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A_   )



@require_tf
class 					A__   (       __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:Optional[Any]						    =   (
          (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
      )
      _UpperCAmelCase	:List[Any]						    =   False
      _UpperCAmelCase	:List[Any]						    =   False





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   TFFunnelModelTester(self       ,							base=A_   )
       UpperCamelCase	:					Dict       	=   ConfigTester(self       ,							config_class=A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					List[str]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_base_model(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*A_   )





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*A_   )







class 					A__   :





      def __init__( self       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Union[str, Any]       	=   set_counts
       UpperCamelCase	:					int       	=   max(A_   )
       UpperCamelCase	:					Optional[Any]       	=   len(A_   )
       UpperCamelCase	:					Union[str, Any]       	=   [1] * num_sets
       UpperCamelCase	:					Union[str, Any]       	=   list(range(A_   )   )





      def        __UpperCamelCase( self       ,							A_       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.get_parent(A_   )
       UpperCamelCase	:					Optional[int]       	=   self.get_parent(A_   )

       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]
        UpperCamelCase	:					int       	=   0
        UpperCamelCase	:					Dict       	=   dst_parent
        if self.ranks[dst_parent] == self.ranks[src_parent]:
         self.ranks[dst_parent] += 1
        UpperCamelCase	:					Optional[int]       	=   self.set_counts[dst_parent]
       else:
        self.set_counts[src_parent] += self.set_counts[dst_parent]
        UpperCamelCase	:					Any       	=   0
        UpperCamelCase	:					Optional[int]       	=   src_parent
        UpperCamelCase	:					int       	=   self.set_counts[src_parent]

       UpperCamelCase	:					Any       	=   max(self.max_set       ,							A_   )
       return True





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       if self.parents[disj_set] == disj_set:
        return disj_set
       UpperCamelCase	:					Optional[int]       	=   self.get_parent(self.parents[disj_set]   )
       return self.parents[disj_set]
