





import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
_a				:	List[Any]           =       logging.get_logger(__name__)



def  UpperCamelCase__     (   _A:    int				):




      '''simple docstring'''







      __lowerCamelCase				   =       ASTConfig()

      if "10-10" in model_name:
            pass
      elif "speech-commands" in model_name:
            __lowerCamelCase				   =       128
      elif "12-12" in model_name:
            __lowerCamelCase				   =       12
            __lowerCamelCase				   =       12
      elif "14-14" in model_name:
            __lowerCamelCase				   =       14
            __lowerCamelCase				   =       14
      elif "16-16" in model_name:
            __lowerCamelCase				   =       16
            __lowerCamelCase				   =       16
      else:
            raise ValueError("""Model not supported"""				)

      __lowerCamelCase				   =       "huggingface/label-files"
      if "speech-commands" in model_name:
            __lowerCamelCase				   =       35
            __lowerCamelCase				   =       "speech-commands-v2-id2label.json"
      else:
            __lowerCamelCase				   =       527
            __lowerCamelCase				   =       "audioset-id2label.json"

      __lowerCamelCase				   =       json.load(open(hf_hub_download(__A							,    __A							,    repo_type="""dataset"""				)							,    """r"""				)				)
      __lowerCamelCase				   =       {int(__A				): v for k, v in idalabel.items()}
      __lowerCamelCase				   =       idalabel
      __lowerCamelCase				   =       {v: k for k, v in idalabel.items()}

      return config



def  UpperCamelCase__     (   _A:    Union[str, Any]				):




      '''simple docstring'''







      if "module.v" in name:
            __lowerCamelCase				   =       name.replace("""module.v"""							,    """audio_spectrogram_transformer"""				)
      if "cls_token" in name:
            __lowerCamelCase				   =       name.replace("""cls_token"""							,    """embeddings.cls_token"""				)
      if "dist_token" in name:
            __lowerCamelCase				   =       name.replace("""dist_token"""							,    """embeddings.distillation_token"""				)
      if "pos_embed" in name:
            __lowerCamelCase				   =       name.replace("""pos_embed"""							,    """embeddings.position_embeddings"""				)
      if "patch_embed.proj" in name:
            __lowerCamelCase				   =       name.replace("""patch_embed.proj"""							,    """embeddings.patch_embeddings.projection"""				)
      # transformer blocks
      if "blocks" in name:
            __lowerCamelCase				   =       name.replace("""blocks"""							,    """encoder.layer"""				)
      if "attn.proj" in name:
            __lowerCamelCase				   =       name.replace("""attn.proj"""							,    """attention.output.dense"""				)
      if "attn" in name:
            __lowerCamelCase				   =       name.replace("""attn"""							,    """attention.self"""				)
      if "norm1" in name:
            __lowerCamelCase				   =       name.replace("""norm1"""							,    """layernorm_before"""				)
      if "norm2" in name:
            __lowerCamelCase				   =       name.replace("""norm2"""							,    """layernorm_after"""				)
      if "mlp.fc1" in name:
            __lowerCamelCase				   =       name.replace("""mlp.fc1"""							,    """intermediate.dense"""				)
      if "mlp.fc2" in name:
            __lowerCamelCase				   =       name.replace("""mlp.fc2"""							,    """output.dense"""				)
      # final layernorm
      if "audio_spectrogram_transformer.norm" in name:
            __lowerCamelCase				   =       name.replace("""audio_spectrogram_transformer.norm"""							,    """audio_spectrogram_transformer.layernorm"""				)
      # classifier head
      if "module.mlp_head.0" in name:
            __lowerCamelCase				   =       name.replace("""module.mlp_head.0"""							,    """classifier.layernorm"""				)
      if "module.mlp_head.1" in name:
            __lowerCamelCase				   =       name.replace("""module.mlp_head.1"""							,    """classifier.dense"""				)

      return name



def  UpperCamelCase__     (   _A:    Optional[Any]							,    _A:    Union[str, Any]				):




      '''simple docstring'''







      for key in orig_state_dict.copy().keys():
            __lowerCamelCase				   =       orig_state_dict.pop(__A				)

            if "qkv" in key:
                  __lowerCamelCase				   =       key.split("""."""				)
                  __lowerCamelCase				   =       int(key_split[3]				)
                  __lowerCamelCase				   =       config.hidden_size
                  if "weight" in key:
                        __lowerCamelCase				   =       val[:dim, :]
                        __lowerCamelCase				   =       val[dim : dim * 2, :]
                        __lowerCamelCase				   =       val[-dim:, :]
                  else:
                        __lowerCamelCase				   =       val[:dim]
                        __lowerCamelCase				   =       val[dim : dim * 2]
                        __lowerCamelCase				   =       val[-dim:]
            else:
                  __lowerCamelCase				   =       val

      return orig_state_dict



def  UpperCamelCase__     (   _A:    Dict				):




      '''simple docstring'''







      __lowerCamelCase				   =       [
          "module.v.head.weight",
          "module.v.head.bias",
          "module.v.head_dist.weight",
          "module.v.head_dist.bias",
      ]
      for k in ignore_keys:
            state_dict.pop(__A							,    __A				)



@torch.no_grad()
def  UpperCamelCase__     (   _A:    int							,    _A:    Union[str, Any]							,    _A:    List[str]=False				):




      '''simple docstring'''







      __lowerCamelCase				   =       get_audio_spectrogram_transformer_config(__A				)

      __lowerCamelCase				   =       {
          "ast-finetuned-audioset-10-10-0.4593": (
              "https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1"
          ),
          "ast-finetuned-audioset-10-10-0.450": (
              "https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1"
          ),
          "ast-finetuned-audioset-10-10-0.448": (
              "https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1"
          ),
          "ast-finetuned-audioset-10-10-0.448-v2": (
              "https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1"
          ),
          "ast-finetuned-audioset-12-12-0.447": (
              "https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1"
          ),
          "ast-finetuned-audioset-14-14-0.443": (
              "https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1"
          ),
          "ast-finetuned-audioset-16-16-0.442": (
              "https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1"
          ),
          "ast-finetuned-speech-commands-v2": (
              "https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1"
          ),
      }

      # load original state_dict
      __lowerCamelCase				   =       model_name_to_url[model_name]
      __lowerCamelCase				   =       torch.hub.load_state_dict_from_url(__A							,    map_location="""cpu"""				)
      # remove some keys
      remove_keys(__A				)
      # rename some keys
      __lowerCamelCase				   =       convert_state_dict(__A							,    __A				)

      # load 🤗 model
      __lowerCamelCase				   =       ASTForAudioClassification(__A				)
      model.eval()

      model.load_state_dict(__A				)

      # verify outputs on dummy input
      # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
      __lowerCamelCase				   =       -4.267_7393 if "speech-commands" not in model_name else -6.84_5978
      __lowerCamelCase				   =       4.568_9974 if "speech-commands" not in model_name else 5.565_4526
      __lowerCamelCase				   =       1024 if "speech-commands" not in model_name else 128
      __lowerCamelCase				   =       ASTFeatureExtractor(mean=__A							,    std=__A							,    max_length=__A				)

      if "speech-commands" in model_name:
            __lowerCamelCase				   =       load_dataset("""speech_commands"""							,    """v0.02"""							,    split="""validation"""				)
            __lowerCamelCase				   =       dataset[0]["audio"]["array"]
      else:
            __lowerCamelCase				   =       hf_hub_download(
                repo_id="""nielsr/audio-spectogram-transformer-checkpoint"""							,    filename="""sample_audio.flac"""							,    repo_type="""dataset"""							,    )

            __lowerCamelCase				   =       torchaudio.load(__A				)
            __lowerCamelCase				   =       waveform.squeeze().numpy()

      __lowerCamelCase				   =       feature_extractor(__A							,    sampling_rate=16000							,    return_tensors="""pt"""				)

      # forward pass
      __lowerCamelCase				   =       model(**__A				)
      __lowerCamelCase				   =       outputs.logits

      if model_name == "ast-finetuned-audioset-10-10-0.4593":
            __lowerCamelCase				   =       torch.tensor([-0.8760, -7.0042, -8.6602]				)
      elif model_name == "ast-finetuned-audioset-10-10-0.450":
            __lowerCamelCase				   =       torch.tensor([-1.1986, -7.0903, -8.2718]				)
      elif model_name == "ast-finetuned-audioset-10-10-0.448":
            __lowerCamelCase				   =       torch.tensor([-2.6128, -8.0080, -9.4344]				)
      elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
            __lowerCamelCase				   =       torch.tensor([-1.5080, -7.4534, -8.8917]				)
      elif model_name == "ast-finetuned-audioset-12-12-0.447":
            __lowerCamelCase				   =       torch.tensor([-0.5050, -6.5833, -8.0843]				)
      elif model_name == "ast-finetuned-audioset-14-14-0.443":
            __lowerCamelCase				   =       torch.tensor([-0.3826, -7.0336, -8.2413]				)
      elif model_name == "ast-finetuned-audioset-16-16-0.442":
            __lowerCamelCase				   =       torch.tensor([-1.2113, -6.9101, -8.3470]				)
      elif model_name == "ast-finetuned-speech-commands-v2":
            __lowerCamelCase				   =       torch.tensor([6.1589, -8.0566, -8.7984]				)
      else:
            raise ValueError("""Unknown model name"""				)
      if not torch.allclose(logits[0, :3]							,    __A							,    atol=1e-4				):
            raise ValueError("""Logits don't match"""				)
      print("""Looks ok!"""				)

      if pytorch_dump_folder_path is not None:
            Path(__A				).mkdir(exist_ok=__A				)
            print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''				)
            model.save_pretrained(__A				)
            print(f'''Saving feature extractor to {pytorch_dump_folder_path}'''				)
            feature_extractor.save_pretrained(__A				)

      if push_to_hub:
            print("""Pushing model and feature extractor to the hub..."""				)
            model.push_to_hub(f'''MIT/{model_name}'''				)
            feature_extractor.push_to_hub(f'''MIT/{model_name}'''				)


if __name__ == "__main__":
   _a				:	Optional[Any]           =       argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--model_name',
       default='ast-finetuned-audioset-10-10-0.4593',
       type=str,
       help='Name of the Audio Spectrogram Transformer 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.'
   )
   parser.add_argument(
       '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
   )

   _a				:	Optional[int]           =       parser.parse_args()
   convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)














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

						"""simple docstring"""



						def __init__(self:					Optional[int] ,					__UpperCAmelCase:					int ,					__UpperCAmelCase:					List[Any]=13 ,					__UpperCAmelCase:					Any=30 ,					__UpperCAmelCase:					List[str]=2 ,					__UpperCAmelCase:					Optional[int]=3 ,					__UpperCAmelCase:					Tuple=True ,					__UpperCAmelCase:					str=True ,					__UpperCAmelCase:					int=32 ,					__UpperCAmelCase:					Tuple=5 ,					__UpperCAmelCase:					Optional[Any]=4 ,					__UpperCAmelCase:					Dict=37 ,					__UpperCAmelCase:					List[Any]="gelu" ,					__UpperCAmelCase:					List[Any]=0.1 ,					__UpperCAmelCase:					Any=0.1 ,					__UpperCAmelCase:					Optional[Any]=10 ,					__UpperCAmelCase:					str=0.02 ,					__UpperCAmelCase:					Any=3 ,					__UpperCAmelCase:					Tuple=0.6 ,					__UpperCAmelCase:					str=None ,					)    ->     List[Any]:




											'''simple docstring'''



											__a					:						int      =					parent
											__a					:						Any      =					batch_size
											__a					:						int      =					image_size
											__a					:						Tuple      =					patch_size
											__a					:						str      =					num_channels
											__a					:						List[Any]      =					is_training
											__a					:						Tuple      =					use_labels
											__a					:						str      =					hidden_size
											__a					:						List[Any]      =					num_hidden_layers
											__a					:						str      =					num_attention_heads
											__a					:						Any      =					intermediate_size
											__a					:						str      =					hidden_act
											__a					:						Union[str, Any]      =					hidden_dropout_prob
											__a					:						Optional[int]      =					attention_probs_dropout_prob
											__a					:						Tuple      =					type_sequence_label_size
											__a					:						Any      =					initializer_range
											__a					:						Dict      =					mask_ratio
											__a					:						Union[str, Any]      =					scope

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



						def        UpperCAmelCase__					(self:					Any   )    ->     str:




											'''simple docstring'''



											__a					:						Dict      =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

											__a					:						Optional[Any]      =					None
											if self.use_labels:
																__a					:						int      =					ids_tensor([self.batch_size] ,					self.type_sequence_label_size   )

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

											return config, pixel_values, labels



						def        UpperCAmelCase__					(self:					Union[str, Any]   )    ->     Optional[int]:




											'''simple docstring'''



											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        UpperCAmelCase__					(self:					List[Any] ,					__UpperCAmelCase:					str ,					__UpperCAmelCase:					Tuple ,					__UpperCAmelCase:					List[str]   )    ->     List[str]:




											'''simple docstring'''



											__a					:						str      =					ViTMAEModel(config=__UpperCAmelCase   )
											model.to(__UpperCAmelCase   )
											model.eval()
											__a					:						int      =					model(__UpperCAmelCase   )
											self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, self.seq_length, self.hidden_size)   )



						def        UpperCAmelCase__					(self:					Tuple ,					__UpperCAmelCase:					Dict ,					__UpperCAmelCase:					Tuple ,					__UpperCAmelCase:					Dict   )    ->     Dict:




											'''simple docstring'''



											__a					:						Optional[Any]      =					ViTMAEForPreTraining(__UpperCAmelCase   )
											model.to(__UpperCAmelCase   )
											model.eval()
											__a					:						List[str]      =					model(__UpperCAmelCase   )
											__a					:						Union[str, Any]      =					(self.image_size // self.patch_size) ** 2
											__a					:						Optional[int]      =					self.patch_size**2 * self.num_channels
											self.parent.assertEqual(result.logits.shape ,					(self.batch_size, num_patches, expected_num_channels)   )

											# test greyscale images
											__a					:						Optional[Any]      =					1
											__a					:						int      =					ViTMAEForPreTraining(__UpperCAmelCase   )
											model.to(__UpperCAmelCase   )
											model.eval()
											__a					:						Dict      =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
											__a					:						Optional[int]      =					model(__UpperCAmelCase   )
											__a					:						List[Any]      =					self.patch_size**2
											self.parent.assertEqual(result.logits.shape ,					(self.batch_size, num_patches, expected_num_channels)   )



						def        UpperCAmelCase__					(self:					Optional[int]   )    ->     Optional[Any]:




											'''simple docstring'''



											__a					:						List[Any]      =					self.prepare_config_and_inputs()
											__a      ,			__a      ,			__a					:						Dict      =					config_and_inputs
											__a					:						Union[str, Any]      =					{"pixel_values": pixel_values}
											return config, inputs_dict




@require_torch
class 							snake_case_       (    __UpperCamelCase		,							__UpperCamelCase		,							unittest.TestCase      ):

						"""simple docstring"""



						snake_case__        =      (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
						snake_case__        =      {"""feature-extraction""": ViTMAEModel} if is_torch_available() else {}

						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False



						def        UpperCAmelCase__					(self:					int   )    ->     Dict:




											'''simple docstring'''



											__a					:						List[str]      =					ViTMAEModelTester(self   )
											__a					:						str      =					ConfigTester(self ,					config_class=__UpperCAmelCase ,					has_text_modality=__UpperCAmelCase ,					hidden_size=37   )



						def        UpperCAmelCase__					(self:					Tuple   )    ->     Dict:




											'''simple docstring'''



											self.config_tester.run_common_tests()



						@unittest.skip(reason="ViTMAE does not use inputs_embeds"   )
						def        UpperCAmelCase__					(self:					str   )    ->     Optional[Any]:




											'''simple docstring'''



											pass



						def        UpperCAmelCase__					(self:					Dict   )    ->     str:




											'''simple docstring'''



											__a      ,			__a					:						Optional[Any]      =					self.model_tester.prepare_config_and_inputs_for_common()

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



						def        UpperCAmelCase__					(self:					Dict   )    ->     Tuple:




											'''simple docstring'''



											__a      ,			__a					:						Union[str, Any]      =					self.model_tester.prepare_config_and_inputs_for_common()

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

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



						def        UpperCAmelCase__					(self:					Any   )    ->     Optional[int]:




											'''simple docstring'''



											__a					:						Union[str, Any]      =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					List[str]   )    ->     Tuple:




											'''simple docstring'''



											__a					:						List[Any]      =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					Dict ,					__UpperCAmelCase:					int ,					__UpperCAmelCase:					Dict ,					__UpperCAmelCase:					str   )    ->     Dict:




											'''simple docstring'''



											np.random.seed(2   )

											__a					:						Union[str, Any]      =					int((pt_model.config.image_size // pt_model.config.patch_size) ** 2   )
											__a					:						List[Any]      =					np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )
											__a					:						Tuple      =					torch.from_numpy(__UpperCAmelCase   )

											# Add `noise` argument.
											# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
											__a					:						Union[str, Any]      =					pt_noise

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



						def        UpperCAmelCase__					(self:					str   )    ->     Dict:




											'''simple docstring'''



											__a      ,			__a					:						Dict      =					self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																__a					:						Union[str, Any]      =					model_class(__UpperCAmelCase   )
																model.to(__UpperCAmelCase   )
																model.eval()
																# make random mask reproducible
																torch.manual_seed(2   )
																with torch.no_grad():
																					__a					:						Any      =					model(**self._prepare_for_class(__UpperCAmelCase ,					__UpperCAmelCase   )   )

																__a					:						Dict      =					outputs[0].cpu().numpy()
																__a					:						str      =					0

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

																					# Make sure we don't have nans
																					__a					:						Dict      =					after_outputs[0].cpu().numpy()
																					__a					:						str      =					0
																					__a					:						str      =					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\n    to get deterministic results."   )
						def        UpperCAmelCase__					(self:					Union[str, Any]   )    ->     Any:




											'''simple docstring'''



											pass



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




											'''simple docstring'''



											pass



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




											'''simple docstring'''



											pass



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




											'''simple docstring'''



											pass



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




											'''simple docstring'''



											pass



						@slow
						def        UpperCAmelCase__					(self:					Tuple   )    ->     int:




											'''simple docstring'''



											for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__a					:						Tuple      =					ViTMAEModel.from_pretrained(__UpperCAmelCase   )
																self.assertIsNotNone(__UpperCAmelCase   )







def     a_   ()				->						Optional[Any]:



					"""simple docstring"""
					__a					:						List[str]      =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
					return image




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

						"""simple docstring"""



						@cached_property
						def        UpperCAmelCase__					(self:					Optional[int]   )    ->     List[str]:




											'''simple docstring'''



											return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"   ) if is_vision_available() else None



						@slow
						def        UpperCAmelCase__					(self:					str   )    ->     Union[str, Any]:




											'''simple docstring'''



											np.random.seed(2   )

											__a					:						int      =					ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"   ).to(__UpperCAmelCase   )

											__a					:						Dict      =					self.default_image_processor
											__a					:						List[str]      =					prepare_img()
											__a					:						Tuple      =					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)
											__a					:						Optional[Any]      =					ViTMAEConfig()
											__a					:						str      =					int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2   )
											__a					:						int      =					np.random.uniform(size=(1, num_patches)   )

											# forward pass
											with torch.no_grad():
																__a					:						Union[str, Any]      =					model(**__UpperCAmelCase ,					noise=torch.from_numpy(__UpperCAmelCase   ).to(device=__UpperCAmelCase   )   )

											# verify the logits
											__a					:						Optional[Any]      =					torch.Size((1, 196, 768)   )
											self.assertEqual(outputs.logits.shape ,					__UpperCAmelCase   )

											__a					:						Any      =					torch.tensor(
											    [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]   )

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





