"""simple docstring"""






import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__  :  Optional[int]						     =	logging.get_logger(__name__)



def   a_						(					lowerCamelCase					,			lowerCamelCase=False       ):
				UpperCAmelCase__             =							[]
				for i in range(config.num_hidden_layers       ):
								# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
								rename_keys.append((f'''blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')       )
								rename_keys.append((f'''blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')       )
								rename_keys.append((f'''blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')       )
								rename_keys.append((f'''blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')       )
								rename_keys.append((f'''blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')       )
								rename_keys.append((f'''blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')       )
								rename_keys.append((f'''blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')       )
								rename_keys.append((f'''blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')       )
								rename_keys.append((f'''blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')       )
								rename_keys.append((f'''blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')       )

				# projection layer + position embeddings
				rename_keys.extend(
				    [
				        ('cls_token', 'vit.embeddings.cls_token'),
				        ('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'),
				        ('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'),
				        ('pos_embed', 'vit.embeddings.position_embeddings'),
				    ]       )

				if base_model:
								# layernorm + pooler
								rename_keys.extend(
								    [
								        ('norm.weight', 'layernorm.weight'),
								        ('norm.bias', 'layernorm.bias'),
								        ('pre_logits.fc.weight', 'pooler.dense.weight'),
								        ('pre_logits.fc.bias', 'pooler.dense.bias'),
								    ]       )

								# if just the base model, we should remove "vit" from all keys that start with "vit"
								UpperCAmelCase__             =							[(pair[0], pair[1][4:]) if pair[1].startswith('vit'       ) else pair for pair in rename_keys]
				else:
								# layernorm + classification head
								rename_keys.extend(
								    [
								        ('norm.weight', 'vit.layernorm.weight'),
								        ('norm.bias', 'vit.layernorm.bias'),
								        ('head.weight', 'classifier.weight'),
								        ('head.bias', 'classifier.bias'),
								    ]       )

				return rename_keys



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase=False       ):
				for i in range(config.num_hidden_layers       ):
								if base_model:
												UpperCAmelCase__             =							''
								else:
												UpperCAmelCase__             =							'vit.'
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								UpperCAmelCase__             =							state_dict.pop(f'''blocks.{i}.attn.qkv.weight'''       )
								UpperCAmelCase__             =							state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''       )
								# next, add query, keys and values (in that order) to the state dict
								UpperCAmelCase__             =							in_proj_weight[
								    : config.hidden_size, :
								]
								UpperCAmelCase__             =							in_proj_bias[: config.hidden_size]
								UpperCAmelCase__             =							in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								UpperCAmelCase__             =							in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								UpperCAmelCase__             =							in_proj_weight[
								    -config.hidden_size :, :
								]
								UpperCAmelCase__             =							in_proj_bias[-config.hidden_size :]



def   a_						(					lowerCamelCase       ):
				UpperCAmelCase__             =							['head.weight', 'head.bias']
				for k in ignore_keys:
								state_dict.pop(lowerCamelCase					,			lowerCamelCase       )



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							dct.pop(lowerCamelCase       )
				UpperCAmelCase__             =							val



def   a_						(					):
				UpperCAmelCase__             =							'http://images.cocodataset.org/val2017/000000039769.jpg'
				UpperCAmelCase__             =							Image.open(requests.get(lowerCamelCase					,			stream=lowerCamelCase       ).raw       )
				return im



@torch.no_grad()
def   a_						(					lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							ViTConfig()
				UpperCAmelCase__             =							False
				# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
				if vit_name[-5:] == "in21k":
								UpperCAmelCase__             =							True
								UpperCAmelCase__             =							int(vit_name[-1_2:-1_0]       )
								UpperCAmelCase__             =							int(vit_name[-9:-6]       )
				else:
								UpperCAmelCase__             =							1_0_0_0
								UpperCAmelCase__             =							'huggingface/label-files'
								UpperCAmelCase__             =							'imagenet-1k-id2label.json'
								UpperCAmelCase__             =							json.load(open(hf_hub_download(lowerCamelCase					,			lowerCamelCase					,			repo_type='dataset'       )					,			'r'       )       )
								UpperCAmelCase__             =							{int(lowerCamelCase       ): v for k, v in idalabel.items()}
								UpperCAmelCase__             =							idalabel
								UpperCAmelCase__             =							{v: k for k, v in idalabel.items()}
								UpperCAmelCase__             =							int(vit_name[-6:-4]       )
								UpperCAmelCase__             =							int(vit_name[-3:]       )
				# size of the architecture
				if "deit" in vit_name:
								if vit_name[9:].startswith('tiny'       ):
												UpperCAmelCase__             =							1_9_2
												UpperCAmelCase__             =							7_6_8
												UpperCAmelCase__             =							1_2
												UpperCAmelCase__             =							3
								elif vit_name[9:].startswith('small'       ):
												UpperCAmelCase__             =							3_8_4
												UpperCAmelCase__             =							1_5_3_6
												UpperCAmelCase__             =							1_2
												UpperCAmelCase__             =							6
								else:
												pass
				else:
								if vit_name[4:].startswith('small'       ):
												UpperCAmelCase__             =							7_6_8
												UpperCAmelCase__             =							2_3_0_4
												UpperCAmelCase__             =							8
												UpperCAmelCase__             =							8
								elif vit_name[4:].startswith('base'       ):
												pass
								elif vit_name[4:].startswith('large'       ):
												UpperCAmelCase__             =							1_0_2_4
												UpperCAmelCase__             =							4_0_9_6
												UpperCAmelCase__             =							2_4
												UpperCAmelCase__             =							1_6
								elif vit_name[4:].startswith('huge'       ):
												UpperCAmelCase__             =							1_2_8_0
												UpperCAmelCase__             =							5_1_2_0
												UpperCAmelCase__             =							3_2
												UpperCAmelCase__             =							1_6

								# load original model from timm
				UpperCAmelCase__             =							timm.create_model(lowerCamelCase					,			pretrained=lowerCamelCase       )
				timm_model.eval()

				# load state_dict of original model, remove and rename some keys
				UpperCAmelCase__             =							timm_model.state_dict()
				if base_model:
								remove_classification_head_(lowerCamelCase       )
				UpperCAmelCase__             =							create_rename_keys(lowerCamelCase					,			lowerCamelCase       )
				for src, dest in rename_keys:
								rename_key(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       )
				read_in_q_k_v(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       )

				# load HuggingFace model
				if vit_name[-5:] == "in21k":
								UpperCAmelCase__             =							ViTModel(lowerCamelCase       ).eval()
				else:
								UpperCAmelCase__             =							ViTForImageClassification(lowerCamelCase       ).eval()
				model.load_state_dict(lowerCamelCase       )

				# Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
				if "deit" in vit_name:
								UpperCAmelCase__             =							DeiTImageProcessor(size=config.image_size       )
				else:
								UpperCAmelCase__             =							ViTImageProcessor(size=config.image_size       )
				UpperCAmelCase__             =							image_processor(images=prepare_img()					,			return_tensors='pt'       )
				UpperCAmelCase__             =							encoding['pixel_values']
				UpperCAmelCase__             =							model(lowerCamelCase       )

				if base_model:
								UpperCAmelCase__             =							timm_model.forward_features(lowerCamelCase       )
								assert timm_pooled_output.shape == outputs.pooler_output.shape
								assert torch.allclose(lowerCamelCase					,			outputs.pooler_output					,			atol=1e-3       )
				else:
								UpperCAmelCase__             =							timm_model(lowerCamelCase       )
								assert timm_logits.shape == outputs.logits.shape
								assert torch.allclose(lowerCamelCase					,			outputs.logits					,			atol=1e-3       )

				Path(lowerCamelCase       ).mkdir(exist_ok=lowerCamelCase       )
				print(f'''Saving model {vit_name} to {pytorch_dump_folder_path}'''       )
				model.save_pretrained(lowerCamelCase       )
				print(f'''Saving image processor to {pytorch_dump_folder_path}'''       )
				image_processor.save_pretrained(lowerCamelCase       )


if __name__ == "__main__":
						lowerCAmelCase__  :  Optional[int]						     =	argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--vit_name',
						    default='vit_base_patch16_224',
						    type=str,
						    help='Name of the ViT timm model you\'d like to convert.',
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
						)

						lowerCAmelCase__  :  Union[str, Any]						     =	parser.parse_args()
						convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)




"""simple docstring"""


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


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


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

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

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

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

      _UpperCAmelCase         =					self.get_config()

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


    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      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    lowerCAmelCase_     (    self		:     List[Any]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Optional[Any]      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     Any    ):
      _UpperCAmelCase         =					AlbertModel(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase      , attention_mask=__lowerCAmelCase      , token_type_ids=__lowerCAmelCase    )
      _UpperCAmelCase         =					model(__lowerCAmelCase      , token_type_ids=__lowerCAmelCase    )
      _UpperCAmelCase         =					model(__lowerCAmelCase    )
      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    lowerCAmelCase_     (    self		:     Any      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     Optional[Any]      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     int    ):
      _UpperCAmelCase         =					AlbertForPreTraining(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(
          __lowerCAmelCase      , attention_mask=__lowerCAmelCase      , token_type_ids=__lowerCAmelCase      , labels=__lowerCAmelCase      , sentence_order_label=__lowerCAmelCase      , )
      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    lowerCAmelCase_     (    self		:     List[str]      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     Tuple    ):
      _UpperCAmelCase         =					AlbertForMaskedLM(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase      , attention_mask=__lowerCAmelCase      , token_type_ids=__lowerCAmelCase      , labels=__lowerCAmelCase    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.vocab_size)    )


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


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     Optional[Any]      , __lowerCAmelCase		:     Optional[int]      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     Optional[Any]    ):
      _UpperCAmelCase         =					self.num_labels
      _UpperCAmelCase         =					AlbertForTokenClassification(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase      , attention_mask=__lowerCAmelCase      , token_type_ids=__lowerCAmelCase      , labels=__lowerCAmelCase    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.num_labels)    )


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


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



@require_torch
class 				a (						lowerCAmelCase_		,							lowerCAmelCase_		,							unittest.TestCase      ):
    _snake_case   :     str   							=							(
        (
            AlbertModel,
            AlbertForPreTraining,
            AlbertForMaskedLM,
            AlbertForMultipleChoice,
            AlbertForSequenceClassification,
            AlbertForTokenClassification,
            AlbertForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    _snake_case   :     Tuple   							=							(
        {
            '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   :     Dict   							=							True


    def    lowerCAmelCase_     (    self		:     str      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     List[Any]      , __lowerCAmelCase		:     List[Any]=False    ):
      _UpperCAmelCase         =					super()._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )

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


    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      _UpperCAmelCase         =					AlbertModelTester(self    )
      _UpperCAmelCase         =					ConfigTester(self      , config_class=__lowerCAmelCase      , hidden_size=37    )


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      self.config_tester.run_common_tests()


    def    lowerCAmelCase_     (    self		:     int    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Dict    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        _UpperCAmelCase         =					type
        self.model_tester.create_and_check_model(*__lowerCAmelCase    )


    @slow
    def    lowerCAmelCase_     (    self		:     Dict    ):
      for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					AlbertModel.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )



@require_torch
class 				a (						unittest.TestCase      ):


    @slow
    def    lowerCAmelCase_     (    self		:     Tuple    ):
      _UpperCAmelCase         =					AlbertModel.from_pretrained("""albert-base-v2"""    )
      _UpperCAmelCase         =					torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
      _UpperCAmelCase         =					torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
      with torch.no_grad():
        _UpperCAmelCase         =					model(__lowerCAmelCase      , attention_mask=__lowerCAmelCase    )[0]
      _UpperCAmelCase         =					torch.Size((1, 11, 768)    )
      self.assertEqual(output.shape      , __lowerCAmelCase    )
      _UpperCAmelCase         =					torch.tensor(
          [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]    )

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