
'''simple docstring'''





import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_PRETRAINING_MAPPING,
						    ElectraForCausalLM,
						    ElectraForMaskedLM,
						    ElectraForMultipleChoice,
						    ElectraForPreTraining,
						    ElectraForQuestionAnswering,
						    ElectraForSequenceClassification,
						    ElectraForTokenClassification,
						    ElectraModel,
						)
						from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST

class     a :




				def __init__(				self						:					Any,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					str=13,     SCREAMING_SNAKE_CASE_						:					Tuple=7,     SCREAMING_SNAKE_CASE_						:					Optional[int]=True,     SCREAMING_SNAKE_CASE_						:					Dict=True,     SCREAMING_SNAKE_CASE_						:					List[Any]=True,     SCREAMING_SNAKE_CASE_						:					Union[str, Any]=True,     SCREAMING_SNAKE_CASE_						:					Any=99,     SCREAMING_SNAKE_CASE_						:					Tuple=32,     SCREAMING_SNAKE_CASE_						:					Tuple=5,     SCREAMING_SNAKE_CASE_						:					List[Any]=4,     SCREAMING_SNAKE_CASE_						:					Dict=37,     SCREAMING_SNAKE_CASE_						:					Union[str, Any]="gelu",     SCREAMING_SNAKE_CASE_						:					Optional[int]=0.1,     SCREAMING_SNAKE_CASE_						:					int=0.1,     SCREAMING_SNAKE_CASE_						:					Optional[int]=5_12,     SCREAMING_SNAKE_CASE_						:					List[str]=16,     SCREAMING_SNAKE_CASE_						:					int=2,     SCREAMING_SNAKE_CASE_						:					Dict=0.02,     SCREAMING_SNAKE_CASE_						:					Optional[int]=3,     SCREAMING_SNAKE_CASE_						:					Optional[int]=4,     SCREAMING_SNAKE_CASE_						:					Optional[int]=None,     ):
										snake_case							:  Dict			=     parent
										snake_case							:  int			=     batch_size
										snake_case							:  int			=     seq_length
										snake_case							:  List[str]			=     is_training
										snake_case							:  List[str]			=     use_input_mask
										snake_case							:  Tuple			=     use_token_type_ids
										snake_case							:  Optional[int]			=     use_labels
										snake_case							:  Tuple			=     vocab_size
										snake_case							:  Optional[Any]			=     hidden_size
										snake_case							:  Dict			=     num_hidden_layers
										snake_case							:  Tuple			=     num_attention_heads
										snake_case							:  Any			=     intermediate_size
										snake_case							:  Union[str, Any]			=     hidden_act
										snake_case							:  List[str]			=     hidden_dropout_prob
										snake_case							:  Any			=     attention_probs_dropout_prob
										snake_case							:  Dict			=     max_position_embeddings
										snake_case							:  int			=     type_vocab_size
										snake_case							:  Tuple			=     type_sequence_label_size
										snake_case							:  Optional[int]			=     initializer_range
										snake_case							:  Optional[Any]			=     num_labels
										snake_case							:  int			=     num_choices
										snake_case							:  Tuple			=     scope




				def 						__snake_case				(				self						:					List[Any]   ):
										snake_case							:  Tuple			=     ids_tensor([self.batch_size, self.seq_length],     self.vocab_size   )

										snake_case							:  Union[str, Any]			=     None
										if self.use_input_mask:
																snake_case							:  Dict			=     random_attention_mask([self.batch_size, self.seq_length]   )

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

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

										snake_case							:  Union[str, Any]			=     self.get_config()

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    fake_token_labels,
										)




				def 						__snake_case				(				self						:					Union[str, Any]   ):
										return ElectraConfig(
										    vocab_size=self.vocab_size,     hidden_size=self.hidden_size,     num_hidden_layers=self.num_hidden_layers,     num_attention_heads=self.num_attention_heads,     intermediate_size=self.intermediate_size,     hidden_act=self.hidden_act,     hidden_dropout_prob=self.hidden_dropout_prob,     attention_probs_dropout_prob=self.attention_probs_dropout_prob,     max_position_embeddings=self.max_position_embeddings,     type_vocab_size=self.type_vocab_size,     is_decoder=SCREAMING_SNAKE_CASE_,     initializer_range=self.initializer_range,     )




				def 						__snake_case				(				self						:					Any   ):
										(
										    (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       
										)							:  Dict			=     self.prepare_config_and_inputs()

										snake_case							:  List[str]			=     True
										snake_case							:  Optional[Any]			=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
										snake_case							:  Optional[int]			=     ids_tensor([self.batch_size, self.seq_length],     vocab_size=2   )

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)




				def 						__snake_case				(				self						:					Tuple,     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					List[str],     SCREAMING_SNAKE_CASE_						:					Optional[int],     ):
										snake_case							:  Any			=     ElectraModel(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  Optional[int]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_   )
										snake_case							:  Dict			=     model(SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_   )
										snake_case							:  int			=     model(SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.last_hidden_state.shape,     (self.batch_size, self.seq_length, self.hidden_size)   )




				def 						__snake_case				(				self						:					List[str],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Dict,     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     ):
										snake_case							:  List[Any]			=     True
										snake_case							:  int			=     ElectraModel(SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  Union[str, Any]			=     model(
										    SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     encoder_hidden_states=SCREAMING_SNAKE_CASE_,     encoder_attention_mask=SCREAMING_SNAKE_CASE_,     )
										snake_case							:  Optional[int]			=     model(
										    SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     encoder_hidden_states=SCREAMING_SNAKE_CASE_,     )
										snake_case							:  List[Any]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.last_hidden_state.shape,     (self.batch_size, self.seq_length, self.hidden_size)   )




				def 						__snake_case				(				self						:					int,     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					Dict,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Tuple,     ):
										snake_case							:  Dict			=     ElectraForMaskedLM(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  Union[str, Any]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.seq_length, self.vocab_size)   )




				def 						__snake_case				(				self						:					Any,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Optional[int],     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Dict,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     ):
										snake_case							:  int			=     ElectraForCausalLM(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  Optional[int]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.seq_length, self.vocab_size)   )




				def 						__snake_case				(				self						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					List[str],     SCREAMING_SNAKE_CASE_						:					Optional[int],     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Dict,     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					Any,     ):
										snake_case							:  Dict			=     self.num_labels
										snake_case							:  Dict			=     ElectraForTokenClassification(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  str			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.seq_length, self.num_labels)   )




				def 						__snake_case				(				self						:					str,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					Dict,     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					Optional[Any],     ):
										snake_case							:  Dict			=     self.num_labels
										snake_case							:  int			=     ElectraForPreTraining(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  List[Any]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.seq_length)   )




				def 						__snake_case				(				self						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					List[str],     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Optional[int],     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					List[str],     ):
										snake_case							:  Tuple			=     self.num_labels
										snake_case							:  List[str]			=     ElectraForSequenceClassification(SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  Optional[Any]			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_   )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.num_labels)   )




				def 						__snake_case				(				self						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					Optional[int],     SCREAMING_SNAKE_CASE_						:					Any,     SCREAMING_SNAKE_CASE_						:					List[str],     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					List[str],     SCREAMING_SNAKE_CASE_						:					int,     ):
										snake_case							:  int			=     ElectraForQuestionAnswering(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  List[Any]			=     model(
										    SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     start_positions=SCREAMING_SNAKE_CASE_,     end_positions=SCREAMING_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						:					str,     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					List[Any],     SCREAMING_SNAKE_CASE_						:					Optional[Any],     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					str,     SCREAMING_SNAKE_CASE_						:					int,     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					Optional[int],     ):
										snake_case							:  str			=     self.num_choices
										snake_case							:  int			=     ElectraForMultipleChoice(config=SCREAMING_SNAKE_CASE_   )
										model.to(SCREAMING_SNAKE_CASE_   )
										model.eval()
										snake_case							:  List[str]			=     input_ids.unsqueeze(1   ).expand(-1,     self.num_choices,     -1   ).contiguous()
										snake_case							:  int			=     token_type_ids.unsqueeze(1   ).expand(-1,     self.num_choices,     -1   ).contiguous()
										snake_case							:  Tuple			=     input_mask.unsqueeze(1   ).expand(-1,     self.num_choices,     -1   ).contiguous()
										snake_case							:  str			=     model(
										    SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_,     token_type_ids=SCREAMING_SNAKE_CASE_,     labels=SCREAMING_SNAKE_CASE_,     )
										self.parent.assertEqual(result.logits.shape,     (self.batch_size, self.num_choices)   )




				def 						__snake_case				(				self						:					Dict   ):
										snake_case							:  Union[str, Any]			=     self.prepare_config_and_inputs()
										(
										    (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       (
										    snake_case
										),       
										)							:  Union[str, Any]			=     config_and_inputs
										snake_case							:  Dict			=     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
										return config, inputs_dict




@require_torch
class     a (					__magic_name__					,__magic_name__					,unittest.TestCase     ):
				_snake_case		     =			(
				    (
				        ElectraModel,
				        ElectraForPreTraining,
				        ElectraForMaskedLM,
				        ElectraForCausalLM,
				        ElectraForMultipleChoice,
				        ElectraForTokenClassification,
				        ElectraForSequenceClassification,
				        ElectraForQuestionAnswering,
				    )
				    if is_torch_available()
				    else ()
				)
				_snake_case		     =			(
				    {
				        '''feature-extraction''': ElectraModel,
				        '''fill-mask''': ElectraForMaskedLM,
				        '''question-answering''': ElectraForQuestionAnswering,
				        '''text-classification''': ElectraForSequenceClassification,
				        '''text-generation''': ElectraForCausalLM,
				        '''token-classification''': ElectraForTokenClassification,
				        '''zero-shot''': ElectraForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_snake_case		     =			True




				def 						__snake_case				(				self						:					str,     SCREAMING_SNAKE_CASE_						:					Union[str, Any],     SCREAMING_SNAKE_CASE_						:					Tuple,     SCREAMING_SNAKE_CASE_						:					int=False   ):
										snake_case							:  Dict			=     super()._prepare_for_class(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     return_labels=SCREAMING_SNAKE_CASE_   )

										if return_labels:
																if model_class in get_values(SCREAMING_SNAKE_CASE_   ):
																						snake_case							:  int			=     torch.zeros(
																						    (self.model_tester.batch_size, self.model_tester.seq_length),     dtype=torch.long,     device=SCREAMING_SNAKE_CASE_   )
										return inputs_dict




				def 						__snake_case				(				self						:					Dict   ):
										snake_case							:  Dict			=     ElectraModelTester(self   )
										snake_case							:  int			=     ConfigTester(self,     config_class=SCREAMING_SNAKE_CASE_,     hidden_size=37   )




				def 						__snake_case				(				self						:					Dict   ):
										self.config_tester.run_common_tests()




				def 						__snake_case				(				self						:					Tuple   ):
										snake_case							:  Optional[int]			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_model(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					str   ):
										snake_case							:  int			=     self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_electra_model_as_decoder(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					List[str]   ):
										snake_case							:  Tuple			=     self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																snake_case							:  List[Any]			=     type
																self.model_tester.create_and_check_electra_model(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					str   ):
										snake_case							:  Dict			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_masked_lm(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					Optional[int]   ):
										snake_case							:  Optional[Any]			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_token_classification(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					List[str]   ):
										snake_case							:  Any			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_pretraining(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					int   ):
										snake_case							:  List[str]			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_sequence_classification(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					Optional[int]   ):
										snake_case							:  str			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_question_answering(*SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					Any   ):
										snake_case							:  Any			=     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_multiple_choice(*SCREAMING_SNAKE_CASE_   )




				@slow
				def 						__snake_case				(				self						:					Tuple   ):
										for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																snake_case							:  Any			=     ElectraModel.from_pretrained(SCREAMING_SNAKE_CASE_   )
																self.assertIsNotNone(SCREAMING_SNAKE_CASE_   )




				def 						__snake_case				(				self						:					Optional[int]   ):
										snake_case							:  Optional[int]			=     self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_electra_for_causal_lm(*SCREAMING_SNAKE_CASE_   )




@require_torch
class     a (					unittest.TestCase     ):




				@slow
				def 						__snake_case				(				self						:					Optional[int]   ):
										snake_case							:  Optional[int]			=     ElectraModel.from_pretrained('''google/electra-small-discriminator'''   )
										snake_case							:  Optional[int]			=     torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]   )
										snake_case							:  List[str]			=     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
										snake_case							:  str			=     model(SCREAMING_SNAKE_CASE_,     attention_mask=SCREAMING_SNAKE_CASE_   )[0]
										snake_case							:  int			=     torch.Size((1, 11, 2_56)   )
										self.assertEqual(output.shape,     SCREAMING_SNAKE_CASE_   )
										snake_case							:  Union[str, Any]			=     torch.tensor(
										    [[[0.4471, 0.6821, -0.3265], [0.4627, 0.5255, -0.3668], [0.4532, 0.3313, -0.4344]]]   )

										self.assertTrue(torch.allclose(output[:, 1:4, 1:4],     SCREAMING_SNAKE_CASE_,     atol=1e-4   )   )









import argparse
import json

import torch

from diffusers import AutoencoderKL, DDPMPipeline, DDPMScheduler, UNetaDModel, VQModel







def 						_UpperCamelCase (     UpperCamelCase_							:       Tuple     ,							UpperCamelCase_							:       Optional[Any]=1		)	-> Optional[int]:

						"""simple docstring"""



						if n_shave_prefix_segments >= 0:
												return ".".join(path.split('.'		)[n_shave_prefix_segments:]		)
						else:
												return ".".join(path.split('.'		)[:n_shave_prefix_segments]		)







def 						_UpperCamelCase (     UpperCamelCase_							:       Union[str, Any]     ,							UpperCamelCase_							:       int=0		)	-> Union[str, Any]:

						"""simple docstring"""



						lowerCAmelCase__         =    []
						for old_item in old_list:
												lowerCAmelCase__         =    old_item
												lowerCAmelCase__         =    new_item.replace('block.'     ,							'resnets.'		)
												lowerCAmelCase__         =    new_item.replace('conv_shorcut'     ,							'conv1'		)
												lowerCAmelCase__         =    new_item.replace('in_shortcut'     ,							'conv_shortcut'		)
												lowerCAmelCase__         =    new_item.replace('temb_proj'     ,							'time_emb_proj'		)

												lowerCAmelCase__         =    shave_segments(UpperCamelCase_     ,							n_shave_prefix_segments=UpperCamelCase_		)

												mapping.append({'old': old_item, 'new': new_item}		)

						return mapping







def 						_UpperCamelCase (     UpperCamelCase_							:       str     ,							UpperCamelCase_							:       Union[str, Any]=0     ,							UpperCamelCase_							:       Dict=False		)	-> List[Any]:

						"""simple docstring"""



						lowerCAmelCase__         =    []
						for old_item in old_list:
												lowerCAmelCase__         =    old_item

												# In `model.mid`, the layer is called `attn`.
												if not in_mid:
																		lowerCAmelCase__         =    new_item.replace('attn'     ,							'attentions'		)
												lowerCAmelCase__         =    new_item.replace('.k.'     ,							'.key.'		)
												lowerCAmelCase__         =    new_item.replace('.v.'     ,							'.value.'		)
												lowerCAmelCase__         =    new_item.replace('.q.'     ,							'.query.'		)

												lowerCAmelCase__         =    new_item.replace('proj_out'     ,							'proj_attn'		)
												lowerCAmelCase__         =    new_item.replace('norm'     ,							'group_norm'		)

												lowerCAmelCase__         =    shave_segments(UpperCamelCase_     ,							n_shave_prefix_segments=UpperCamelCase_		)
												mapping.append({'old': old_item, 'new': new_item}		)

						return mapping







def 						_UpperCamelCase (     UpperCamelCase_							:       int     ,							UpperCamelCase_							:       Dict     ,							UpperCamelCase_							:       Union[str, Any]     ,							UpperCamelCase_							:       Any=None     ,							UpperCamelCase_							:       str=None     ,							UpperCamelCase_							:       str=None		)	-> Dict:

						"""simple docstring"""



						assert isinstance(UpperCamelCase_     ,							UpperCamelCase_		), "Paths should be a list of dicts containing 'old' and 'new' keys."

						if attention_paths_to_split is not None:
												if config is None:
																		raise ValueError('Please specify the config if setting \'attention_paths_to_split\' to \'True\'.'		)

												for path, path_map in attention_paths_to_split.items():
																		lowerCAmelCase__         =    old_checkpoint[path]
																		lowerCAmelCase__         =    old_tensor.shape[0] // 3

																		lowerCAmelCase__         =    (-1, channels) if len(old_tensor.shape		) == 3 else (-1)

																		lowerCAmelCase__         =    old_tensor.shape[0] // config.get('num_head_channels'     ,							1		) // 3

																		lowerCAmelCase__         =    old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]		)
																		lowerCAmelCase__    ,							lowerCAmelCase__    ,							lowerCAmelCase__         =    old_tensor.split(channels // num_heads     ,							dim=1		)

																		lowerCAmelCase__         =    query.reshape(UpperCamelCase_		).squeeze()
																		lowerCAmelCase__         =    key.reshape(UpperCamelCase_		).squeeze()
																		lowerCAmelCase__         =    value.reshape(UpperCamelCase_		).squeeze()

						for path in paths:
												lowerCAmelCase__         =    path['new']

												if attention_paths_to_split is not None and new_path in attention_paths_to_split:
																		continue

												lowerCAmelCase__         =    new_path.replace('down.'     ,							'down_blocks.'		)
												lowerCAmelCase__         =    new_path.replace('up.'     ,							'up_blocks.'		)

												if additional_replacements is not None:
																		for replacement in additional_replacements:
																								lowerCAmelCase__         =    new_path.replace(replacement['old']     ,							replacement['new']		)

												if "attentions" in new_path:
																		lowerCAmelCase__         =    old_checkpoint[path['old']].squeeze()
												else:
																		lowerCAmelCase__         =    old_checkpoint[path['old']]







def 						_UpperCamelCase (     UpperCamelCase_							:       str     ,							UpperCamelCase_							:       Dict		)	-> Any:

						"""simple docstring"""



						lowerCAmelCase__         =    {}

						lowerCAmelCase__         =    checkpoint['temb.dense.0.weight']
						lowerCAmelCase__         =    checkpoint['temb.dense.0.bias']
						lowerCAmelCase__         =    checkpoint['temb.dense.1.weight']
						lowerCAmelCase__         =    checkpoint['temb.dense.1.bias']

						lowerCAmelCase__         =    checkpoint['norm_out.weight']
						lowerCAmelCase__         =    checkpoint['norm_out.bias']

						lowerCAmelCase__         =    checkpoint['conv_in.weight']
						lowerCAmelCase__         =    checkpoint['conv_in.bias']
						lowerCAmelCase__         =    checkpoint['conv_out.weight']
						lowerCAmelCase__         =    checkpoint['conv_out.bias']

						lowerCAmelCase__         =    len({'.'.join(layer.split('.'		)[:2]		) for layer in checkpoint if 'down' in layer}		)
						lowerCAmelCase__         =    {
						    layer_id: [key for key in checkpoint if F"down.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
						}

						lowerCAmelCase__         =    len({'.'.join(layer.split('.'		)[:2]		) for layer in checkpoint if 'up' in layer}		)
						lowerCAmelCase__         =    {layer_id: [key for key in checkpoint if F"up.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)}

						for i in range(UpperCamelCase_		):
												lowerCAmelCase__         =    (i - 1) // (config['layers_per_block'] + 1)

												if any('downsample' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    checkpoint[
																		    F"down.{i}.downsample.op.weight"
																		]
																		lowerCAmelCase__         =    checkpoint[F"down.{i}.downsample.op.bias"]
												#            new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.weight'] = checkpoint[f'down.{i}.downsample.conv.weight']
												#            new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.bias'] = checkpoint[f'down.{i}.downsample.conv.bias']

												if any('block' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							2		).split('.'		)[:2]		) for layer in down_blocks[i] if 'block' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in down_blocks[i] if F"block.{layer_id}" in key]
																		    for layer_id in range(UpperCamelCase_		)
																		}

																		if num_blocks > 0:
																								for j in range(config['layers_per_block']		):
																														lowerCAmelCase__         =    renew_resnet_paths(blocks[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_		)

												if any('attn' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							2		).split('.'		)[:2]		) for layer in down_blocks[i] if 'attn' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in down_blocks[i] if F"attn.{layer_id}" in key]
																		    for layer_id in range(UpperCamelCase_		)
																		}

																		if num_attn > 0:
																								for j in range(config['layers_per_block']		):
																														lowerCAmelCase__         =    renew_attention_paths(attns[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							config=UpperCamelCase_		)

						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.block_1' in key]
						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.block_2' in key]
						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.attn_1' in key]

						# Mid new 2
						lowerCAmelCase__         =    renew_resnet_paths(UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'block_1', 'new': 'resnets.0'}]     ,							)

						lowerCAmelCase__         =    renew_resnet_paths(UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'block_2', 'new': 'resnets.1'}]     ,							)

						lowerCAmelCase__         =    renew_attention_paths(UpperCamelCase_     ,							in_mid=UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'attn_1', 'new': 'attentions.0'}]     ,							)

						for i in range(UpperCamelCase_		):
												lowerCAmelCase__         =    num_up_blocks - 1 - i

												if any('upsample' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    checkpoint[
																		    F"up.{i}.upsample.conv.weight"
																		]
																		lowerCAmelCase__         =    checkpoint[F"up.{i}.upsample.conv.bias"]

												if any('block' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							2		).split('.'		)[:2]		) for layer in up_blocks[i] if 'block' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in up_blocks[i] if F"block.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
																		}

																		if num_blocks > 0:
																								for j in range(config['layers_per_block'] + 1		):
																														lowerCAmelCase__         =    {'old': F"up_blocks.{i}", 'new': F"up_blocks.{block_id}"}
																														lowerCAmelCase__         =    renew_resnet_paths(blocks[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[replace_indices]		)

												if any('attn' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							2		).split('.'		)[:2]		) for layer in up_blocks[i] if 'attn' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in up_blocks[i] if F"attn.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
																		}

																		if num_attn > 0:
																								for j in range(config['layers_per_block'] + 1		):
																														lowerCAmelCase__         =    {'old': F"up_blocks.{i}", 'new': F"up_blocks.{block_id}"}
																														lowerCAmelCase__         =    renew_attention_paths(attns[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[replace_indices]		)

						lowerCAmelCase__         =    {k.replace('mid_new_2'     ,							'mid_block'		): v for k, v in new_checkpoint.items()}
						return new_checkpoint







def 						_UpperCamelCase (     UpperCamelCase_							:       Tuple     ,							UpperCamelCase_							:       int		)	-> Optional[int]:

						"""simple docstring"""



						lowerCAmelCase__         =    {}

						lowerCAmelCase__         =    checkpoint['encoder.norm_out.weight']
						lowerCAmelCase__         =    checkpoint['encoder.norm_out.bias']

						lowerCAmelCase__         =    checkpoint['encoder.conv_in.weight']
						lowerCAmelCase__         =    checkpoint['encoder.conv_in.bias']
						lowerCAmelCase__         =    checkpoint['encoder.conv_out.weight']
						lowerCAmelCase__         =    checkpoint['encoder.conv_out.bias']

						lowerCAmelCase__         =    checkpoint['decoder.norm_out.weight']
						lowerCAmelCase__         =    checkpoint['decoder.norm_out.bias']

						lowerCAmelCase__         =    checkpoint['decoder.conv_in.weight']
						lowerCAmelCase__         =    checkpoint['decoder.conv_in.bias']
						lowerCAmelCase__         =    checkpoint['decoder.conv_out.weight']
						lowerCAmelCase__         =    checkpoint['decoder.conv_out.bias']

						lowerCAmelCase__         =    len({'.'.join(layer.split('.'		)[:3]		) for layer in checkpoint if 'down' in layer}		)
						lowerCAmelCase__         =    {
						    layer_id: [key for key in checkpoint if F"down.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
						}

						lowerCAmelCase__         =    len({'.'.join(layer.split('.'		)[:3]		) for layer in checkpoint if 'up' in layer}		)
						lowerCAmelCase__         =    {layer_id: [key for key in checkpoint if F"up.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)}

						for i in range(UpperCamelCase_		):
												lowerCAmelCase__         =    (i - 1) // (config['layers_per_block'] + 1)

												if any('downsample' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    checkpoint[
																		    F"encoder.down.{i}.downsample.conv.weight"
																		]
																		lowerCAmelCase__         =    checkpoint[
																		    F"encoder.down.{i}.downsample.conv.bias"
																		]

												if any('block' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							3		).split('.'		)[:3]		) for layer in down_blocks[i] if 'block' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in down_blocks[i] if F"block.{layer_id}" in key]
																		    for layer_id in range(UpperCamelCase_		)
																		}

																		if num_blocks > 0:
																								for j in range(config['layers_per_block']		):
																														lowerCAmelCase__         =    renew_resnet_paths(blocks[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_		)

												if any('attn' in layer for layer in down_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							3		).split('.'		)[:3]		) for layer in down_blocks[i] if 'attn' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in down_blocks[i] if F"attn.{layer_id}" in key]
																		    for layer_id in range(UpperCamelCase_		)
																		}

																		if num_attn > 0:
																								for j in range(config['layers_per_block']		):
																														lowerCAmelCase__         =    renew_attention_paths(attns[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							config=UpperCamelCase_		)

						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.block_1' in key]
						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.block_2' in key]
						lowerCAmelCase__         =    [key for key in checkpoint if 'mid.attn_1' in key]

						# Mid new 2
						lowerCAmelCase__         =    renew_resnet_paths(UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'block_1', 'new': 'resnets.0'}]     ,							)

						lowerCAmelCase__         =    renew_resnet_paths(UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'block_2', 'new': 'resnets.1'}]     ,							)

						lowerCAmelCase__         =    renew_attention_paths(UpperCamelCase_     ,							in_mid=UpperCamelCase_		)
						assign_to_checkpoint(
						    UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[{'old': 'mid.', 'new': 'mid_new_2.'}, {'old': 'attn_1', 'new': 'attentions.0'}]     ,							)

						for i in range(UpperCamelCase_		):
												lowerCAmelCase__         =    num_up_blocks - 1 - i

												if any('upsample' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    checkpoint[
																		    F"decoder.up.{i}.upsample.conv.weight"
																		]
																		lowerCAmelCase__         =    checkpoint[
																		    F"decoder.up.{i}.upsample.conv.bias"
																		]

												if any('block' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							3		).split('.'		)[:3]		) for layer in up_blocks[i] if 'block' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in up_blocks[i] if F"block.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
																		}

																		if num_blocks > 0:
																								for j in range(config['layers_per_block'] + 1		):
																														lowerCAmelCase__         =    {'old': F"up_blocks.{i}", 'new': F"up_blocks.{block_id}"}
																														lowerCAmelCase__         =    renew_resnet_paths(blocks[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[replace_indices]		)

												if any('attn' in layer for layer in up_blocks[i]		):
																		lowerCAmelCase__         =    len(
																		    {'.'.join(shave_segments(UpperCamelCase_     ,							3		).split('.'		)[:3]		) for layer in up_blocks[i] if 'attn' in layer}		)
																		lowerCAmelCase__         =    {
																		    layer_id: [key for key in up_blocks[i] if F"attn.{layer_id}" in key] for layer_id in range(UpperCamelCase_		)
																		}

																		if num_attn > 0:
																								for j in range(config['layers_per_block'] + 1		):
																														lowerCAmelCase__         =    {'old': F"up_blocks.{i}", 'new': F"up_blocks.{block_id}"}
																														lowerCAmelCase__         =    renew_attention_paths(attns[j]		)
																														assign_to_checkpoint(UpperCamelCase_     ,							UpperCamelCase_     ,							UpperCamelCase_     ,							additional_replacements=[replace_indices]		)

						lowerCAmelCase__         =    {k.replace('mid_new_2'     ,							'mid_block'		): v for k, v in new_checkpoint.items()}
						lowerCAmelCase__         =    checkpoint['quant_conv.weight']
						lowerCAmelCase__         =    checkpoint['quant_conv.bias']
						if "quantize.embedding.weight" in checkpoint:
												lowerCAmelCase__         =    checkpoint['quantize.embedding.weight']
						lowerCAmelCase__         =    checkpoint['post_quant_conv.weight']
						lowerCAmelCase__         =    checkpoint['post_quant_conv.bias']

						return new_checkpoint


if __name__ == "__main__":
					__snake_case     :     Tuple       =  argparse.ArgumentParser()

					parser.add_argument(
					    """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
					)

					parser.add_argument(
					    """--config_file""",
					    default=None,
					    type=str,
					    required=True,
					    help="""The config json file corresponding to the architecture.""",
					)

					parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")

					__snake_case     :     str       =  parser.parse_args()
					__snake_case     :     Any       =  torch.load(args.checkpoint_path)

					with open(args.config_file) as f:
										__snake_case     :     Optional[int]       =  json.loads(f.read())

					# unet case
					__snake_case     :     Any       =  {key.split(""".""")[0] for key in checkpoint.keys()}
					if "encoder" in key_prefix_set and "decoder" in key_prefix_set:
										__snake_case     :     Optional[int]       =  convert_vq_autoenc_checkpoint(checkpoint, config)
					else:
										__snake_case     :     List[str]       =  convert_ddpm_checkpoint(checkpoint, config)

					if "ddpm" in config:
										del config["ddpm"]

					if config["_class_name"] == "VQModel":
										__snake_case     :     List[str]       =  VQModel(**config)
										model.load_state_dict(converted_checkpoint)
										model.save_pretrained(args.dump_path)
					elif config["_class_name"] == "AutoencoderKL":
										__snake_case     :     List[Any]       =  AutoencoderKL(**config)
										model.load_state_dict(converted_checkpoint)
										model.save_pretrained(args.dump_path)
					else:
										__snake_case     :     Any       =  UNetaDModel(**config)
										model.load_state_dict(converted_checkpoint)

										__snake_case     :     List[str]       =  DDPMScheduler.from_config("""/""".join(args.checkpoint_path.split("""/""")[:-1]))

										__snake_case     :     Tuple       =  DDPMPipeline(unet=model, scheduler=scheduler)
										pipe.save_pretrained(args.dump_path)





