



"""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 AutoImageProcessor, SwinvaConfig, SwinvaForImageClassification

def     UpperCamelCase_   (   lowerCAmelCase__				:    str						)						->       Optional[int]:

		"""simple docstring"""




		lowerCAmelCase_       :    Union[str, Any]								=			SwinvaConfig()
		lowerCAmelCase_       :    Tuple								=			swinva_name.split('_'						)

		lowerCAmelCase_       :    int								=			name_split[1]
		if "to" in name_split[3]:
				lowerCAmelCase_       :    Any								=			int(name_split[3][-3:]						)
		else:
				lowerCAmelCase_       :    Optional[int]								=			int(name_split[3]						)
		if "to" in name_split[2]:
				lowerCAmelCase_       :    List[Any]								=			int(name_split[2][-2:]						)
		else:
				lowerCAmelCase_       :    Optional[Any]								=			int(name_split[2][6:]						)

		if model_size == "tiny":
				lowerCAmelCase_       :    Optional[Any]								=			96
				lowerCAmelCase_       :    List[str]								=			(2, 2, 6, 2)
				lowerCAmelCase_       :    Tuple								=			(3, 6, 12, 24)
		elif model_size == "small":
				lowerCAmelCase_       :    str								=			96
				lowerCAmelCase_       :    Dict								=			(2, 2, 18, 2)
				lowerCAmelCase_       :    Optional[int]								=			(3, 6, 12, 24)
		elif model_size == "base":
				lowerCAmelCase_       :    str								=			128
				lowerCAmelCase_       :    Tuple								=			(2, 2, 18, 2)
				lowerCAmelCase_       :    Dict								=			(4, 8, 16, 32)
		else:
				lowerCAmelCase_       :    List[str]								=			192
				lowerCAmelCase_       :    Optional[int]								=			(2, 2, 18, 2)
				lowerCAmelCase_       :    Optional[int]								=			(6, 12, 24, 48)

		if "to" in swinva_name:
				lowerCAmelCase_       :    Optional[int]								=			(12, 12, 12, 6)

		if ("22k" in swinva_name) and ("to" not in swinva_name):
				lowerCAmelCase_       :    List[str]								=			2_1841
				lowerCAmelCase_       :    Union[str, Any]								=			'huggingface/label-files'
				lowerCAmelCase_       :    Any								=			'imagenet-22k-id2label.json'
				lowerCAmelCase_       :    Dict								=			json.load(open(hf_hub_download(__UpperCamelCase   ,       __UpperCamelCase   ,       repo_type='dataset'						)   ,       'r'						)						)
				lowerCAmelCase_       :    str								=			{int(__UpperCamelCase						): v for k, v in idalabel.items()}
				lowerCAmelCase_       :    Optional[Any]								=			idalabel
				lowerCAmelCase_       :    Tuple								=			{v: k for k, v in idalabel.items()}

		else:
				lowerCAmelCase_       :    Tuple								=			1000
				lowerCAmelCase_       :    int								=			'huggingface/label-files'
				lowerCAmelCase_       :    str								=			'imagenet-1k-id2label.json'
				lowerCAmelCase_       :    int								=			json.load(open(hf_hub_download(__UpperCamelCase   ,       __UpperCamelCase   ,       repo_type='dataset'						)   ,       'r'						)						)
				lowerCAmelCase_       :    Optional[int]								=			{int(__UpperCamelCase						): v for k, v in idalabel.items()}
				lowerCAmelCase_       :    Tuple								=			idalabel
				lowerCAmelCase_       :    Union[str, Any]								=			{v: k for k, v in idalabel.items()}

		lowerCAmelCase_       :    Any								=			img_size
		lowerCAmelCase_       :    Optional[int]								=			num_classes
		lowerCAmelCase_       :    str								=			embed_dim
		lowerCAmelCase_       :    List[str]								=			depths
		lowerCAmelCase_       :    str								=			num_heads
		lowerCAmelCase_       :    List[str]								=			window_size

		return config

def     UpperCamelCase_   (   lowerCAmelCase__				:    Union[str, Any]						)						->       Union[str, Any]:

		"""simple docstring"""




		if "patch_embed.proj" in name:
				lowerCAmelCase_       :    Optional[Any]								=			name.replace('patch_embed.proj'   ,       'embeddings.patch_embeddings.projection'						)
		if "patch_embed.norm" in name:
				lowerCAmelCase_       :    List[str]								=			name.replace('patch_embed.norm'   ,       'embeddings.norm'						)
		if "layers" in name:
				lowerCAmelCase_       :    str								=			'encoder.' + name
		if "attn.proj" in name:
				lowerCAmelCase_       :    Any								=			name.replace('attn.proj'   ,       'attention.output.dense'						)
		if "attn" in name:
				lowerCAmelCase_       :    Optional[int]								=			name.replace('attn'   ,       'attention.self'						)
		if "norm1" in name:
				lowerCAmelCase_       :    List[Any]								=			name.replace('norm1'   ,       'layernorm_before'						)
		if "norm2" in name:
				lowerCAmelCase_       :    Optional[int]								=			name.replace('norm2'   ,       'layernorm_after'						)
		if "mlp.fc1" in name:
				lowerCAmelCase_       :    List[str]								=			name.replace('mlp.fc1'   ,       'intermediate.dense'						)
		if "mlp.fc2" in name:
				lowerCAmelCase_       :    Union[str, Any]								=			name.replace('mlp.fc2'   ,       'output.dense'						)
		if "q_bias" in name:
				lowerCAmelCase_       :    Dict								=			name.replace('q_bias'   ,       'query.bias'						)
		if "k_bias" in name:
				lowerCAmelCase_       :    Union[str, Any]								=			name.replace('k_bias'   ,       'key.bias'						)
		if "v_bias" in name:
				lowerCAmelCase_       :    List[Any]								=			name.replace('v_bias'   ,       'value.bias'						)
		if "cpb_mlp" in name:
				lowerCAmelCase_       :    Optional[Any]								=			name.replace('cpb_mlp'   ,       'continuous_position_bias_mlp'						)
		if name == "norm.weight":
				lowerCAmelCase_       :    Tuple								=			'layernorm.weight'
		if name == "norm.bias":
				lowerCAmelCase_       :    Dict								=			'layernorm.bias'

		if "head" in name:
				lowerCAmelCase_       :    List[str]								=			name.replace('head'   ,       'classifier'						)
		else:
				lowerCAmelCase_       :    Optional[Any]								=			'swinv2.' + name

		return name

def     UpperCamelCase_   (   lowerCAmelCase__				:    str   ,       lowerCAmelCase__				:    Union[str, Any]						)						->       Dict:

		"""simple docstring"""




		for key in orig_state_dict.copy().keys():
				lowerCAmelCase_       :    Union[str, Any]								=			orig_state_dict.pop(__UpperCamelCase						)

				if "mask" in key:
						continue
				elif "qkv" in key:
						lowerCAmelCase_       :    List[Any]								=			key.split('.'						)
						lowerCAmelCase_       :    Tuple								=			int(key_split[1]						)
						lowerCAmelCase_       :    int								=			int(key_split[3]						)
						lowerCAmelCase_       :    Tuple								=			model.swinva.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

						if "weight" in key:
								lowerCAmelCase_       :    int								=			val[:dim, :]
								lowerCAmelCase_       :    Tuple								=			val[dim : dim * 2, :]
								lowerCAmelCase_       :    Any								=			val[-dim:, :]
						else:
								lowerCAmelCase_       :    Union[str, Any]								=			val[:dim]
								lowerCAmelCase_       :    Optional[Any]								=			val[
								    dim : dim * 2
								]
								lowerCAmelCase_       :    List[Any]								=			val[-dim:]
				else:
						lowerCAmelCase_       :    List[Any]								=			val

		return orig_state_dict

def     UpperCamelCase_   (   lowerCAmelCase__				:    Union[str, Any]   ,       lowerCAmelCase__				:    Any						)						->       List[Any]:

		"""simple docstring"""




		lowerCAmelCase_       :    Optional[int]								=			timm.create_model(__UpperCamelCase   ,       pretrained=__UpperCamelCase						)
		timm_model.eval()

		lowerCAmelCase_       :    Tuple								=			get_swinva_config(__UpperCamelCase						)
		lowerCAmelCase_       :    int								=			SwinvaForImageClassification(__UpperCamelCase						)
		model.eval()

		lowerCAmelCase_       :    str								=			convert_state_dict(timm_model.state_dict()   ,       __UpperCamelCase						)
		model.load_state_dict(__UpperCamelCase						)

		lowerCAmelCase_       :    Optional[Any]								=			'http://images.cocodataset.org/val2017/000000039769.jpg'

		lowerCAmelCase_       :    Optional[int]								=			AutoImageProcessor.from_pretrained('microsoft/{}'.format(swinva_name.replace('_'   ,       '-'						)						)						)
		lowerCAmelCase_       :    List[str]								=			Image.open(requests.get(__UpperCamelCase   ,       stream=__UpperCamelCase						).raw						)
		lowerCAmelCase_       :    str								=			image_processor(images=__UpperCamelCase   ,       return_tensors='pt'						)

		lowerCAmelCase_       :    Optional[Any]								=			timm_model(inputs['pixel_values']						)
		lowerCAmelCase_       :    List[str]								=			model(**__UpperCamelCase						).logits

		assert torch.allclose(__UpperCamelCase   ,       __UpperCamelCase   ,       atol=1e-3						)

		print(f"Saving model {swinva_name} to {pytorch_dump_folder_path}"						)
		model.save_pretrained(__UpperCamelCase						)

		print(f"Saving image processor to {pytorch_dump_folder_path}"						)
		image_processor.save_pretrained(__UpperCamelCase						)

		model.push_to_hub(
		    repo_path_or_name=Path(__UpperCamelCase   ,       __UpperCamelCase						)   ,       organization='nandwalritik'   ,       commit_message='Add model'   ,       )


if __name__ == "__main__":
					lowercase__		:    Optional[int]				   =     argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--swinv2_name""",
					    default="""swinv2_tiny_patch4_window8_256""",
					    type=str,
					    help="""Name of the Swinv2 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."""
					)

					lowercase__		:    List[str]				   =     parser.parse_args()
					convert_swinva_checkpoint(args.swinva_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 				UpperCamelCase__       :




      """simple docstring"""




      def __init__(				self    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict=1_3   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=7   ,      SCREAMING_SNAKE_CASE_    :				Dict=True   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				str=True   ,      SCREAMING_SNAKE_CASE_    :				List[str]=9_9   ,      SCREAMING_SNAKE_CASE_    :				int=1_6   ,      SCREAMING_SNAKE_CASE_    :				List[str]=3_6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=6   ,      SCREAMING_SNAKE_CASE_    :				Tuple=6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=6   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=3_7   ,      SCREAMING_SNAKE_CASE_    :				Tuple="gelu"   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				int=0.1   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=5_1_2   ,      SCREAMING_SNAKE_CASE_    :				List[str]=1_6   ,      SCREAMING_SNAKE_CASE_    :				List[str]=2   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.02   ,      SCREAMING_SNAKE_CASE_    :				Dict=3   ,      SCREAMING_SNAKE_CASE_    :				int=4   ,      SCREAMING_SNAKE_CASE_    :				Tuple=None   ,      ):
        lowerCAmelCase_       :    Any								=			parent
        lowerCAmelCase_       :    Optional[int]								=			batch_size
        lowerCAmelCase_       :    Dict								=			seq_length
        lowerCAmelCase_       :    Tuple								=			is_training
        lowerCAmelCase_       :    str								=			use_input_mask
        lowerCAmelCase_       :    Union[str, Any]								=			use_token_type_ids
        lowerCAmelCase_       :    Tuple								=			use_labels
        lowerCAmelCase_       :    Optional[int]								=			vocab_size
        lowerCAmelCase_       :    Any								=			embedding_size
        lowerCAmelCase_       :    Optional[Any]								=			hidden_size
        lowerCAmelCase_       :    str								=			num_hidden_layers
        lowerCAmelCase_       :    Optional[Any]								=			num_hidden_groups
        lowerCAmelCase_       :    Dict								=			num_attention_heads
        lowerCAmelCase_       :    Optional[Any]								=			intermediate_size
        lowerCAmelCase_       :    Any								=			hidden_act
        lowerCAmelCase_       :    Union[str, Any]								=			hidden_dropout_prob
        lowerCAmelCase_       :    int								=			attention_probs_dropout_prob
        lowerCAmelCase_       :    int								=			max_position_embeddings
        lowerCAmelCase_       :    List[Any]								=			type_vocab_size
        lowerCAmelCase_       :    Any								=			type_sequence_label_size
        lowerCAmelCase_       :    Optional[int]								=			initializer_range
        lowerCAmelCase_       :    Tuple								=			num_labels
        lowerCAmelCase_       :    Dict								=			num_choices
        lowerCAmelCase_       :    Tuple								=			scope



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

        lowerCAmelCase_       :    str								=			None
        if self.use_input_mask:
          lowerCAmelCase_       :    int								=			random_attention_mask([self.batch_size, self.seq_length]	)

        lowerCAmelCase_       :    List[Any]								=			None
        if self.use_token_type_ids:
          lowerCAmelCase_       :    str								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.type_vocab_size	)

        lowerCAmelCase_       :    List[str]								=			None
        lowerCAmelCase_       :    Union[str, Any]								=			None
        lowerCAmelCase_       :    str								=			None
        if self.use_labels:
          lowerCAmelCase_       :    Any								=			ids_tensor([self.batch_size]   ,      self.type_sequence_label_size	)
          lowerCAmelCase_       :    List[str]								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels	)
          lowerCAmelCase_       :    Dict								=			ids_tensor([self.batch_size]   ,      self.num_choices	)

        lowerCAmelCase_       :    List[Any]								=			self.get_config()

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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict	):
        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 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				str	):
        lowerCAmelCase_       :    Union[str, Any]								=			AlbertModel(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    List[str]								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    List[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_	)
        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 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Tuple	):
        lowerCAmelCase_       :    Optional[Any]								=			AlbertForPreTraining(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Optional[Any]								=			model(
            SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_   ,      sentence_order_label=SCREAMING_SNAKE_CASE_   ,      )
        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 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Dict	):
        lowerCAmelCase_       :    str								=			AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    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.vocab_size)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]	):
        lowerCAmelCase_       :    List[str]								=			AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    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 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict	):
        lowerCAmelCase_       :    Union[str, Any]								=			self.num_labels
        lowerCAmelCase_       :    Union[str, Any]								=			AlbertForSequenceClassification(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    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.num_labels)	)



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





      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    Optional[int]								=			self.prepare_config_and_inputs()
        (
            (
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,
        )       :    Optional[int]								=			config_and_inputs
        lowerCAmelCase_       :    Dict								=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_torch
class 				UpperCamelCase__       (     lowercase_, lowercase_, unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					(
          (
              AlbertModel,
              AlbertForPreTraining,
              AlbertForMaskedLM,
              AlbertForMultipleChoice,
              AlbertForSequenceClassification,
              AlbertForTokenClassification,
              AlbertForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_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 {}
      )
      _SCREAMING_SNAKE_CASE      =					True



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				str=False	):
        lowerCAmelCase_       :    List[str]								=			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_	):
            lowerCAmelCase_       :    List[str]								=			torch.zeros(
                (self.model_tester.batch_size, self.model_tester.seq_length)   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Any								=			torch.zeros(
                self.model_tester.batch_size   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)
        return inputs_dict



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    str								=			AlbertModelTester(self	)
        lowerCAmelCase_       :    Optional[int]								=			ConfigTester(self   ,      config_class=SCREAMING_SNAKE_CASE_   ,      hidden_size=3_7	)



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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    Any								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    Tuple								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    str								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        lowerCAmelCase_       :    Optional[int]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_	)



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





      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase_       :    Optional[Any]								=			AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)







@require_torch
class 				UpperCamelCase__       (     unittest.TestCase  ):




      """simple docstring"""




      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    Any								=			AlbertModel.from_pretrained('albert-base-v2'	)
        lowerCAmelCase_       :    Tuple								=			torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]	)
        lowerCAmelCase_       :    Tuple								=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
        with torch.no_grad():
          lowerCAmelCase_       :    List[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_	)[0]
        lowerCAmelCase_       :    str								=			torch.Size((1, 1_1, 7_6_8)	)
        self.assertEqual(output.shape   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    int								=			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]   ,      SCREAMING_SNAKE_CASE_   ,      atol=1E-4	)	)





