

"""simple docstring"""




import argparse
import hashlib
import os
import urllib
import warnings

import torch
from torch import nn
from tqdm import tqdm

from transformers import WhisperConfig, WhisperForConditionalGeneration


A_				:      List[str]           =			{
    "tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
    "tiny": "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
    "base.en": "https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt",
    "base": "https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt",
    "small.en": "https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt",
    "small": "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt",
    "medium.en": "https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt",
    "medium": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
    "large": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large.pt",
    "large-v2": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
}

def 						A							(  snake_case__     ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__										=   ["""layers""", """blocks"""]
			for k in ignore_keys:
						state_dict.pop(snake_case__		,					snake_case__     )


A_				:      List[Any]           =			{
    "blocks": "layers",
    "mlp.0": "fc1",
    "mlp.2": "fc2",
    "mlp_ln": "final_layer_norm",
    ".attn.query": ".self_attn.q_proj",
    ".attn.key": ".self_attn.k_proj",
    ".attn.value": ".self_attn.v_proj",
    ".attn_ln": ".self_attn_layer_norm",
    ".attn.out": ".self_attn.out_proj",
    ".cross_attn.query": ".encoder_attn.q_proj",
    ".cross_attn.key": ".encoder_attn.k_proj",
    ".cross_attn.value": ".encoder_attn.v_proj",
    ".cross_attn_ln": ".encoder_attn_layer_norm",
    ".cross_attn.out": ".encoder_attn.out_proj",
    "decoder.ln.": "decoder.layer_norm.",
    "encoder.ln.": "encoder.layer_norm.",
    "token_embedding": "embed_tokens",
    "encoder.positional_embedding": "encoder.embed_positions.weight",
    "decoder.positional_embedding": "decoder.embed_positions.weight",
    "ln_post": "layer_norm",
}

def 						A							(  snake_case__     ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__										=   list(s_dict.keys()     )
			for key in keys:
						SCREAMING_SNAKE_CASE__										=   key
						for k, v in WHISPER_MAPPING.items():
									if k in key:
												SCREAMING_SNAKE_CASE__										=   new_key.replace(snake_case__		,					snake_case__     )

						print(f"""{key} -> {new_key}"""     )

						SCREAMING_SNAKE_CASE__										=   s_dict.pop(snake_case__     )
			return s_dict

def 						A							(  snake_case__     ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   emb.weight.shape
			SCREAMING_SNAKE_CASE__										=   nn.Linear(snake_case__		,					snake_case__		,					bias=snake_case__     )
			SCREAMING_SNAKE_CASE__										=   emb.weight.data
			return lin_layer

def 						A							(  snake_case__		,					snake_case__     ):




			'''simple docstring'''


			os.makedirs(snake_case__		,					exist_ok=snake_case__     )
			SCREAMING_SNAKE_CASE__										=   os.path.basename(snake_case__     )

			SCREAMING_SNAKE_CASE__										=   url.split("""/"""     )[-2]
			SCREAMING_SNAKE_CASE__										=   os.path.join(snake_case__		,					snake_case__     )

			if os.path.exists(snake_case__     ) and not os.path.isfile(snake_case__     ):
						raise RuntimeError(f"""{download_target} exists and is not a regular file"""     )

			if os.path.isfile(snake_case__     ):
						SCREAMING_SNAKE_CASE__										=   open(snake_case__		,					"""rb"""     ).read()
						if hashlib.shaaaa(snake_case__     ).hexdigest() == expected_shaaaa:
									return model_bytes
						else:
									warnings.warn(f"""{download_target} exists, but the SHA256 checksum does not match; re-downloading the file"""     )

			with urllib.request.urlopen(snake_case__     ) as source, open(snake_case__		,					"""wb"""     ) as output:
						with tqdm(
						    total=int(source.info().get("""Content-Length"""     )     )		,					ncols=80		,					unit="""iB"""		,					unit_scale=snake_case__		,					unit_divisor=10_24     ) as loop:
									while True:
												SCREAMING_SNAKE_CASE__										=   source.read(81_92     )
												if not buffer:
															break

												output.write(snake_case__     )
												loop.update(len(snake_case__     )     )

			SCREAMING_SNAKE_CASE__										=   open(snake_case__		,					"""rb"""     ).read()
			if hashlib.shaaaa(snake_case__     ).hexdigest() != expected_shaaaa:
						raise RuntimeError(
						    """Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model."""     )

			return model_bytes

def 						A							(  snake_case__		,					snake_case__     ):




			'''simple docstring'''


			if ".pt" not in checkpoint_path:
						SCREAMING_SNAKE_CASE__										=   _download(_MODELS[checkpoint_path]     )
			else:
						SCREAMING_SNAKE_CASE__										=   torch.load(snake_case__		,					map_location="""cpu"""     )
			SCREAMING_SNAKE_CASE__										=   original_checkpoint["""dims"""]
			SCREAMING_SNAKE_CASE__										=   original_checkpoint["""model_state_dict"""]
			SCREAMING_SNAKE_CASE__										=   state_dict["""decoder.token_embedding.weight"""]
			remove_ignore_keys_(snake_case__     )
			rename_keys(snake_case__     )
			SCREAMING_SNAKE_CASE__										=   True
			SCREAMING_SNAKE_CASE__										=   state_dict["""decoder.layers.0.fc1.weight"""].shape[0]

			SCREAMING_SNAKE_CASE__										=   WhisperConfig(
			    vocab_size=dimensions["""n_vocab"""]		,					encoder_ffn_dim=snake_case__		,					decoder_ffn_dim=snake_case__		,					num_mel_bins=dimensions["""n_mels"""]		,					d_model=dimensions["""n_audio_state"""]		,					max_target_positions=dimensions["""n_text_ctx"""]		,					encoder_layers=dimensions["""n_audio_layer"""]		,					encoder_attention_heads=dimensions["""n_audio_head"""]		,					decoder_layers=dimensions["""n_text_layer"""]		,					decoder_attention_heads=dimensions["""n_text_state"""]		,					max_source_positions=dimensions["""n_audio_ctx"""]		,					)

			SCREAMING_SNAKE_CASE__										=   WhisperForConditionalGeneration(snake_case__     )
			SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   model.model.load_state_dict(snake_case__		,					strict=snake_case__     )
			if len(snake_case__     ) > 0 and not set(snake_case__     ) <= {
			    "encoder.embed_positions.weights",
			    "decoder.embed_positions.weights",
			}:
						raise ValueError(
						    """Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,"""
						    f""" but all the following weights are missing {missing}"""     )

			if tie_embeds:
						SCREAMING_SNAKE_CASE__										=   make_linear_from_emb(model.model.decoder.embed_tokens     )
			else:
						SCREAMING_SNAKE_CASE__										=   proj_out_weights

			model.save_pretrained(snake_case__     )


if __name__ == "__main__":
						A_				:      Tuple           =			argparse.ArgumentParser()
						# # Required parameters
						parser.add_argument("--checkpoint_path", type=str, help="Patht to the downloaded checkpoints")
						parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
						A_				:      Tuple           =			parser.parse_args()

						convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)





"""simple docstring"""




import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch
						from torch import nn

						from transformers import ViTMAEForPreTraining, ViTMAEModel
						from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import ViTImageProcessor





class  lowerCamelCase    :

			def __init__(       self					:    Dict     ,					__UpperCAmelCase					:    Union[str, Any]     ,					__UpperCAmelCase					:    Dict=1_3     ,					__UpperCAmelCase					:    Optional[Any]=3_0     ,					__UpperCAmelCase					:    Optional[Any]=2     ,					__UpperCAmelCase					:    Tuple=3     ,					__UpperCAmelCase					:    List[Any]=True     ,					__UpperCAmelCase					:    Optional[Any]=True     ,					__UpperCAmelCase					:    List[Any]=3_2     ,					__UpperCAmelCase					:    int=5     ,					__UpperCAmelCase					:    Union[str, Any]=4     ,					__UpperCAmelCase					:    Union[str, Any]=3_7     ,					__UpperCAmelCase					:    Optional[int]="gelu"     ,					__UpperCAmelCase					:    Dict=0.1     ,					__UpperCAmelCase					:    List[str]=0.1     ,					__UpperCAmelCase					:    Optional[Any]=1_0     ,					__UpperCAmelCase					:    List[str]=0.02     ,					__UpperCAmelCase					:    int=3     ,					__UpperCAmelCase					:    Any=0.6     ,					__UpperCAmelCase					:    Dict=None     ,					) ->     str:
						SCREAMING_SNAKE_CASE__										=   parent
						SCREAMING_SNAKE_CASE__										=   batch_size
						SCREAMING_SNAKE_CASE__										=   image_size
						SCREAMING_SNAKE_CASE__										=   patch_size
						SCREAMING_SNAKE_CASE__										=   num_channels
						SCREAMING_SNAKE_CASE__										=   is_training
						SCREAMING_SNAKE_CASE__										=   use_labels
						SCREAMING_SNAKE_CASE__										=   hidden_size
						SCREAMING_SNAKE_CASE__										=   num_hidden_layers
						SCREAMING_SNAKE_CASE__										=   num_attention_heads
						SCREAMING_SNAKE_CASE__										=   intermediate_size
						SCREAMING_SNAKE_CASE__										=   hidden_act
						SCREAMING_SNAKE_CASE__										=   hidden_dropout_prob
						SCREAMING_SNAKE_CASE__										=   attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE__										=   type_sequence_label_size
						SCREAMING_SNAKE_CASE__										=   initializer_range
						SCREAMING_SNAKE_CASE__										=   mask_ratio
						SCREAMING_SNAKE_CASE__										=   scope

						# in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
						# (we add 1 for the [CLS] token)
						SCREAMING_SNAKE_CASE__										=   (image_size // patch_size) ** 2
						SCREAMING_SNAKE_CASE__										=   int(math.ceil((1 - mask_ratio) * (num_patches + 1)							)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     str:
						SCREAMING_SNAKE_CASE__										=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

						SCREAMING_SNAKE_CASE__										=   None
						if self.use_labels:
									SCREAMING_SNAKE_CASE__										=   ids_tensor([self.batch_size]     ,					self.type_sequence_label_size							)

						SCREAMING_SNAKE_CASE__										=   self.get_config()

						return config, pixel_values, labels

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Union[str, Any]:
						return ViTMAEConfig(
						    image_size=self.image_size     ,					patch_size=self.patch_size     ,					num_channels=self.num_channels     ,					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     ,					is_decoder=__UpperCAmelCase     ,					initializer_range=self.initializer_range     ,					mask_ratio=self.mask_ratio     ,					)

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]     ,					__UpperCAmelCase					:    Any     ,					__UpperCAmelCase					:    Union[str, Any]     ,					__UpperCAmelCase					:    str							) ->     Dict:
						SCREAMING_SNAKE_CASE__										=   ViTMAEModel(config=__UpperCAmelCase							)
						model.to(__UpperCAmelCase							)
						model.eval()
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)
						self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, self.seq_length, self.hidden_size)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]     ,					__UpperCAmelCase					:    Optional[int]     ,					__UpperCAmelCase					:    List[str]     ,					__UpperCAmelCase					:    Dict							) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   ViTMAEForPreTraining(__UpperCAmelCase							)
						model.to(__UpperCAmelCase							)
						model.eval()
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   (self.image_size // self.patch_size) ** 2
						SCREAMING_SNAKE_CASE__										=   self.patch_size**2 * self.num_channels
						self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, num_patches, expected_num_channels)							)

						# test greyscale images
						SCREAMING_SNAKE_CASE__										=   1
						SCREAMING_SNAKE_CASE__										=   ViTMAEForPreTraining(__UpperCAmelCase							)
						model.to(__UpperCAmelCase							)
						model.eval()
						SCREAMING_SNAKE_CASE__										=   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   self.patch_size**2
						self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, num_patches, expected_num_channels)							)





			def        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   self.prepare_config_and_inputs()
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   config_and_inputs
						SCREAMING_SNAKE_CASE__										=   {"""pixel_values""": pixel_values}
						return config, inputs_dict





@require_torch
class  lowerCamelCase    (A__      ,A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		Optional[int]       		=				(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
			lowerCamelCase__ :		str       		=				{'feature-extraction': ViTMAEModel} if is_torch_available() else {}

			lowerCamelCase__ :		Tuple       		=				False
			lowerCamelCase__ :		List[str]       		=				False
			lowerCamelCase__ :		str       		=				False
			lowerCamelCase__ :		List[Any]       		=				False

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   ViTMAEModelTester(self							)
						SCREAMING_SNAKE_CASE__										=   ConfigTester(self     ,					config_class=__UpperCAmelCase     ,					has_text_modality=__UpperCAmelCase     ,					hidden_size=3_7							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     List[Any]:
						self.config_tester.run_common_tests()

			@unittest.skip(reason="""ViTMAE does not use inputs_embeds"""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     int:
						pass

			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     List[Any]:
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE__										=   model_class(__UpperCAmelCase							)
									self.assertIsInstance(model.get_input_embeddings()     ,					(nn.Module)							)
									SCREAMING_SNAKE_CASE__										=   model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(__UpperCAmelCase     ,					nn.Linear							)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Optional[int]:
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE__										=   model_class(__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   inspect.signature(model.forward							)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									SCREAMING_SNAKE_CASE__										=   [*signature.parameters.keys()]

									SCREAMING_SNAKE_CASE__										=   ["""pixel_values"""]
									self.assertListEqual(arg_names[:1]     ,					__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     int:
						SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     int:
						SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    int     ,					__UpperCAmelCase					:    int     ,					__UpperCAmelCase					:    List[Any]     ,					__UpperCAmelCase					:    Optional[Any]							) ->     List[Any]:
						# make masks reproducible
						np.random.seed(2							)

						SCREAMING_SNAKE_CASE__										=   int((pt_model.config.image_size // pt_model.config.patch_size) ** 2							)
						SCREAMING_SNAKE_CASE__										=   np.random.uniform(size=(self.model_tester.batch_size, num_patches)							)
						SCREAMING_SNAKE_CASE__										=   torch.from_numpy(__UpperCAmelCase							)

						# Add `noise` argument.
						# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
						SCREAMING_SNAKE_CASE__										=   pt_noise

						super().check_pt_tf_models(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     int:
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE__										=   model_class(__UpperCAmelCase							)
									model.to(__UpperCAmelCase							)
									model.eval()
									# make random mask reproducible
									torch.manual_seed(2							)
									with torch.no_grad():
												SCREAMING_SNAKE_CASE__										=   model(**self._prepare_for_class(__UpperCAmelCase     ,					__UpperCAmelCase							)							)

									SCREAMING_SNAKE_CASE__										=   outputs[0].cpu().numpy()
									SCREAMING_SNAKE_CASE__										=   0

									with tempfile.TemporaryDirectory() as tmpdirname:
												model.save_pretrained(__UpperCAmelCase							)
												SCREAMING_SNAKE_CASE__										=   model_class.from_pretrained(__UpperCAmelCase							)
												model.to(__UpperCAmelCase							)
												# make random mask reproducible
												torch.manual_seed(2							)
												with torch.no_grad():
															SCREAMING_SNAKE_CASE__										=   model(**self._prepare_for_class(__UpperCAmelCase     ,					__UpperCAmelCase							)							)

												# Make sure we don't have nans
												SCREAMING_SNAKE_CASE__										=   after_outputs[0].cpu().numpy()
												SCREAMING_SNAKE_CASE__										=   0
												SCREAMING_SNAKE_CASE__										=   np.amax(np.abs(out_a - out_a							)							)
												self.assertLessEqual(__UpperCAmelCase     ,					1e-5							)

			@unittest.skip(
			    reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     Union[str, Any]:
						pass

			@unittest.skip(
			    reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     Optional[int]:
						pass

			@unittest.skip(
			    reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Union[str, Any]:
						pass

			@unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     List[str]:
						pass

			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""							)
			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     List[Any]:
						pass





			@slow
			def        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     Union[str, Any]:
						for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE__										=   ViTMAEModel.from_pretrained(__UpperCAmelCase							)
									self.assertIsNotNone(__UpperCAmelCase							)

def 						A							(  ):




			'''simple docstring'''


			SCREAMING_SNAKE_CASE__										=   Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
			return image





@require_torch
@require_vision
class  lowerCamelCase    (unittest.TestCase      ):

			@cached_property
			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     str:
						return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""							) if is_vision_available() else None





			@slow
			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     List[str]:
						# make random mask reproducible across the PT and TF model
						np.random.seed(2							)

						SCREAMING_SNAKE_CASE__										=   ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""							).to(__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   self.default_image_processor
						SCREAMING_SNAKE_CASE__										=   prepare_img()
						SCREAMING_SNAKE_CASE__										=   image_processor(images=__UpperCAmelCase     ,					return_tensors="""pt"""							).to(__UpperCAmelCase							)

						# prepare a noise vector that will be also used for testing the TF model
						# (this way we can ensure that the PT and TF models operate on the same inputs)
						SCREAMING_SNAKE_CASE__										=   ViTMAEConfig()
						SCREAMING_SNAKE_CASE__										=   int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2							)
						SCREAMING_SNAKE_CASE__										=   np.random.uniform(size=(1, num_patches)							)

						# forward pass
						with torch.no_grad():
									SCREAMING_SNAKE_CASE__										=   model(**__UpperCAmelCase     ,					noise=torch.from_numpy(__UpperCAmelCase							).to(device=__UpperCAmelCase							)							)

						# verify the logits
						SCREAMING_SNAKE_CASE__										=   torch.Size((1, 1_9_6, 7_6_8)							)
						self.assertEqual(outputs.logits.shape     ,					__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   torch.tensor(
						    [[-0.0_548, -1.7_023, -0.9_325], [0.3_721, -0.5_670, -0.2_233], [0.8_235, -1.3_878, -0.3_524]]							)

						self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]     ,					expected_slice.to(__UpperCAmelCase							)     ,					atol=1e-4							)							)


