




"""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   UpperCAmelCase_    :
		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_=5_12				,						UpperCamelCase_=3				,						UpperCamelCase_=0.0_2				,						UpperCamelCase_=3				,						UpperCamelCase_=4				,						UpperCamelCase_=None				,						UpperCamelCase_=False				,						)						->       Tuple:
					__lowercase			:  int								=       parent
					__lowercase			:  Optional[int]								=       batch_size
					__lowercase			:  Dict								=       seq_length
					__lowercase			:  int								=       is_training
					__lowercase			:  int								=       use_input_mask
					__lowercase			:  str								=       use_token_type_ids
					__lowercase			:  List[Any]								=       use_labels
					__lowercase			:  Union[str, Any]								=       vocab_size
					__lowercase			:  Any								=       block_sizes
					__lowercase			:  Optional[Any]								=       num_decoder_layers
					__lowercase			:  Union[str, Any]								=       d_model
					__lowercase			:  str								=       n_head
					__lowercase			:  List[str]								=       d_head
					__lowercase			:  str								=       d_inner
					__lowercase			:  Tuple								=       hidden_act
					__lowercase			:  str								=       hidden_dropout
					__lowercase			:  Optional[Any]								=       attention_dropout
					__lowercase			:  Dict								=       activation_dropout
					__lowercase			:  Dict								=       max_position_embeddings
					__lowercase			:  int								=       type_vocab_size
					__lowercase			:  List[Any]								=       2
					__lowercase			:  Optional[int]								=       num_labels
					__lowercase			:  int								=       num_choices
					__lowercase			:  List[Any]								=       scope
					__lowercase			:  Optional[Any]								=       initializer_std

					# Used in the tests to check the size of the first attention layer
					__lowercase			:  List[str]								=       n_head
					# Used in the tests to check the size of the first hidden state
					__lowercase			:  Optional[int]								=       self.d_model
					# Used in the tests to check the number of output hidden states/attentions
					__lowercase			:  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:
								__lowercase			:  Dict								=       self.num_hidden_layers + 2
		def 	_lowerCamelCase							(						self       )						->       Optional[Any]:
					__lowercase			:  Optional[int]								=       ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size       )

					__lowercase			:  Optional[int]								=       None
					if self.use_input_mask:
								__lowercase			:  Union[str, Any]								=       random_attention_mask([self.batch_size, self.seq_length]       )

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

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

					__lowercase			:  Optional[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 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Dict:
					__lowercase			:  Any								=       TFFunnelModel(config=UpperCamelCase_       )
					__lowercase			:  Tuple								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  str								=       model(UpperCamelCase_       )

					__lowercase			:  int								=       [input_ids, input_mask]
					__lowercase			:  Optional[int]								=       model(UpperCamelCase_       )

					__lowercase			:  Any								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.d_model)       )

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

					__lowercase			:  Tuple								=       False
					__lowercase			:  Any								=       TFFunnelModel(config=UpperCamelCase_       )
					__lowercase			:  str								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.d_model)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Union[str, Any]:
					__lowercase			:  Optional[int]								=       TFFunnelBaseModel(config=UpperCamelCase_       )
					__lowercase			:  Tuple								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  Any								=       model(UpperCamelCase_       )

					__lowercase			:  List[Any]								=       [input_ids, input_mask]
					__lowercase			:  Any								=       model(UpperCamelCase_       )

					__lowercase			:  Any								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, 2, self.d_model)       )

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

					__lowercase			:  Any								=       False
					__lowercase			:  str								=       TFFunnelBaseModel(config=UpperCamelCase_       )
					__lowercase			:  Tuple								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, 2, self.d_model)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       List[Any]:
					__lowercase			:  List[str]								=       TFFunnelForPreTraining(config=UpperCamelCase_       )
					__lowercase			:  Tuple								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  List[Any]								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Any:
					__lowercase			:  str								=       TFFunnelForMaskedLM(config=UpperCamelCase_       )
					__lowercase			:  str								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  Any								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Tuple:
					__lowercase			:  int								=       self.num_labels
					__lowercase			:  List[str]								=       TFFunnelForSequenceClassification(config=UpperCamelCase_       )
					__lowercase			:  Any								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  Tuple								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.num_labels)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Dict:
					__lowercase			:  Optional[int]								=       self.num_choices
					__lowercase			:  List[Any]								=       TFFunnelForMultipleChoice(config=UpperCamelCase_       )
					__lowercase			:  Union[str, Any]								=       tf.tile(tf.expand_dims(UpperCamelCase_				,						1       )				,						(1, self.num_choices, 1)       )
					__lowercase			:  Tuple								=       tf.tile(tf.expand_dims(UpperCamelCase_				,						1       )				,						(1, self.num_choices, 1)       )
					__lowercase			:  str								=       tf.tile(tf.expand_dims(UpperCamelCase_				,						1       )				,						(1, self.num_choices, 1)       )
					__lowercase			:  Optional[Any]								=       {
					    '''input_ids''': multiple_choice_inputs_ids,
					    '''attention_mask''': multiple_choice_input_mask,
					    '''token_type_ids''': multiple_choice_token_type_ids,
					}
					__lowercase			:  List[str]								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.num_choices)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Union[str, Any]:
					__lowercase			:  str								=       self.num_labels
					__lowercase			:  Optional[Any]								=       TFFunnelForTokenClassification(config=UpperCamelCase_       )
					__lowercase			:  Dict								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  str								=       model(UpperCamelCase_       )
					self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.num_labels)       )
		def 	_lowerCamelCase							(						self				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						UpperCamelCase_				,						)						->       Optional[Any]:
					__lowercase			:  Any								=       TFFunnelForQuestionAnswering(config=UpperCamelCase_       )
					__lowercase			:  str								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					__lowercase			:  List[Any]								=       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 	_lowerCamelCase							(						self       )						->       Optional[Any]:
					__lowercase			:  Tuple								=       self.prepare_config_and_inputs()
					(
					    (
					    __lowercase
					)			,(
					    __lowercase
					)			,(
					    __lowercase
					)			,(
					    __lowercase
					)			,(
					    __lowercase
					)			,(
					    __lowercase
					)			,(
					    __lowercase
					)			,
					)			:  str								=       config_and_inputs
					__lowercase			:  List[Any]								=       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
					return config, inputs_dict


@require_tf
class   UpperCAmelCase_    (			snake_case   ,					snake_case   ,					unittest.TestCase		):
		UpperCamelCase							=(
		    (
		        TFFunnelModel,
		        TFFunnelForMaskedLM,
		        TFFunnelForPreTraining,
		        TFFunnelForQuestionAnswering,
		        TFFunnelForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		UpperCamelCase							=(
		    {
		        "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							=False
		UpperCamelCase							=False
		def 	_lowerCamelCase							(						self       )						->       int:
					__lowercase			:  Optional[int]								=       TFFunnelModelTester(self       )
					__lowercase			:  Any								=       ConfigTester(self				,						config_class=UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       Optional[Any]:
					self.config_tester.run_common_tests()
		def 	_lowerCamelCase							(						self       )						->       Any:
					__lowercase			:  List[str]								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       Dict:
					__lowercase			:  Any								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_pretraining(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       List[str]:
					__lowercase			:  Dict								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					__lowercase			:  Union[str, Any]								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       List[Any]:
					__lowercase			:  List[Any]								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*UpperCamelCase_       )


@require_tf
class   UpperCAmelCase_    (			snake_case   ,					unittest.TestCase		):
		UpperCamelCase							=(
		    (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
		)
		UpperCamelCase							=False
		UpperCamelCase							=False
		def 	_lowerCamelCase							(						self       )						->       str:
					__lowercase			:  List[str]								=       TFFunnelModelTester(self				,						base=UpperCamelCase_       )
					__lowercase			:  Optional[Any]								=       ConfigTester(self				,						config_class=UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       List[str]:
					self.config_tester.run_common_tests()
		def 	_lowerCamelCase							(						self       )						->       Optional[Any]:
					__lowercase			:  List[Any]								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_base_model(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       List[str]:
					__lowercase			:  int								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase_       )
		def 	_lowerCamelCase							(						self       )						->       Tuple:
					__lowercase			:  str								=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase_       )









import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


SCREAMING_SNAKE_CASE__     :     Union[str, Any]		    =       logging.get_logger(__name__)



class       a_       (     SCREAMING_SNAKE_CASE__     ):





	def __init__(							self      ,     SCREAMING_SNAKE_CASE      )  ->					Union[str, Any]:




				"""simple docstring"""



				super().__init__()
				SCREAMING_SNAKE_CASE_								=						nn.ModuleList(SCREAMING_SNAKE_CASE      )





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = True      ,     )  ->					Union[ControlNetOutput, Tuple]:




				"""simple docstring"""



				for i, (image, scale, controlnet) in enumerate(zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     self.nets      )      ):
							SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						controlnet(
							    SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     )

							# merge samples
							if i == 0:
										SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						down_samples, mid_sample
							else:
										SCREAMING_SNAKE_CASE_								=						[
										    samples_prev + samples_curr
										    for samples_prev, samples_curr in zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      )
										]
										mid_block_res_sample += mid_sample

				return down_block_res_samples, mid_block_res_sample





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = True      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = None      ,     )  ->					int:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						save_directory
				for controlnet in self.nets:
							controlnet.save_pretrained(
							    SCREAMING_SNAKE_CASE      ,     is_main_process=SCREAMING_SNAKE_CASE      ,     save_function=SCREAMING_SNAKE_CASE      ,     safe_serialization=SCREAMING_SNAKE_CASE      ,     variant=SCREAMING_SNAKE_CASE      ,     )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						model_path_to_save + f'_{idx}'





	@classmethod
	def 						A_(							cls      ,     SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )  ->					Tuple:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						[]

				# load controlnet and append to list until no controlnet directory exists anymore
				# first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
				# second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
				SCREAMING_SNAKE_CASE_								=						pretrained_model_path
				while os.path.isdir(SCREAMING_SNAKE_CASE      ):
							SCREAMING_SNAKE_CASE_								=						ControlNetModel.from_pretrained(SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )
							controlnets.append(SCREAMING_SNAKE_CASE      )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						pretrained_model_path + f'_{idx}'

				logger.info(f'{len(SCREAMING_SNAKE_CASE      )} controlnets loaded from {pretrained_model_path}.'      )

				if len(SCREAMING_SNAKE_CASE      ) == 0:
							raise ValueError(
							    f'No ControlNets found under {os.path.dirname(SCREAMING_SNAKE_CASE      )}. Expected at least {pretrained_model_path + "_0"}.'      )

				return cls(SCREAMING_SNAKE_CASE      )
