

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				:			Optional[Any]   ,   a				:			Any   ,   a				:			str=13   ,   a				:			Union[str, Any]=7   ,   a				:			int=True   ,   a				:			Optional[Any]=True   ,   a				:			Optional[int]=True   ,   a				:			int=True   ,   a				:			List[Any]=99   ,   a				:			List[Any]=16   ,   a				:			str=36   ,   a				:			Dict=6   ,   a				:			str=6   ,   a				:			List[str]=6   ,   a				:			Any=37   ,   a				:			Optional[Any]="gelu"   ,   a				:			List[str]=0.1   ,   a				:			Tuple=0.1   ,   a				:			int=512   ,   a				:			Optional[Any]=16   ,   a				:			List[str]=2   ,   a				:			List[str]=0.02   ,   a				:			int=3   ,   a				:			int=4   ,   a				:			Union[str, Any]=None   ,   )				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							parent
				SCREAMING_SNAKE_CASE :   List[str]					  =							batch_size
				SCREAMING_SNAKE_CASE :   int					  =							seq_length
				SCREAMING_SNAKE_CASE :   int					  =							is_training
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							use_input_mask
				SCREAMING_SNAKE_CASE :   str					  =							use_token_type_ids
				SCREAMING_SNAKE_CASE :   int					  =							use_labels
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							vocab_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							embedding_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							hidden_size
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							num_hidden_layers
				SCREAMING_SNAKE_CASE :   List[str]					  =							num_hidden_groups
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_attention_heads
				SCREAMING_SNAKE_CASE :   int					  =							intermediate_size
				SCREAMING_SNAKE_CASE :   List[Any]					  =							hidden_act
				SCREAMING_SNAKE_CASE :   List[Any]					  =							hidden_dropout_prob
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE :   Tuple					  =							max_position_embeddings
				SCREAMING_SNAKE_CASE :   Tuple					  =							type_vocab_size
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_sequence_label_size
				SCREAMING_SNAKE_CASE :   str					  =							initializer_range
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							num_labels
				SCREAMING_SNAKE_CASE :   int					  =							num_choices
				SCREAMING_SNAKE_CASE :   str					  =							scope
			def 				__UpperCamelCase			(      self				:			Any						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							ids_tensor([self.batch_size, self.seq_length]   ,   self.vocab_size						)

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

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

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							None
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							None
				SCREAMING_SNAKE_CASE :   List[Any]					  =							None
				if self.use_labels:
					SCREAMING_SNAKE_CASE :   int					  =							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 :   List[str]					  =							self.get_config()

				return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Tuple:

				"""simple docstring"""

				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				:			Union[str, Any]   ,   a				:			Optional[int]   ,   a				:			Optional[Any]   ,   a				:			int   ,   a				:			Dict   ,   a				:			int   ,   a				:			Union[str, Any]   ,   a				:			int						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							AlbertModel(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Any					  =							model(a   ,   attention_mask=a   ,   token_type_ids=a						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							model(a   ,   token_type_ids=a						)
				SCREAMING_SNAKE_CASE :   Dict					  =							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				:			str   ,   a				:			int   ,   a				:			Optional[int]   ,   a				:			Optional[Any]   ,   a				:			Dict   ,   a				:			str   ,   a				:			int   ,   a				:			Tuple						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							AlbertForPreTraining(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							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				:			Any   ,   a				:			Any   ,   a				:			Any   ,   a				:			Optional[int]   ,   a				:			str   ,   a				:			int   ,   a				:			str   ,   a				:			Any						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							AlbertForMaskedLM(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Dict					  =							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				:			str   ,   a				:			Tuple   ,   a				:			Optional[Any]   ,   a				:			Union[str, Any]   ,   a				:			Optional[int]   ,   a				:			Optional[int]   ,   a				:			List[Any]   ,   a				:			Tuple						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							AlbertForQuestionAnswering(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							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				:			List[Any]   ,   a				:			Tuple   ,   a				:			Tuple   ,   a				:			int   ,   a				:			List[Any]   ,   a				:			List[Any]   ,   a				:			Dict   ,   a				:			Any						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							AlbertForSequenceClassification(a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   int					  =							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				:			List[str]   ,   a				:			List[str]   ,   a				:			str   ,   a				:			Optional[int]   ,   a				:			Optional[int]   ,   a				:			str   ,   a				:			List[Any]   ,   a				:			List[str]						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							AlbertForTokenClassification(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   str					  =							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				:			str   ,   a				:			Optional[int]   ,   a				:			Tuple   ,   a				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			Union[str, Any]   ,   a				:			List[Any]   ,   a				:			str						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							self.num_choices
				SCREAMING_SNAKE_CASE :   Any					  =							AlbertForMultipleChoice(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							input_ids.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							token_type_ids.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   List[str]					  =							input_mask.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   List[str]					  =							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				:			int						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							self.prepare_config_and_inputs()
				(
				    (
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,(
				    SCREAMING_SNAKE_CASE
				)  ,
				) :   int					  =							config_and_inputs
				SCREAMING_SNAKE_CASE :   List[str]					  =							{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
				return config, inputs_dict






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







			lowerCamelCase__							     =(
			    (
			        AlbertModel,
			        AlbertForPreTraining,
			        AlbertForMaskedLM,
			        AlbertForMultipleChoice,
			        AlbertForSequenceClassification,
			        AlbertForTokenClassification,
			        AlbertForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__							     =(
			    {
			        'feature-extraction': AlbertModel,
			        'fill-mask': AlbertForMaskedLM,
			        'question-answering': AlbertForQuestionAnswering,
			        'text-classification': AlbertForSequenceClassification,
			        'token-classification': AlbertForTokenClassification,
			        'zero-shot': AlbertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__							     =True
			def 				__UpperCamelCase			(      self				:			List[str]   ,   a				:			int   ,   a				:			List[str]   ,   a				:			Optional[int]=False						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							super()._prepare_for_class(a   ,   a   ,   return_labels=a						)

				if return_labels:
					if model_class in get_values(a						):
						SCREAMING_SNAKE_CASE :   int					  =							torch.zeros(
						    (self.model_tester.batch_size, self.model_tester.seq_length)   ,   dtype=torch.long   ,   device=a						)
						SCREAMING_SNAKE_CASE :   Optional[Any]					  =							torch.zeros(
						    self.model_tester.batch_size   ,   dtype=torch.long   ,   device=a						)
				return inputs_dict
			def 				__UpperCamelCase			(      self				:			Dict						)				->							Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							AlbertModelTester(self						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							ConfigTester(self   ,   config_class=a   ,   hidden_size=37						)
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Optional[int]:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*a						)
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_pretraining(*a						)
			def 				__UpperCamelCase			(      self				:			Optional[int]						)				->							Union[str, Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*a						)
			def 				__UpperCamelCase			(      self				:			str						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*a						)
			def 				__UpperCamelCase			(      self				:			int						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*a						)
			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*a						)
			def 				__UpperCamelCase			(      self				:			Any						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
					SCREAMING_SNAKE_CASE :   Any					  =							type
					self.model_tester.create_and_check_model(*a						)






			@slow
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							str:

				"""simple docstring"""

				for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							AlbertModel.from_pretrained(a						)
					self.assertIsNotNone(a						)






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







			@slow
			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							AlbertModel.from_pretrained("albert-base-v2"						)
				SCREAMING_SNAKE_CASE :   int					  =							torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
				with torch.no_grad():
					SCREAMING_SNAKE_CASE :   List[Any]					  =							model(a   ,   attention_mask=a						)[0]
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.Size((1, 11, 768)						)
				self.assertEqual(output.shape   ,   a						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							torch.tensor(
				    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]						)

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



import importlib

import torch
import yaml
from omegaconf import OmegaConf
from taming.models.vqgan import VQModel




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_=False      )		->     Optional[Any]:




						'''simple docstring'''







						__UpperCAmelCase			:							Optional[Any]	    =	OmegaConf.load(lowercase_      )
						if display:
												print(yaml.dump(OmegaConf.to_container(lowercase_      )      )      )
						return config




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_=None      ,       lowercase_=None      )		->     Union[str, Any]:




						'''simple docstring'''







						if conf_path is None:
												__UpperCAmelCase			:							Optional[int]	    =	'''./model_checkpoints/vqgan_only.yaml'''
						__UpperCAmelCase			:							Dict	    =	load_config(lowercase_      ,       display=lowercase_      )
						__UpperCAmelCase			:							Tuple	    =	VQModel(**config.model.params      )
						if ckpt_path is None:
												__UpperCAmelCase			:							Dict	    =	'''./model_checkpoints/vqgan_only.pt'''
						__UpperCAmelCase			:							int	    =	torch.load(lowercase_      ,       map_location=lowercase_      )
						if ".ckpt" in ckpt_path:
												__UpperCAmelCase			:							Union[str, Any]	    =	sd['''state_dict''']
						model.load_state_dict(lowercase_      ,       strict=lowercase_      )
						model.to(lowercase_      )
						del sd
						return model




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      )		->     List[Any]:




						'''simple docstring'''







						__UpperCAmelCase    ,      __UpperCAmelCase    ,      __UpperCAmelCase			:							Optional[Any]	    =	model.encode(lowercase_      )
						print(f"VQGAN --- {model.__class__.__name__}: latent shape: {z.shape[2:]}"      )
						__UpperCAmelCase			:							List[str]	    =	model.decode(lowercase_      )
						return xrec




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_=False      )		->     int:




						'''simple docstring'''







						__UpperCAmelCase    ,      __UpperCAmelCase			:							Optional[int]	    =	string.rsplit('''.'''      ,       1      )
						if reload:
												__UpperCAmelCase			:							Optional[int]	    =	importlib.import_module(lowercase_      )
												importlib.reload(lowercase_      )
						return getattr(importlib.import_module(lowercase_      ,       package=lowercase_      )      ,       cls      )




def     __SCREAMING_SNAKE_CASE							(       lowercase_      )		->     Union[str, Any]:




						'''simple docstring'''







						if "target" not in config:
												raise KeyError('''Expected key `target` to instantiate.'''      )
						return get_obj_from_str(config['''target''']      )(**config.get('''params'''      ,       {}      )      )




def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      ,       lowercase_=True      ,       lowercase_=True      )		->     int:




						'''simple docstring'''







						__UpperCAmelCase			:							Dict	    =	instantiate_from_config(lowercase_      )
						if sd is not None:
												model.load_state_dict(lowercase_      )
						if gpu:
												model.cuda()
						if eval_mode:
												model.eval()
						return {"model": model}





def     __SCREAMING_SNAKE_CASE							(       lowercase_      ,       lowercase_      ,       lowercase_      ,       lowercase_      )		->     Optional[Any]:




						'''simple docstring'''







						if ckpt:
												__UpperCAmelCase			:							Optional[int]	    =	torch.load(lowercase_      ,       map_location='''cpu'''      )
												__UpperCAmelCase			:							Union[str, Any]	    =	pl_sd['''global_step''']
												print(f"loaded model from global step {global_step}."      )
						else:
												__UpperCAmelCase			:							Union[str, Any]	    =	{'''state_dict''': None}
												__UpperCAmelCase			:							Optional[int]	    =	None
						__UpperCAmelCase			:							str	    =	load_model_from_config(config.model      ,       pl_sd['''state_dict''']      ,       gpu=lowercase_      ,       eval_mode=lowercase_      )['''model''']
						return model, global_step

