



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




       """simple docstring"""






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

             """simple docstring"""





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

             # Used in the tests to check the size of the first attention layer
             UpperCAmelCase_  :	Any				      =		n_head
             # Used in the tests to check the size of the first hidden state
             UpperCAmelCase_  :	Any				      =		self.d_model
             # Used in the tests to check the number of output hidden states/attentions
             UpperCAmelCase_  :	Dict				      =		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_  :	Dict				      =		self.num_hidden_layers + 2



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size)

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

             UpperCAmelCase_  :	str				      =		None
             if self.use_token_type_ids:
                   UpperCAmelCase_  :	Optional[Any]				      =		ids_tensor([self.batch_size, self.seq_length]     ,self.type_vocab_size)

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

             UpperCAmelCase_  :	Optional[Any]				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		TFFunnelModel(config=lowercase)
             UpperCAmelCase_  :	Optional[Any]				      =		{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
             UpperCAmelCase_  :	str				      =		model(lowercase)

             UpperCAmelCase_  :	Dict				      =		[input_ids, input_mask]
             UpperCAmelCase_  :	List[str]				      =		model(lowercase)

             UpperCAmelCase_  :	Optional[int]				      =		model(lowercase)
             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=lowercase)
             UpperCAmelCase_  :	int				      =		model(lowercase)
             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=lowercase)
             UpperCAmelCase_  :	str				      =		model(lowercase)
             self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, self.seq_length, self.d_model))



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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

             UpperCAmelCase_  :	List[str]				      =		[input_ids, input_mask]
             UpperCAmelCase_  :	Any				      =		model(lowercase)

             UpperCAmelCase_  :	int				      =		model(lowercase)
             self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, 2, self.d_model))

             UpperCAmelCase_  :	Any				      =		False
             UpperCAmelCase_  :	List[Any]				      =		TFFunnelBaseModel(config=lowercase)
             UpperCAmelCase_  :	str				      =		model(lowercase)
             self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, 3, self.d_model))

             UpperCAmelCase_  :	int				      =		False
             UpperCAmelCase_  :	str				      =		TFFunnelBaseModel(config=lowercase)
             UpperCAmelCase_  :	Tuple				      =		model(lowercase)
             self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, 2, self.d_model))



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.num_labels
             UpperCAmelCase_  :	Tuple				      =		TFFunnelForSequenceClassification(config=lowercase)
             UpperCAmelCase_  :	List[str]				      =		{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
             UpperCAmelCase_  :	Any				      =		model(lowercase)
             self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels))



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		TFFunnelForQuestionAnswering(config=lowercase)
             UpperCAmelCase_  :	Tuple				      =		{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
             UpperCAmelCase_  :	Any				      =		model(lowercase)
             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  A_   (     self):

             """simple docstring"""





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


@require_tf
class   snake_case_ (lowercase__ , lowercase__ , unittest.TestCase  ):




       """simple docstring"""






       _lowerCamelCase			    =				(
           (
               TFFunnelModel,
               TFFunnelForMaskedLM,
               TFFunnelForPreTraining,
               TFFunnelForQuestionAnswering,
               TFFunnelForTokenClassification,
           )
           if is_tf_available()
           else ()
       )
       _lowerCamelCase			    =				(
           {
               """feature-extraction""": (TFFunnelBaseModel, TFFunnelModel),
               """fill-mask""": TFFunnelForMaskedLM,
               """question-answering""": TFFunnelForQuestionAnswering,
               """text-classification""": TFFunnelForSequenceClassification,
               """token-classification""": TFFunnelForTokenClassification,
               """zero-shot""": TFFunnelForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       _lowerCamelCase			    =				False
       _lowerCamelCase			    =				False



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		TFFunnelModelTester(self)
             UpperCAmelCase_  :	Tuple				      =		ConfigTester(self     ,config_class=lowercase)



       def  A_   (     self):

             """simple docstring"""





             self.config_tester.run_common_tests()



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_pretraining(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_question_answering(*lowercase)


@require_tf
class   snake_case_ (lowercase__ , unittest.TestCase  ):




       """simple docstring"""






       _lowerCamelCase			    =				(
           (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
       )
       _lowerCamelCase			    =				False
       _lowerCamelCase			    =				False



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		TFFunnelModelTester(self     ,base=lowercase)
             UpperCAmelCase_  :	Optional[int]				      =		ConfigTester(self     ,config_class=lowercase)



       def  A_   (     self):

             """simple docstring"""





             self.config_tester.run_common_tests()



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_base_model(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_sequence_classification(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_multiple_choice(*lowercase)





from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''google/umt5-small''': '''https://huggingface.co/google/umt5-small/resolve/main/config.json''',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""umt5"""
       _lowerCamelCase			    =				["""past_key_values"""]



       def __init__(     self     ,lowercase=250112     ,lowercase=512     ,lowercase=64     ,lowercase=1024     ,lowercase=8     ,lowercase=None     ,lowercase=6     ,lowercase=32     ,lowercase=128     ,lowercase=0.1     ,lowercase=1E-6     ,lowercase=1.0     ,lowercase="gated-gelu"     ,lowercase=True     ,lowercase=True     ,lowercase="T5Tokenizer"     ,lowercase=True     ,lowercase=0     ,lowercase=1     ,lowercase=0     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(
                 is_encoder_decoder=lowercase     ,tokenizer_class=lowercase     ,tie_word_embeddings=lowercase     ,pad_token_id=lowercase     ,eos_token_id=lowercase     ,decoder_start_token_id=lowercase     ,**lowercase     ,)
             UpperCAmelCase_  :	Optional[int]				      =		vocab_size
             UpperCAmelCase_  :	Any				      =		d_model
             UpperCAmelCase_  :	Any				      =		d_kv
             UpperCAmelCase_  :	int				      =		d_ff
             UpperCAmelCase_  :	Tuple				      =		num_layers
             UpperCAmelCase_  :	int				      =		(
                 num_decoder_layers if num_decoder_layers is not None else self.num_layers
             )  # default = symmetry
             UpperCAmelCase_  :	Optional[int]				      =		num_heads
             UpperCAmelCase_  :	str				      =		relative_attention_num_buckets
             UpperCAmelCase_  :	Any				      =		relative_attention_max_distance
             UpperCAmelCase_  :	Optional[Any]				      =		dropout_rate
             UpperCAmelCase_  :	Union[str, Any]				      =		layer_norm_epsilon
             UpperCAmelCase_  :	Optional[Any]				      =		initializer_factor
             UpperCAmelCase_  :	int				      =		feed_forward_proj
             UpperCAmelCase_  :	str				      =		use_cache

             UpperCAmelCase_  :	List[str]				      =		self.feed_forward_proj.split("-")
             UpperCAmelCase_  :	Any				      =		act_info[-1]
             UpperCAmelCase_  :	Optional[int]				      =		act_info[0] == "gated"

             if len(lowercase) > 1 and act_info[0] != "gated" or len(lowercase) > 2:
                   raise ValueError(
                       F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
                       "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
                       "'gated-gelu' or 'relu'")

             if feed_forward_proj == "gated-gelu":
                   UpperCAmelCase_  :	Tuple				      =		"gelu_new"



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.d_model



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.num_heads



       @property
       def  A_   (     self):

             """simple docstring"""





             return self.num_layers


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		{
                 "input_ids": {0: "batch", 1: "encoder_sequence"},
                 "attention_mask": {0: "batch", 1: "encoder_sequence"},
             }
             if self.use_past:
                   UpperCAmelCase_  :	Union[str, Any]				      =		"past_encoder_sequence + sequence"
                   UpperCAmelCase_  :	Optional[int]				      =		{0: "batch"}
                   UpperCAmelCase_  :	Union[str, Any]				      =		{0: "batch", 1: "past_decoder_sequence + sequence"}
             else:
                   UpperCAmelCase_  :	Optional[int]				      =		{0: "batch", 1: "decoder_sequence"}
                   UpperCAmelCase_  :	Dict				      =		{0: "batch", 1: "decoder_sequence"}

             if self.use_past:
                   self.fill_with_past_key_values_(lowercase     ,direction="inputs")

             return common_inputs



       @property
       # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
       def  A_   (     self):

             """simple docstring"""





             return 13



       @property
       def  A_   (     self):

             """simple docstring"""





             return 5E-4
