
import unittest

from transformers import AlbertConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    MODEL_FOR_PRETRAINING_MAPPING,
					    AlbertForMaskedLM,
					    AlbertForMultipleChoice,
					    AlbertForPreTraining,
					    AlbertForQuestionAnswering,
					    AlbertForSequenceClassification,
					    AlbertForTokenClassification,
					    AlbertModel,
					)
					from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST





class 							UpperCAmelCase      :




	'''simple docstring'''





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






	def 				UpperCamelCase_  (       self					: Union[str, Any] ):
					__A   			= ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size )

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

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

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

					__A   			= self.get_config()

					return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






	def 				UpperCamelCase_  (       self					: Tuple ):
					return AlbertConfig(
					    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							,initializer_range=self.initializer_range							,num_hidden_groups=self.num_hidden_groups							,)






	def 				UpperCamelCase_  (       self					: str							,A					: Union[str, Any]							,A					: Union[str, Any]							,A					: Any							,A					: str							,A					: Tuple							,A					: List[Any]							,A					: str ):
					__A   			= AlbertModel(config=A )
					model.to(A )
					model.eval()
					__A   			= model(A							,attention_mask=A							,token_type_ids=A )
					__A   			= model(A							,token_type_ids=A )
					__A   			= model(A )
					self.parent.assertEqual(result.last_hidden_state.shape							,(self.batch_size, self.seq_length, self.hidden_size) )
					self.parent.assertEqual(result.pooler_output.shape							,(self.batch_size, self.hidden_size) )






	def 				UpperCamelCase_  (       self					: Dict							,A					: Union[str, Any]							,A					: Dict							,A					: Optional[int]							,A					: int							,A					: Tuple							,A					: Optional[int]							,A					: str ):
					__A   			= AlbertForPreTraining(config=A )
					model.to(A )
					model.eval()
					__A   			= model(
					    A							,attention_mask=A							,token_type_ids=A							,labels=A							,sentence_order_label=A							,)
					self.parent.assertEqual(result.prediction_logits.shape							,(self.batch_size, self.seq_length, self.vocab_size) )
					self.parent.assertEqual(result.sop_logits.shape							,(self.batch_size, config.num_labels) )






	def 				UpperCamelCase_  (       self					: Tuple							,A					: Union[str, Any]							,A					: Dict							,A					: Dict							,A					: Union[str, Any]							,A					: List[str]							,A					: List[str]							,A					: List[str] ):
					__A   			= AlbertForMaskedLM(config=A )
					model.to(A )
					model.eval()
					__A   			= model(A							,attention_mask=A							,token_type_ids=A							,labels=A )
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size) )






	def 				UpperCamelCase_  (       self					: List[Any]							,A					: Union[str, Any]							,A					: Any							,A					: Union[str, Any]							,A					: Any							,A					: Tuple							,A					: str							,A					: Optional[int] ):
					__A   			= AlbertForQuestionAnswering(config=A )
					model.to(A )
					model.eval()
					__A   			= model(
					    A							,attention_mask=A							,token_type_ids=A							,start_positions=A							,end_positions=A							,)
					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 				UpperCamelCase_  (       self					: Any							,A					: Dict							,A					: Dict							,A					: Dict							,A					: Any							,A					: Optional[Any]							,A					: Any							,A					: str ):
					__A   			= self.num_labels
					__A   			= AlbertForSequenceClassification(A )
					model.to(A )
					model.eval()
					__A   			= model(A							,attention_mask=A							,token_type_ids=A							,labels=A )
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels) )






	def 				UpperCamelCase_  (       self					: Any							,A					: int							,A					: Dict							,A					: Dict							,A					: Union[str, Any]							,A					: List[str]							,A					: str							,A					: int ):
					__A   			= self.num_labels
					__A   			= AlbertForTokenClassification(config=A )
					model.to(A )
					model.eval()
					__A   			= model(A							,attention_mask=A							,token_type_ids=A							,labels=A )
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.num_labels) )






	def 				UpperCamelCase_  (       self					: Any							,A					: Union[str, Any]							,A					: Optional[Any]							,A					: int							,A					: Dict							,A					: int							,A					: List[Any]							,A					: Any ):
					__A   			= self.num_choices
					__A   			= AlbertForMultipleChoice(config=A )
					model.to(A )
					model.eval()
					__A   			= input_ids.unsqueeze(1 ).expand(-1							,self.num_choices							,-1 ).contiguous()
					__A   			= token_type_ids.unsqueeze(1 ).expand(-1							,self.num_choices							,-1 ).contiguous()
					__A   			= input_mask.unsqueeze(1 ).expand(-1							,self.num_choices							,-1 ).contiguous()
					__A   			= model(
					    A							,attention_mask=A							,token_type_ids=A							,labels=A							,)
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_choices) )







	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= self.prepare_config_and_inputs()
					(
					    (
					    __A
					)			,							(
					    __A
					)			,							(
					    __A
					)			,							(
					    __A
					)			,							(
					    __A
					)			,							(
					    __A
					)			,							(
					    __A
					)			,							
					)   			= config_and_inputs
					__A   			= {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
					return config, inputs_dict





@require_torch
class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			unittest.TestCase  ):




	'''simple docstring'''





	snake_case_					    =     (
	    (
	        AlbertModel,
	        AlbertForPreTraining,
	        AlbertForMaskedLM,
	        AlbertForMultipleChoice,
	        AlbertForSequenceClassification,
	        AlbertForTokenClassification,
	        AlbertForQuestionAnswering,
	    )
	    if is_torch_available()
	    else ()
	)
	snake_case_					    =     (
	    {
	        "feature-extraction": AlbertModel,
	        "fill-mask": AlbertForMaskedLM,
	        "question-answering": AlbertForQuestionAnswering,
	        "text-classification": AlbertForSequenceClassification,
	        "token-classification": AlbertForTokenClassification,
	        "zero-shot": AlbertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	snake_case_					    =     True






	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: List[str]							,A					: Dict							,A					: str=False ):
					__A   			= super()._prepare_for_class(A							,A							,return_labels=A )

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






	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= AlbertModelTester(self )
					__A   			= ConfigTester(self							,config_class=A							,hidden_size=37 )






	def 				UpperCamelCase_  (       self					: Optional[Any] ):
					self.config_tester.run_common_tests()






	def 				UpperCamelCase_  (       self					: Dict ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*A )






	def 				UpperCamelCase_  (       self					: List[Any] ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_pretraining(*A )






	def 				UpperCamelCase_  (       self					: Union[str, Any] ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*A )






	def 				UpperCamelCase_  (       self					: Any ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*A )






	def 				UpperCamelCase_  (       self					: str ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*A )






	def 				UpperCamelCase_  (       self					: List[Any] ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*A )






	def 				UpperCamelCase_  (       self					: Optional[Any] ):
					__A   			= self.model_tester.prepare_config_and_inputs()
					for type in ["absolute", "relative_key", "relative_key_query"]:
									__A   			= type
									self.model_tester.create_and_check_model(*A )







	@slow
	def 				UpperCamelCase_  (       self					: Dict ):
					for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__A   			= AlbertModel.from_pretrained(A )
									self.assertIsNotNone(A )





@require_torch
class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	@slow
	def 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= AlbertModel.from_pretrained("albert-base-v2" )
					__A   			= torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]] )
					__A   			= torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
					with torch.no_grad():
									__A   			= model(A							,attention_mask=A )[0]
					__A   			= torch.Size((1, 11, 7_68) )
					self.assertEqual(output.shape							,A )
					__A   			= torch.tensor(
					    [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]] )

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


import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


SCREAMING_SNAKE_CASE							:Union[str, Any]							 =				False
SCREAMING_SNAKE_CASE							:Any							 =				True
SCREAMING_SNAKE_CASE							:Tuple							 =				False


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

					parser.add_argument(
					    '--repo_path',
					    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.')

					SCREAMING_SNAKE_CASE							:Union[str, Any]							 =				parser.parse_args()

					SCREAMING_SNAKE_CASE							:Dict							 =				{
					    'image_size': 'sample_size',
					    'num_res_blocks': 'layers_per_block',
					    'block_channels': 'block_out_channels',
					    'down_blocks': 'down_block_types',
					    'up_blocks': 'up_block_types',
					    'downscale_freq_shift': 'freq_shift',
					    'resnet_num_groups': 'norm_num_groups',
					    'resnet_act_fn': 'act_fn',
					    'resnet_eps': 'norm_eps',
					    'num_head_channels': 'attention_head_dim',
					}

					SCREAMING_SNAKE_CASE							:Optional[int]							 =				{
					    'time_steps': 'time_proj',
					    'mid': 'mid_block',
					    'downsample_blocks': 'down_blocks',
					    'upsample_blocks': 'up_blocks',
					}

					SCREAMING_SNAKE_CASE							:int							 =				'' if has_file(args.repo_path, 'config.json') else 'unet'

					with open(os.path.join(args.repo_path, subfolder, 'config.json'), 'r', encoding='utf-8') as reader:
										SCREAMING_SNAKE_CASE							:Dict							 =				reader.read()
										SCREAMING_SNAKE_CASE							:List[str]							 =				json.loads(text)

					if do_only_config:
										for key in config_parameters_to_change.keys():
															config.pop(key, None)

					if has_file(args.repo_path, 'config.json'):
										SCREAMING_SNAKE_CASE							:Optional[int]							 =				UNetaDModel(**config)
					else:
										SCREAMING_SNAKE_CASE							:Optional[Any]							 =				UNetaDConditionModel if 'ldm-text2im-large-256' in args.repo_path else UNetaDModel
										SCREAMING_SNAKE_CASE							:List[str]							 =				class_name(**config)

					if do_only_config:
										model.save_config(os.path.join(args.repo_path, subfolder))

					SCREAMING_SNAKE_CASE							:List[str]							 =				dict(model.config)

					if do_only_renaming:
										for key, value in config_parameters_to_change.items():
															if key in config:
																				SCREAMING_SNAKE_CASE							:Optional[Any]							 =				config[key]
																				del config[key]

										SCREAMING_SNAKE_CASE							:Optional[Any]							 =				[k.replace('UNetRes', '') for k in config['down_block_types']]
										SCREAMING_SNAKE_CASE							:List[Any]							 =				[k.replace('UNetRes', '') for k in config['up_block_types']]

					if do_only_weights:
										SCREAMING_SNAKE_CASE							:Tuple							 =				torch.load(os.path.join(args.repo_path, subfolder, 'diffusion_pytorch_model.bin'))

										SCREAMING_SNAKE_CASE							:Any							 =				{}
										for param_key, param_value in state_dict.items():
															if param_key.endswith('.op.bias') or param_key.endswith('.op.weight'):
																				continue
															SCREAMING_SNAKE_CASE							:List[str]							 =				False
															for key, new_key in key_parameters_to_change.items():
																				if not has_changed and param_key.split('.')[0] == key:
																									SCREAMING_SNAKE_CASE							:List[Any]							 =				param_value
																									SCREAMING_SNAKE_CASE							:str							 =				True
															if not has_changed:
																				SCREAMING_SNAKE_CASE							:List[str]							 =				param_value

										model.load_state_dict(new_state_dict)
										model.save_pretrained(os.path.join(args.repo_path, subfolder))
