"""simple docstring"""


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,				_UpperCAmelCase,				_UpperCAmelCase=13,				_UpperCAmelCase=7,				_UpperCAmelCase=True,				_UpperCAmelCase=True,				_UpperCAmelCase=True,				_UpperCAmelCase=True,				_UpperCAmelCase=99,				_UpperCAmelCase=[1, 1, 2],				_UpperCAmelCase=1,				_UpperCAmelCase=32,				_UpperCAmelCase=4,				_UpperCAmelCase=8,				_UpperCAmelCase=37,				_UpperCAmelCase="gelu_new",				_UpperCAmelCase=0.1,				_UpperCAmelCase=0.1,				_UpperCAmelCase=0.0,				_UpperCAmelCase=512,				_UpperCAmelCase=3,				_UpperCAmelCase=0.02,				_UpperCAmelCase=3,				_UpperCAmelCase=4,				_UpperCAmelCase=None,				_UpperCAmelCase=False,				):
        '''simple docstring'''



        lowercase__				 =							parent
        lowercase__				 =							batch_size
        lowercase__				 =							seq_length
        lowercase__				 =							is_training
        lowercase__				 =							use_input_mask
        lowercase__				 =							use_token_type_ids
        lowercase__				 =							use_labels
        lowercase__				 =							vocab_size
        lowercase__				 =							block_sizes
        lowercase__				 =							num_decoder_layers
        lowercase__				 =							d_model
        lowercase__				 =							n_head
        lowercase__				 =							d_head
        lowercase__				 =							d_inner
        lowercase__				 =							hidden_act
        lowercase__				 =							hidden_dropout
        lowercase__				 =							attention_dropout
        lowercase__				 =							activation_dropout
        lowercase__				 =							max_position_embeddings
        lowercase__				 =							type_vocab_size
        lowercase__				 =							2
        lowercase__				 =							num_labels
        lowercase__				 =							num_choices
        lowercase__				 =							scope
        lowercase__				 =							initializer_std

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


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



        lowercase__				 =							ids_tensor([self.batch_size, self.seq_length],				self.vocab_size  )

        lowercase__				 =							None
        if self.use_input_mask:
             lowercase__				 =							random_attention_mask([self.batch_size, self.seq_length]  )

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

        lowercase__				 =							None
        lowercase__				 =							None
        lowercase__				 =							None
        if self.use_labels:
             lowercase__				 =							ids_tensor([self.batch_size],				self.type_sequence_label_size  )
             lowercase__				 =							ids_tensor([self.batch_size, self.seq_length],				self.num_labels  )
             lowercase__				 =							ids_tensor([self.batch_size],				self.num_choices  )

        lowercase__				 =							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 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							TFFunnelModel(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )

        lowercase__				 =							[input_ids, input_mask]
        lowercase__				 =							model(_UpperCAmelCase  )

        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.last_hidden_state.shape,				(self.batch_size, self.seq_length, self.d_model)  )

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

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


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							TFFunnelBaseModel(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )

        lowercase__				 =							[input_ids, input_mask]
        lowercase__				 =							model(_UpperCAmelCase  )

        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.last_hidden_state.shape,				(self.batch_size, 2, self.d_model)  )

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

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


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							TFFunnelForPreTraining(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.seq_length)  )


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							TFFunnelForMaskedLM(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.seq_length, self.vocab_size)  )


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							self.num_labels
        lowercase__				 =							TFFunnelForSequenceClassification(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.num_labels)  )


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							self.num_choices
        lowercase__				 =							TFFunnelForMultipleChoice(config=_UpperCAmelCase  )
        lowercase__				 =							tf.tile(tf.expand_dims(_UpperCAmelCase,				1  ),				(1, self.num_choices, 1)  )
        lowercase__				 =							tf.tile(tf.expand_dims(_UpperCAmelCase,				1  ),				(1, self.num_choices, 1)  )
        lowercase__				 =							tf.tile(tf.expand_dims(_UpperCAmelCase,				1  ),				(1, self.num_choices, 1)  )
        lowercase__				 =							{
            "input_ids": multiple_choice_inputs_ids,
            "attention_mask": multiple_choice_input_mask,
            "token_type_ids": multiple_choice_token_type_ids,
        }
        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.num_choices)  )


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							self.num_labels
        lowercase__				 =							TFFunnelForTokenClassification(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )
        self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.seq_length, self.num_labels)  )


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				):
        '''simple docstring'''



        lowercase__				 =							TFFunnelForQuestionAnswering(config=_UpperCAmelCase  )
        lowercase__				 =							{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
        lowercase__				 =							model(_UpperCAmelCase  )
        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 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							self.prepare_config_and_inputs()
        (
            (
            lowercase__
        ) ,						(
            lowercase__
        ) ,						(
            lowercase__
        ) ,						(
            lowercase__
        ) ,						(
            lowercase__
        ) ,						(
            lowercase__
        ) ,						(
            lowercase__
        ) ,						
        )				 =							config_and_inputs
        lowercase__				 =							{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict






@require_tf
class    a__    (   _a					,					_a					,					unittest.TestCase	):
   snake_case_							=      (
       (
           TFFunnelModel,
           TFFunnelForMaskedLM,
           TFFunnelForPreTraining,
           TFFunnelForQuestionAnswering,
           TFFunnelForTokenClassification,
       )
       if is_tf_available()
       else ()
   )
   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 {}
   )
   snake_case_							=      False
   snake_case_							=      False


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



        lowercase__				 =							TFFunnelModelTester(self  )
        lowercase__				 =							ConfigTester(self,				config_class=_UpperCAmelCase  )


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



        self.config_tester.run_common_tests()


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*_UpperCAmelCase  )






@require_tf
class    a__    (   _a					,					unittest.TestCase	):
   snake_case_							=      (
       (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
   )
   snake_case_							=      False
   snake_case_							=      False


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



        lowercase__				 =							TFFunnelModelTester(self,				base=_UpperCAmelCase  )
        lowercase__				 =							ConfigTester(self,				config_class=_UpperCAmelCase  )


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



        self.config_tester.run_common_tests()


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_base_model(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*_UpperCAmelCase  )


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



        lowercase__				 =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*_UpperCAmelCase  )





"""simple docstring"""


import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()






class    a__    (   _a					,					unittest.TestCase	):
   snake_case_							=      PriorTransformer
   snake_case_							=      "hidden_states"


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							4
        lowercase__				 =							8
        lowercase__				 =							7

        lowercase__				 =							floats_tensor((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )

        lowercase__				 =							floats_tensor((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )
        lowercase__				 =							floats_tensor((batch_size, num_embeddings, embedding_dim)  ).to(_UpperCAmelCase  )

        return {
            "hidden_states": hidden_states,
            "timestep": 2,
            "proj_embedding": proj_embedding,
            "encoder_hidden_states": encoder_hidden_states,
        }


   def 	snake_case__	(     self,				_UpperCAmelCase=0  ):
        '''simple docstring'''



        torch.manual_seed(_UpperCAmelCase  )
        lowercase__				 =							4
        lowercase__				 =							8
        lowercase__				 =							7

        lowercase__				 =							torch.randn((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )

        lowercase__				 =							torch.randn((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )
        lowercase__				 =							torch.randn((batch_size, num_embeddings, embedding_dim)  ).to(_UpperCAmelCase  )

        return {
            "hidden_states": hidden_states,
            "timestep": 2,
            "proj_embedding": proj_embedding,
            "encoder_hidden_states": encoder_hidden_states,
        }


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return (4, 8)


   @property
   def 	snake_case__	(     self  ):
        '''simple docstring'''



        return (4, 8)


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



        lowercase__				 =							{
            "num_attention_heads": 2,
            "attention_head_dim": 4,
            "num_layers": 2,
            "embedding_dim": 8,
            "num_embeddings": 7,
            "additional_embeddings": 4,
        }
        lowercase__				 =							self.dummy_input
        return init_dict, inputs_dict


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



        lowercase__ ,						lowercase__				 =							PriorTransformer.from_pretrained(
            "hf-internal-testing/prior-dummy",				output_loading_info=_UpperCAmelCase  )
        self.assertIsNotNone(_UpperCAmelCase  )
        self.assertEqual(len(loading_info["missing_keys"]  ),				0  )

        model.to(_UpperCAmelCase  )
        lowercase__				 =							model(**self.dummy_input  )[0]

        assert hidden_states is not None, "Make sure output is not None"


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



        lowercase__ ,						lowercase__				 =							self.prepare_init_args_and_inputs_for_common()

        lowercase__				 =							self.model_class(**_UpperCAmelCase  )
        lowercase__				 =							inspect.signature(model.forward  )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowercase__				 =							[*signature.parameters.keys()]

        lowercase__				 =							["hidden_states", "timestep"]
        self.assertListEqual(arg_names[:2],				_UpperCAmelCase  )


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



        lowercase__				 =							PriorTransformer.from_pretrained("hf-internal-testing/prior-dummy"  )
        lowercase__				 =							model.to(_UpperCAmelCase  )

        if hasattr(_UpperCAmelCase,				"set_default_attn_processor"  ):
             model.set_default_attn_processor()

        lowercase__				 =							self.get_dummy_seed_input()

        with torch.no_grad():
             lowercase__				 =							model(**_UpperCAmelCase  )[0]

        lowercase__				 =							output[0, :5].flatten().cpu()
        print(_UpperCAmelCase  )

        # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
        # the expected output slices are not the same for CPU and GPU.
        lowercase__				 =							torch.tensor([-1.3_436, -0.2_870, 0.7_538, 0.4_368, -0.0_239]  )
        self.assertTrue(torch_all_close(_UpperCAmelCase,				_UpperCAmelCase,				rtol=1E-2  )  )






@slow
class    a__    (   unittest.TestCase	):


   def 	snake_case__	(     self,				_UpperCAmelCase=1,				_UpperCAmelCase=768,				_UpperCAmelCase=77,				_UpperCAmelCase=0  ):
        '''simple docstring'''



        torch.manual_seed(_UpperCAmelCase  )
        lowercase__				 =							batch_size
        lowercase__				 =							embedding_dim
        lowercase__				 =							num_embeddings

        lowercase__				 =							torch.randn((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )

        lowercase__				 =							torch.randn((batch_size, embedding_dim)  ).to(_UpperCAmelCase  )
        lowercase__				 =							torch.randn((batch_size, num_embeddings, embedding_dim)  ).to(_UpperCAmelCase  )

        return {
            "hidden_states": hidden_states,
            "timestep": 2,
            "proj_embedding": proj_embedding,
            "encoder_hidden_states": encoder_hidden_states,
        }


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



        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()


   @parameterized.expand(
       [
           # fmt: off
           [13, [-0.5_861,  0.1_283, -0.0_931,  0.0_882,  0.4_476,  0.1_329, -0.0_498,  0.0_640]],
           [37, [-0.4_913,  0.0_110, -0.0_483,  0.0_541,  0.4_954, -0.0_170,  0.0_354,  0.1_651]],
           # fmt: on
       ]  )
   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        lowercase__				 =							PriorTransformer.from_pretrained("kandinsky-community/kandinsky-2-1-prior",				subfolder="prior"  )
        model.to(_UpperCAmelCase  )
        lowercase__				 =							self.get_dummy_seed_input(seed=_UpperCAmelCase  )

        with torch.no_grad():
             lowercase__				 =							model(**_UpperCAmelCase  )[0]

        assert list(sample.shape  ) == [1, 768]

        lowercase__				 =							sample[0, :8].flatten().cpu()
        print(_UpperCAmelCase  )
        lowercase__				 =							torch.tensor(_UpperCAmelCase  )

        assert torch_all_close(_UpperCAmelCase,				_UpperCAmelCase,				atol=1E-3  )




