






import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
				import torch

				from transformers import BertGenerationDecoder, BertGenerationEncoder


class 							__a							:





		'''simple docstring'''



		def __init__(						self ,    UpperCamelCase__ ,    UpperCamelCase__=13 ,    UpperCamelCase__=7 ,    UpperCamelCase__=True ,    UpperCamelCase__=True ,    UpperCamelCase__=99 ,    UpperCamelCase__=32 ,    UpperCamelCase__=5 ,    UpperCamelCase__=4 ,    UpperCamelCase__=37 ,    UpperCamelCase__="gelu" ,    UpperCamelCase__=0.1 ,    UpperCamelCase__=0.1 ,    UpperCamelCase__=50 ,    UpperCamelCase__=0.02 ,    UpperCamelCase__=True ,    UpperCamelCase__=None ,    ):
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   parent
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   batch_size
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   seq_length
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   is_training
						SCREAMING_SNAKE_CASE_   :	int													=   use_input_mask
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   vocab_size
						SCREAMING_SNAKE_CASE_   :	int													=   hidden_size
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   num_hidden_layers
						SCREAMING_SNAKE_CASE_   :	int													=   num_attention_heads
						SCREAMING_SNAKE_CASE_   :	Dict													=   intermediate_size
						SCREAMING_SNAKE_CASE_   :	List[str]													=   hidden_act
						SCREAMING_SNAKE_CASE_   :	int													=   hidden_dropout_prob
						SCREAMING_SNAKE_CASE_   :	str													=   attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   max_position_embeddings
						SCREAMING_SNAKE_CASE_   :	List[str]													=   initializer_range
						SCREAMING_SNAKE_CASE_   :	int													=   use_labels
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   scope


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

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

						if self.use_labels:
										SCREAMING_SNAKE_CASE_   :	int													=   ids_tensor([self.batch_size, self.seq_length] ,    self.vocab_size       )

						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   self.get_config()

						return config, input_ids, input_mask, token_labels


		def   __snake_case			(						self       ):
						return BertGenerationConfig(
						    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 ,    is_decoder=UpperCamelCase__ ,    initializer_range=self.initializer_range ,    )


		def   __snake_case			(						self       ):
						(
						    (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       
						)   :	int													=   self.prepare_config_and_inputs()

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

						return (
						    config,
						    input_ids,
						    input_mask,
						    token_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    **UpperCamelCase__ ,    ):
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   BertGenerationEncoder(config=UpperCamelCase__       )
						model.to(UpperCamelCase__       )
						model.eval()
						SCREAMING_SNAKE_CASE_   :	List[str]													=   model(UpperCamelCase__ ,    attention_mask=UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	List[str]													=   model(UpperCamelCase__       )
						self.parent.assertEqual(result.last_hidden_state.shape ,    (self.batch_size, self.seq_length, self.hidden_size)       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    **UpperCamelCase__ ,    ):
						SCREAMING_SNAKE_CASE_   :	List[str]													=   True
						SCREAMING_SNAKE_CASE_   :	str													=   BertGenerationEncoder(config=UpperCamelCase__       )
						model.to(UpperCamelCase__       )
						model.eval()
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   model(
						    UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    encoder_hidden_states=UpperCamelCase__ ,    encoder_attention_mask=UpperCamelCase__ ,    )
						SCREAMING_SNAKE_CASE_   :	Dict													=   model(
						    UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    encoder_hidden_states=UpperCamelCase__ ,    )
						self.parent.assertEqual(result.last_hidden_state.shape ,    (self.batch_size, self.seq_length, self.hidden_size)       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    **UpperCamelCase__ ,    ):
						SCREAMING_SNAKE_CASE_   :	List[str]													=   True
						SCREAMING_SNAKE_CASE_   :	Tuple													=   True
						SCREAMING_SNAKE_CASE_   :	Dict													=   BertGenerationDecoder(config=UpperCamelCase__       ).to(UpperCamelCase__       ).eval()

						# first forward pass
						SCREAMING_SNAKE_CASE_   :	int													=   model(
						    UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    encoder_hidden_states=UpperCamelCase__ ,    encoder_attention_mask=UpperCamelCase__ ,    use_cache=UpperCamelCase__ ,    )
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   outputs.past_key_values

						# create hypothetical multiple next token and extent to next_input_ids
						SCREAMING_SNAKE_CASE_   :	Dict													=   ids_tensor((self.batch_size, 3) ,    config.vocab_size       )
						SCREAMING_SNAKE_CASE_   :	Dict													=   ids_tensor((self.batch_size, 3) ,    vocab_size=2       )

						# append to next input_ids and
						SCREAMING_SNAKE_CASE_   :	Tuple													=   torch.cat([input_ids, next_tokens] ,    dim=-1       )
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   torch.cat([input_mask, next_mask] ,    dim=-1       )

						SCREAMING_SNAKE_CASE_   :	int													=   model(
						    UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    encoder_hidden_states=UpperCamelCase__ ,    encoder_attention_mask=UpperCamelCase__ ,    output_hidden_states=UpperCamelCase__ ,    )['hidden_states'][0]
						SCREAMING_SNAKE_CASE_   :	List[str]													=   model(
						    UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    encoder_hidden_states=UpperCamelCase__ ,    encoder_attention_mask=UpperCamelCase__ ,    past_key_values=UpperCamelCase__ ,    output_hidden_states=UpperCamelCase__ ,    )['hidden_states'][0]

						# select random slice
						SCREAMING_SNAKE_CASE_   :	Tuple													=   ids_tensor((1,) ,    output_from_past.shape[-1]       ).item()
						SCREAMING_SNAKE_CASE_   :	List[str]													=   output_from_no_past[:, -3:, random_slice_idx].detach()
						SCREAMING_SNAKE_CASE_   :	int													=   output_from_past[:, :, random_slice_idx].detach()

						self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]       )

						# test that outputs are equal for slice
						self.parent.assertTrue(torch.allclose(UpperCamelCase__ ,    UpperCamelCase__ ,    atol=1E-3       )       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    *UpperCamelCase__ ,    ):
						SCREAMING_SNAKE_CASE_   :	Dict													=   BertGenerationDecoder(UpperCamelCase__       )
						model.to(UpperCamelCase__       )
						model.eval()
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   model(UpperCamelCase__ ,    attention_mask=UpperCamelCase__ ,    labels=UpperCamelCase__       )
						self.parent.assertEqual(result.logits.shape ,    (self.batch_size, self.seq_length, self.vocab_size)       )






		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	List[str]													=   self.prepare_config_and_inputs()
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   {'input_ids': input_ids, 'attention_mask': input_mask}
						return config, inputs_dict






@require_torch
class 							__a							(  __A ,	__A ,	__A ,	unittest.TestCase						):





		'''simple docstring'''



		UpperCAmelCase__      :				List[str]				    =						(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
		UpperCAmelCase__      :				Dict				    =						(BertGenerationDecoder,) if is_torch_available() else ()
		UpperCAmelCase__      :				Dict				    =						(
		    {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
		    if is_torch_available()
		    else {}
		)


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   BertGenerationEncoderTester(self       )
						SCREAMING_SNAKE_CASE_   :	Any													=   ConfigTester(self ,    config_class=UpperCamelCase__ ,    hidden_size=37       )


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


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Tuple													=   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase__       )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	List[Any]													=   self.model_tester.prepare_config_and_inputs()
						SCREAMING_SNAKE_CASE_   :	List[str]													=   'bert'
						self.model_tester.create_and_check_model(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__       )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_decoder_model_past_large_inputs(*UpperCamelCase__       )


		def   __snake_case			(						self       ):
						# This regression test was failing with PyTorch < 1.3
						(
						    (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       (
						    SCREAMING_SNAKE_CASE_
						)							,       
						)   :	List[str]													=   self.model_tester.prepare_config_and_inputs_for_decoder()

						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   None

						self.model_tester.create_and_check_model_as_decoder(
						    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__ ,    )


		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Any													=   self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_for_causal_lm(*UpperCamelCase__       )






		@slow
		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Any													=   BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'       )
						self.assertIsNotNone(UpperCamelCase__       )






@require_torch
class 							__a							(  unittest.TestCase						):





		'''simple docstring'''



		@slow
		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'       )
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]]       )
						with torch.no_grad():
										SCREAMING_SNAKE_CASE_   :	List[Any]													=   model(UpperCamelCase__       )[0]
						SCREAMING_SNAKE_CASE_   :	List[str]													=   torch.Size([1, 8, 1024]       )
						self.assertEqual(output.shape ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	int													=   torch.tensor(
						    [[[0.17_75, 0.00_83, -0.03_21], [1.60_02, 0.12_87, 0.39_12], [2.14_73, 0.57_91, 0.60_66]]]       )
						self.assertTrue(torch.allclose(output[:, :3, :3] ,    UpperCamelCase__ ,    atol=1E-4       )       )






@require_torch
class 							__a							(  unittest.TestCase						):





		'''simple docstring'''



		@slow
		def   __snake_case			(						self       ):
						SCREAMING_SNAKE_CASE_   :	Dict													=   BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'       )
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]]       )
						with torch.no_grad():
										SCREAMING_SNAKE_CASE_   :	int													=   model(UpperCamelCase__       )[0]
						SCREAMING_SNAKE_CASE_   :	int													=   torch.Size([1, 8, 50358]       )
						self.assertEqual(output.shape ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.tensor(
						    [[[-0.57_88, -2.59_94, -3.70_54], [0.04_38, 4.79_97, 1.87_95], [1.58_62, 6.64_09, 4.46_38]]]       )
						self.assertTrue(torch.allclose(output[:, :3, :3] ,    UpperCamelCase__ ,    atol=1E-4       )       )







import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor



def 					_lowerCamelCase(			lowerCAmelCase__				:		Optional[Any]      ):





				'''simple docstring'''







				if "cls_token" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   name.replace('cls_token'    ,    'vit.embeddings.cls_token'      )
				if "mask_token" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   name.replace('mask_token'    ,    'decoder.mask_token'      )
				if "decoder_pos_embed" in name:
								SCREAMING_SNAKE_CASE_   :	Any													=   name.replace('decoder_pos_embed'    ,    'decoder.decoder_pos_embed'      )
				if "pos_embed" in name and "decoder" not in name:
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   name.replace('pos_embed'    ,    'vit.embeddings.position_embeddings'      )
				if "patch_embed.proj" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   name.replace('patch_embed.proj'    ,    'vit.embeddings.patch_embeddings.projection'      )
				if "patch_embed.norm" in name:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   name.replace('patch_embed.norm'    ,    'vit.embeddings.norm'      )
				if "decoder_blocks" in name:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   name.replace('decoder_blocks'    ,    'decoder.decoder_layers'      )
				if "blocks" in name:
								SCREAMING_SNAKE_CASE_   :	Any													=   name.replace('blocks'    ,    'vit.encoder.layer'      )
				if "attn.proj" in name:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   name.replace('attn.proj'    ,    'attention.output.dense'      )
				if "attn" in name:
								SCREAMING_SNAKE_CASE_   :	Any													=   name.replace('attn'    ,    'attention.self'      )
				if "norm1" in name:
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   name.replace('norm1'    ,    'layernorm_before'      )
				if "norm2" in name:
								SCREAMING_SNAKE_CASE_   :	Any													=   name.replace('norm2'    ,    'layernorm_after'      )
				if "mlp.fc1" in name:
								SCREAMING_SNAKE_CASE_   :	Dict													=   name.replace('mlp.fc1'    ,    'intermediate.dense'      )
				if "mlp.fc2" in name:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   name.replace('mlp.fc2'    ,    'output.dense'      )
				if "decoder_embed" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   name.replace('decoder_embed'    ,    'decoder.decoder_embed'      )
				if "decoder_norm" in name:
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   name.replace('decoder_norm'    ,    'decoder.decoder_norm'      )
				if "decoder_pred" in name:
								SCREAMING_SNAKE_CASE_   :	str													=   name.replace('decoder_pred'    ,    'decoder.decoder_pred'      )
				if "norm.weight" in name and "decoder" not in name:
								SCREAMING_SNAKE_CASE_   :	Tuple													=   name.replace('norm.weight'    ,    'vit.layernorm.weight'      )
				if "norm.bias" in name and "decoder" not in name:
								SCREAMING_SNAKE_CASE_   :	List[str]													=   name.replace('norm.bias'    ,    'vit.layernorm.bias'      )

				return name



def 					_lowerCamelCase(			lowerCAmelCase__				:		List[str]    ,    lowerCAmelCase__				:		Optional[int]      ):





				'''simple docstring'''







				for key in orig_state_dict.copy().keys():
								SCREAMING_SNAKE_CASE_   :	Tuple													=   orig_state_dict.pop(lowerCAmelCase__      )

								if "qkv" in key:
												SCREAMING_SNAKE_CASE_   :	Optional[int]													=   key.split('.'      )
												SCREAMING_SNAKE_CASE_   :	List[Any]													=   int(key_split[1]      )
												if "decoder_blocks" in key:
																SCREAMING_SNAKE_CASE_   :	Any													=   config.decoder_hidden_size
																SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   'decoder.decoder_layers.'
																if "weight" in key:
																				SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   val[:dim, :]
																				SCREAMING_SNAKE_CASE_   :	List[str]													=   val[dim : dim * 2, :]
																				SCREAMING_SNAKE_CASE_   :	Any													=   val[-dim:, :]
																elif "bias" in key:
																				SCREAMING_SNAKE_CASE_   :	str													=   val[:dim]
																				SCREAMING_SNAKE_CASE_   :	List[Any]													=   val[dim : dim * 2]
																				SCREAMING_SNAKE_CASE_   :	List[Any]													=   val[-dim:]
												else:
																SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   config.hidden_size
																SCREAMING_SNAKE_CASE_   :	str													=   'vit.encoder.layer.'
																if "weight" in key:
																				SCREAMING_SNAKE_CASE_   :	Any													=   val[:dim, :]
																				SCREAMING_SNAKE_CASE_   :	Tuple													=   val[dim : dim * 2, :]
																				SCREAMING_SNAKE_CASE_   :	int													=   val[-dim:, :]
																elif "bias" in key:
																				SCREAMING_SNAKE_CASE_   :	List[Any]													=   val[:dim]
																				SCREAMING_SNAKE_CASE_   :	Dict													=   val[dim : dim * 2]
																				SCREAMING_SNAKE_CASE_   :	str													=   val[-dim:]

								else:
												SCREAMING_SNAKE_CASE_   :	Optional[int]													=   val

				return orig_state_dict



def 					_lowerCamelCase(			lowerCAmelCase__				:		Optional[Any]    ,    lowerCAmelCase__				:		List[str]      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Dict													=   ViTMAEConfig()
				if "large" in checkpoint_url:
								SCREAMING_SNAKE_CASE_   :	Any													=   1024
								SCREAMING_SNAKE_CASE_   :	Any													=   4096
								SCREAMING_SNAKE_CASE_   :	List[str]													=   24
								SCREAMING_SNAKE_CASE_   :	List[str]													=   16
				elif "huge" in checkpoint_url:
								SCREAMING_SNAKE_CASE_   :	str													=   14
								SCREAMING_SNAKE_CASE_   :	List[str]													=   1280
								SCREAMING_SNAKE_CASE_   :	int													=   5120
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   32
								SCREAMING_SNAKE_CASE_   :	List[str]													=   16

				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   ViTMAEForPreTraining(lowerCAmelCase__      )

				SCREAMING_SNAKE_CASE_   :	List[str]													=   torch.hub.load_state_dict_from_url(lowerCAmelCase__    ,    map_location='cpu'      )['model']

				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   ViTMAEImageProcessor(size=config.image_size      )

				SCREAMING_SNAKE_CASE_   :	int													=   convert_state_dict(lowerCAmelCase__    ,    lowerCAmelCase__      )

				model.load_state_dict(lowerCAmelCase__      )
				model.eval()

				SCREAMING_SNAKE_CASE_   :	Any													=   'https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'

				SCREAMING_SNAKE_CASE_   :	Dict													=   Image.open(requests.get(lowerCAmelCase__    ,    stream=lowerCAmelCase__      ).raw      )
				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   ViTMAEImageProcessor(size=config.image_size      )
				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   image_processor(images=lowerCAmelCase__    ,    return_tensors='pt'      )

				# forward pass
				torch.manual_seed(2      )
				SCREAMING_SNAKE_CASE_   :	Any													=   model(**lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   outputs.logits

				if "large" in checkpoint_url:
								SCREAMING_SNAKE_CASE_   :	Any													=   torch.tensor(
								    [[-0.7_309, -0.7_128, -1.0_169], [-1.0_161, -0.9_058, -1.1_878], [-1.0_478, -0.9_411, -1.1_911]]      )
				elif "huge" in checkpoint_url:
								SCREAMING_SNAKE_CASE_   :	Dict													=   torch.tensor(
								    [[-1.1_599, -0.9_199, -1.2_221], [-1.1_952, -0.9_269, -1.2_307], [-1.2_143, -0.9_337, -1.2_262]]      )
				else:
								SCREAMING_SNAKE_CASE_   :	int													=   torch.tensor(
								    [[-0.9_192, -0.8_481, -1.1_259], [-1.1_349, -1.0_034, -1.2_599], [-1.1_757, -1.0_429, -1.2_726]]      )

				# verify logits
				assert torch.allclose(logits[0, :3, :3]    ,    lowerCAmelCase__    ,    atol=1E-4      )

				print(F'''Saving model 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__":
				A        =			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--checkpoint_url',
				    default='https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth',
				    type=str,
				    help='URL of the checkpoint you\'d like to convert.',
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
				)

				A        =			parser.parse_args()
				convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)