




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






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

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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size )

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

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

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

							SCREAMING_SNAKE_CASE_   : int					  =      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      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : Dict					  =      TFFunnelModel(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : int					  =      model(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      [input_ids, input_mask]
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      model(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      model(snake_case__ )
							self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.d_model) )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      False
							SCREAMING_SNAKE_CASE_   : int					  =      TFFunnelModel(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      model(snake_case__ )
							self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.d_model) )

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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : Dict					  =      TFFunnelBaseModel(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      model(snake_case__ )

							SCREAMING_SNAKE_CASE_   : int					  =      [input_ids, input_mask]
							SCREAMING_SNAKE_CASE_   : str					  =      model(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Tuple					  =      model(snake_case__ )
							self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, 2, self.d_model) )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      False
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      TFFunnelBaseModel(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      model(snake_case__ )
							self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, 3, self.d_model) )

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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : int					  =      TFFunnelForPreTraining(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : Dict					  =      model(snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length) )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : int					  =      TFFunnelForMaskedLM(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : str					  =      model(snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size) )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.num_labels
							SCREAMING_SNAKE_CASE_   : Dict					  =      TFFunnelForSequenceClassification(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      model(snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels) )






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






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : Dict					  =      self.num_labels
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      TFFunnelForTokenClassification(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      model(snake_case__ )
							self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels) )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__      ,):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      TFFunnelForQuestionAnswering(config=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
							SCREAMING_SNAKE_CASE_   : List[str]					  =      model(snake_case__ )
							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 ):
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self.prepare_config_and_inputs()
							(
							    SCREAMING_SNAKE_CASE_
							)   : Dict					  =      config_and_inputs
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
							return config, inputs_dict







@require_tf
class 				lowerCAmelCase_       (					__A						,     __A						,     unittest.TestCase							):
				__a     :       Dict						=     (
				    (
				        TFFunnelModel,
				        TFFunnelForMaskedLM,
				        TFFunnelForPreTraining,
				        TFFunnelForQuestionAnswering,
				        TFFunnelForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				__a     :       Any						=     (
				    {
				        "feature-extraction": (TFFunnelBaseModel, TFFunnelModel),
				        "fill-mask": TFFunnelForMaskedLM,
				        "question-answering": TFFunnelForQuestionAnswering,
				        "text-classification": TFFunnelForSequenceClassification,
				        "token-classification": TFFunnelForTokenClassification,
				        "zero-shot": TFFunnelForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				__a     :       Optional[Any]						=     False
				__a     :       Optional[Any]						=     False






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      TFFunnelModelTester(self )
							SCREAMING_SNAKE_CASE_   : Any					  =      ConfigTester(self      ,config_class=snake_case__ )






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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Any					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*snake_case__ )







@require_tf
class 				lowerCAmelCase_       (					__A						,     unittest.TestCase							):
				__a     :       Dict						=     (
				    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
				)
				__a     :       List[str]						=     False
				__a     :       Union[str, Any]						=     False






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      TFFunnelModelTester(self      ,base=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      ConfigTester(self      ,config_class=snake_case__ )






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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_base_model(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*snake_case__ )






import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			@property
			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Union[str, Any]:

				"""simple docstring"""

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   str					  =							UNetaDModel(
				    block_out_channels=(32, 64)   ,   layers_per_block=2   ,   sample_size=32   ,   in_channels=3   ,   out_channels=3   ,   down_block_types=("DownBlock2D", "AttnDownBlock2D")   ,   up_block_types=("AttnUpBlock2D", "UpBlock2D")   ,   )
				return model






			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.dummy_uncond_unet
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Any					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"   ,   return_dict=a						)[0]

				SCREAMING_SNAKE_CASE :   List[Any]					  =							image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE :   Any					  =							image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]						)

				assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < 1e-2






@slow
@require_torch
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			int						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							"google/ncsnpp-celebahq-256"
				SCREAMING_SNAKE_CASE :   List[Any]					  =							UNetaDModel.from_pretrained(a						)
				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Any					  =							pipe(num_inference_steps=20   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							image[0, -3:, -3:, -1]
				assert image.shape == (1, 256, 256, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586]						)

				assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2