



"""simple docstring"""
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()
__SCREAMING_SNAKE_CASE  :						Tuple    							=      logging.get_logger(__name__)



def        lowerCAmelCase_(   lowercase_     :							Tuple							)	->       Any:
				_lowerCamelCase        =				ASTConfig()

				if "10-10" in model_name:
								pass
				elif "speech-commands" in model_name:
								_lowerCamelCase        =				1_28
				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        =				5_27
								_lowerCamelCase        =				'''audioset-id2label.json'''

				_lowerCamelCase        =				json.load(open(hf_hub_download(lowercase_						,   lowercase_						,   repo_type='''dataset'''							)						,   '''r'''							)							)
				_lowerCamelCase        =				{int(lowercase_							): v for k, v in idalabel.items()}
				_lowerCamelCase        =				idalabel
				_lowerCamelCase        =				{v: k for k, v in idalabel.items()}

				return config



def        lowerCAmelCase_(   lowercase_     :							Optional[int]							)	->       Tuple:
				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        lowerCAmelCase_(   lowercase_     :							List[str]						,   lowercase_     :							Any							)	->       str:
				for key in orig_state_dict.copy().keys():
								_lowerCamelCase        =				orig_state_dict.pop(lowercase_							)

								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        lowerCAmelCase_(   lowercase_     :							Any							)	->       Tuple:
				_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(lowercase_						,   lowercase_							)



@torch.no_grad()
def        lowerCAmelCase_(   lowercase_     :							Any						,   lowercase_     :							Any						,   lowercase_     :							Any=False							)	->       Optional[Any]:
				_lowerCamelCase        =				get_audio_spectrogram_transformer_config(lowercase_							)

				_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(lowercase_						,   map_location='''cpu'''							)
				# remove some keys
				remove_keys(lowercase_							)
				# rename some keys
				_lowerCamelCase        =				convert_state_dict(lowercase_						,   lowercase_							)

				# load 🤗 model
				_lowerCamelCase        =				ASTForAudioClassification(lowercase_							)
				model.eval()

				model.load_state_dict(lowercase_							)

				# verify outputs on dummy input
				# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
				_lowerCamelCase        =				-4.2_6_7_7_3_9_3 if '''speech-commands''' not in model_name else -6.8_4_5_9_7_8
				_lowerCamelCase        =				4.5_6_8_9_9_7_4 if '''speech-commands''' not in model_name else 5.5_6_5_4_5_2_6
				_lowerCamelCase        =				10_24 if '''speech-commands''' not in model_name else 1_28
				_lowerCamelCase        =				ASTFeatureExtractor(mean=lowercase_						,   std=lowercase_						,   max_length=lowercase_							)

				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  ,     _lowerCamelCase        =				torchaudio.load(lowercase_							)
								_lowerCamelCase        =				waveform.squeeze().numpy()

				_lowerCamelCase        =				feature_extractor(lowercase_						,   sampling_rate=1_60_00						,   return_tensors='''pt'''							)

				# forward pass
				_lowerCamelCase        =				model(**lowercase_							)
				_lowerCamelCase        =				outputs.logits

				if model_name == "ast-finetuned-audioset-10-10-0.4593":
								_lowerCamelCase        =				torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.450":
								_lowerCamelCase        =				torch.tensor([-1.1_9_8_6, -7.0_9_0_3, -8.2_7_1_8]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.448":
								_lowerCamelCase        =				torch.tensor([-2.6_1_2_8, -8.0_0_8_0, -9.4_3_4_4]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
								_lowerCamelCase        =				torch.tensor([-1.5_0_8_0, -7.4_5_3_4, -8.8_9_1_7]							)
				elif model_name == "ast-finetuned-audioset-12-12-0.447":
								_lowerCamelCase        =				torch.tensor([-0.5_0_5_0, -6.5_8_3_3, -8.0_8_4_3]							)
				elif model_name == "ast-finetuned-audioset-14-14-0.443":
								_lowerCamelCase        =				torch.tensor([-0.3_8_2_6, -7.0_3_3_6, -8.2_4_1_3]							)
				elif model_name == "ast-finetuned-audioset-16-16-0.442":
								_lowerCamelCase        =				torch.tensor([-1.2_1_1_3, -6.9_1_0_1, -8.3_4_7_0]							)
				elif model_name == "ast-finetuned-speech-commands-v2":
								_lowerCamelCase        =				torch.tensor([6.1_5_8_9, -8.0_5_6_6, -8.7_9_8_4]							)
				else:
								raise ValueError('''Unknown model name'''							)
				if not torch.allclose(logits[0, :3]						,   lowercase_						,   atol=1e-4							):
								raise ValueError('''Logits don\'t match'''							)
				print('''Looks ok!'''							)

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

				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__":
			__SCREAMING_SNAKE_CASE  :						str    							=      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.'''
			)

			__SCREAMING_SNAKE_CASE  :						Any    							=      parser.parse_args()
			convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





"""simple docstring"""
from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import tensorflow as tf

			from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
			from PIL import Image

			from transformers import ViTImageProcessor


class 							lowerCamelCase_:





					'''simple docstring'''



					def __init__(    self				,       lowerCamelCase__				,       lowerCamelCase__=1_3				,       lowerCamelCase__=3_0				,       lowerCamelCase__=2				,       lowerCamelCase__=3				,       lowerCamelCase__=True				,       lowerCamelCase__=True				,       lowerCamelCase__=3_2				,       lowerCamelCase__=2				,       lowerCamelCase__=4				,       lowerCamelCase__=3_7				,       lowerCamelCase__="gelu"				,       lowerCamelCase__=0.1				,       lowerCamelCase__=0.1				,       lowerCamelCase__=1_0				,       lowerCamelCase__=0.0_2				,       lowerCamelCase__=3				,       lowerCamelCase__=0.6				,       lowerCamelCase__=None				,       ):
									_lowerCamelCase        =				parent
									_lowerCamelCase        =				batch_size
									_lowerCamelCase        =				image_size
									_lowerCamelCase        =				patch_size
									_lowerCamelCase        =				num_channels
									_lowerCamelCase        =				is_training
									_lowerCamelCase        =				use_labels
									_lowerCamelCase        =				hidden_size
									_lowerCamelCase        =				num_hidden_layers
									_lowerCamelCase        =				num_attention_heads
									_lowerCamelCase        =				intermediate_size
									_lowerCamelCase        =				hidden_act
									_lowerCamelCase        =				hidden_dropout_prob
									_lowerCamelCase        =				attention_probs_dropout_prob
									_lowerCamelCase        =				type_sequence_label_size
									_lowerCamelCase        =				initializer_range
									_lowerCamelCase        =				mask_ratio
									_lowerCamelCase        =				scope

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

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

									_lowerCamelCase        =				self.get_config()

									return config, pixel_values, labels

					def 			snake_case__      (    self   ):
									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				,       decoder_hidden_size=self.hidden_size				,       decoder_num_hidden_layers=self.num_hidden_layers				,       decoder_num_attention_heads=self.num_attention_heads				,       decoder_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=lowerCamelCase__				,       initializer_range=self.initializer_range				,       mask_ratio=self.mask_ratio				,       )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				TFViTMAEModel(config=lowerCamelCase__   )
									_lowerCamelCase        =				model(lowerCamelCase__				,       training=lowerCamelCase__   )
									self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									_lowerCamelCase        =				TFViTMAEForPreTraining(lowerCamelCase__   )
									_lowerCamelCase        =				model(lowerCamelCase__				,       training=lowerCamelCase__   )
									# expected sequence length = num_patches
									_lowerCamelCase        =				(self.image_size // self.patch_size) ** 2
									_lowerCamelCase        =				self.patch_size**2 * self.num_channels
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, num_patches, expected_num_channels)   )

									# test greyscale images
									_lowerCamelCase        =				1
									_lowerCamelCase        =				TFViTMAEForPreTraining(lowerCamelCase__   )

									_lowerCamelCase        =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
									_lowerCamelCase        =				model(lowerCamelCase__				,       training=lowerCamelCase__   )
									_lowerCamelCase        =				self.patch_size**2
									self.parent.assertEqual(result.logits.shape				,       (self.batch_size, num_patches, expected_num_channels)   )





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.prepare_config_and_inputs()
									((_lowerCamelCase)  ,     (_lowerCamelCase)  ,     (_lowerCamelCase))        =				config_and_inputs
									_lowerCamelCase        =				{'''pixel_values''': pixel_values}
									return config, inputs_dict


@require_tf
class 							lowerCamelCase_(   A__,				A__,				unittest.TestCase    ):





					'''simple docstring'''



					lowercase__   :       Optional[Any]										= (TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
					lowercase__   :       Dict										= {'feature-extraction': TFViTMAEModel} if is_tf_available() else {}

					lowercase__   :       Optional[Any]										= False
					lowercase__   :       Union[str, Any]										= False
					lowercase__   :       str										= False
					lowercase__   :       List[str]										= False

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				TFViTMAEModelTester(self   )
									_lowerCamelCase        =				ConfigTester(self				,       config_class=lowerCamelCase__				,       has_text_modality=lowerCamelCase__				,       hidden_size=3_7   )

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

					@unittest.skip(reason='''ViTMAE does not use inputs_embeds'''   )
					def 			snake_case__      (    self   ):
									pass

					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
													_lowerCamelCase        =				model_class(lowerCamelCase__   )
													self.assertIsInstance(model.get_input_embeddings()				,       (tf.keras.layers.Layer)   )
													_lowerCamelCase        =				model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(lowerCamelCase__				,       tf.keras.layers.Layer   )   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

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

													_lowerCamelCase        =				['''pixel_values''']
													self.assertListEqual(arg_names[:1]				,       lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									# make the mask reproducible
									np.random.seed(2   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									_lowerCamelCase        =				int((config.image_size // config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )

									for model_class in self.all_model_classes:
													_lowerCamelCase        =				model_class(lowerCamelCase__   )
													_lowerCamelCase        =				self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )

													_lowerCamelCase        =				model(lowerCamelCase__				,       noise=lowerCamelCase__   )

													_lowerCamelCase        =				copy.deepcopy(self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )   )
													_lowerCamelCase        =				model(**lowerCamelCase__				,       noise=lowerCamelCase__   )
													_lowerCamelCase        =				outputs_dict[0].numpy()
													_lowerCamelCase        =				outputs_keywords[0].numpy()

													self.assertLess(np.sum(np.abs(output_dict - output_keywords   )   )				,       1e-6   )

					def 			snake_case__      (    self   ):
									# make the mask reproducible
									np.random.seed(2   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									_lowerCamelCase        =				int((config.image_size // config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )

									def prepare_numpy_arrays(lowerCamelCase__   ):
													_lowerCamelCase        =				{}
													for k, v in inputs_dict.items():
																	if tf.is_tensor(lowerCamelCase__   ):
																					_lowerCamelCase        =				v.numpy()
																	else:
																					_lowerCamelCase        =				np.array(lowerCamelCase__   )

													return inputs_np_dict

									for model_class in self.all_model_classes:
													_lowerCamelCase        =				model_class(lowerCamelCase__   )

													_lowerCamelCase        =				self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )
													_lowerCamelCase        =				prepare_numpy_arrays(lowerCamelCase__   )

													_lowerCamelCase        =				model(lowerCamelCase__				,       noise=lowerCamelCase__   )
													_lowerCamelCase        =				model(**lowerCamelCase__				,       noise=lowerCamelCase__   )
													self.assert_outputs_same(lowerCamelCase__				,       lowerCamelCase__   )

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
									# make masks reproducible
									np.random.seed(2   )

									_lowerCamelCase        =				int((tf_model.config.image_size // tf_model.config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )
									_lowerCamelCase        =				tf.constant(lowerCamelCase__   )

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

									super().check_pt_tf_models(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									# make mask reproducible
									np.random.seed(2   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									_lowerCamelCase        =				{
									    module_member
									    for model_class in self.all_model_classes
									    for module in (import_module(model_class.__module__   ),)
									    for module_member_name in dir(lowerCamelCase__   )
									    if module_member_name.endswith('''MainLayer'''   )
									    # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
									    and module_member_name[: -len('''MainLayer'''   )] == model_class.__name__[: -len('''Model'''   )]
									    for module_member in (getattr(lowerCamelCase__				,       lowerCamelCase__   ),)
									    if isinstance(lowerCamelCase__				,       lowerCamelCase__   )
									    and tf.keras.layers.Layer in module_member.__bases__
									    and getattr(lowerCamelCase__				,       '''_keras_serializable'''				,       lowerCamelCase__   )
									}

									_lowerCamelCase        =				int((config.image_size // config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )
									_lowerCamelCase        =				tf.convert_to_tensor(lowerCamelCase__   )
									inputs_dict.update({'''noise''': noise}   )

									for main_layer_class in tf_main_layer_classes:
													_lowerCamelCase        =				main_layer_class(lowerCamelCase__   )

													_lowerCamelCase        =				{
													    name: tf.keras.Input(tensor.shape[1:]				,       dtype=tensor.dtype   ) for name, tensor in inputs_dict.items()
													}

													_lowerCamelCase        =				tf.keras.Model(lowerCamelCase__				,       outputs=main_layer(lowerCamelCase__   )   )
													_lowerCamelCase        =				model(lowerCamelCase__   )

													with tempfile.TemporaryDirectory() as tmpdirname:
																	_lowerCamelCase        =				os.path.join(lowerCamelCase__				,       '''keras_model.h5'''   )
																	model.save(lowerCamelCase__   )
																	_lowerCamelCase        =				tf.keras.models.load_model(
																	    lowerCamelCase__				,       custom_objects={main_layer_class.__name__: main_layer_class}   )
																	assert isinstance(lowerCamelCase__				,       tf.keras.Model   )
																	_lowerCamelCase        =				model(lowerCamelCase__   )
																	self.assert_outputs_same(lowerCamelCase__				,       lowerCamelCase__   )

					@slow
					def 			snake_case__      (    self   ):
									# make mask reproducible
									np.random.seed(2   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									_lowerCamelCase        =				int((config.image_size // config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )

									for model_class in self.all_model_classes:
													_lowerCamelCase        =				model_class(lowerCamelCase__   )
													_lowerCamelCase        =				self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )
													_lowerCamelCase        =				model(lowerCamelCase__				,       noise=lowerCamelCase__   )

													if model_class.__name__ == "TFViTMAEModel":
																	_lowerCamelCase        =				outputs.last_hidden_state.numpy()
																	_lowerCamelCase        =				0
													else:
																	_lowerCamelCase        =				outputs.logits.numpy()
																	_lowerCamelCase        =				0

													with tempfile.TemporaryDirectory() as tmpdirname:
																	model.save_pretrained(lowerCamelCase__				,       saved_model=lowerCamelCase__   )
																	_lowerCamelCase        =				model_class.from_pretrained(lowerCamelCase__   )
																	_lowerCamelCase        =				model(lowerCamelCase__				,       noise=lowerCamelCase__   )

																	if model_class.__name__ == "TFViTMAEModel":
																					_lowerCamelCase        =				after_outputs['''last_hidden_state'''].numpy()
																					_lowerCamelCase        =				0
																	else:
																					_lowerCamelCase        =				after_outputs['''logits'''].numpy()
																					_lowerCamelCase        =				0

																	_lowerCamelCase        =				np.amax(np.abs(out_a - out_a   )   )
																	self.assertLessEqual(lowerCamelCase__				,       1e-5   )

					def 			snake_case__      (    self   ):
									# make mask reproducible
									np.random.seed(2   )

									_lowerCamelCase  ,     _lowerCamelCase        =				self.model_tester.prepare_config_and_inputs_for_common()

									_lowerCamelCase        =				int((config.image_size // config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(self.model_tester.batch_size, num_patches)   )

									for model_class in self.all_model_classes:
													_lowerCamelCase        =				model_class(lowerCamelCase__   )
													_lowerCamelCase        =				self._prepare_for_class(lowerCamelCase__				,       lowerCamelCase__   )

													_lowerCamelCase        =				model(lowerCamelCase__				,       noise=lowerCamelCase__   )
													_lowerCamelCase        =				model.get_config()
													# make sure that returned config is jsonifiable, which is required by keras
													json.dumps(lowerCamelCase__   )
													_lowerCamelCase        =				model_class.from_config(model.get_config()   )
													# make sure it also accepts a normal config
													_lowerCamelCase        =				model_class.from_config(model.config   )
													_lowerCamelCase        =				new_model(lowerCamelCase__   )  # Build model
													new_model.set_weights(model.get_weights()   )
													_lowerCamelCase        =				new_model(lowerCamelCase__				,       noise=lowerCamelCase__   )

													self.assert_outputs_same(lowerCamelCase__				,       lowerCamelCase__   )

					@unittest.skip(
					    reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results.'''   )
					def 			snake_case__      (    self   ):
									pass

					@unittest.skip(reason='''ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'''   )
					def 			snake_case__      (    self   ):
									pass





					@slow
					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				TFViTMAEModel.from_pretrained('''google/vit-base-patch16-224'''   )
									self.assertIsNotNone(lowerCamelCase__   )



def        lowerCAmelCase_(   )	->       List[Any]:
				_lowerCamelCase        =				Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
				return image




@require_tf
@require_vision
class 							lowerCamelCase_(   unittest.TestCase    ):





					'''simple docstring'''



					@cached_property
					def 			snake_case__      (    self   ):
									return ViTImageProcessor.from_pretrained('''facebook/vit-mae-base'''   ) if is_vision_available() else None





					@slow
					def 			snake_case__      (    self   ):
									# make random mask reproducible across the PT and TF model
									np.random.seed(2   )

									_lowerCamelCase        =				TFViTMAEForPreTraining.from_pretrained('''facebook/vit-mae-base'''   )

									_lowerCamelCase        =				self.default_image_processor
									_lowerCamelCase        =				prepare_img()
									_lowerCamelCase        =				image_processor(images=lowerCamelCase__				,       return_tensors='''tf'''   )

									# 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)
									_lowerCamelCase        =				ViTMAEConfig()
									_lowerCamelCase        =				int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2   )
									_lowerCamelCase        =				np.random.uniform(size=(1, num_patches)   )

									# forward pass
									_lowerCamelCase        =				model(**lowerCamelCase__				,       noise=lowerCamelCase__   )

									# verify the logits
									_lowerCamelCase        =				tf.convert_to_tensor([1, 1_9_6, 7_6_8]   )
									self.assertEqual(outputs.logits.shape				,       lowerCamelCase__   )

									_lowerCamelCase        =				tf.convert_to_tensor(
									    [[-0.0_5_4_8, -1.7_0_2_3, -0.9_3_2_5], [0.3_7_2_1, -0.5_6_7_0, -0.2_2_3_3], [0.8_2_3_5, -1.3_8_7_8, -0.3_5_2_4]]   )

									tf.debugging.assert_near(outputs.logits[0, :3, :3]				,       lowerCamelCase__				,       atol=1e-4   )
