





"""simple docstring"""




import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    Tuple					) ->str:






	'''simple docstring'''






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

	return name






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    List[str]    ,       _lowercase						:    Any					) ->Tuple:






	'''simple docstring'''






	for key in orig_state_dict.copy().keys():
		a      :		Tuple	=				orig_state_dict.pop(_lowercase					)

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

		else:
			a      :		Optional[int]	=				val

	return orig_state_dict






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    Tuple    ,       _lowercase						:    List[str]					) ->Any:






	'''simple docstring'''






	a      :		Any	=				ViTMAEConfig()
	if "large" in checkpoint_url:
		a      :		Optional[Any]	=				1024
		a      :		Optional[Any]	=				4096
		a      :		Any	=				24
		a      :		Union[str, Any]	=				16
	elif "huge" in checkpoint_url:
		a      :		Optional[int]	=				14
		a      :		Union[str, Any]	=				1280
		a      :		str	=				5120
		a      :		Tuple	=				32
		a      :		Union[str, Any]	=				16

	a      :		Optional[Any]	=				ViTMAEForPreTraining(_lowercase					)

	a      :		Tuple	=				torch.hub.load_state_dict_from_url(_lowercase    ,       map_location="cpu"					)["model"]

	a      :		Dict	=				ViTMAEImageProcessor(size=config.image_size					)

	a      :		Optional[int]	=				convert_state_dict(_lowercase    ,       _lowercase					)

	model.load_state_dict(_lowercase					)
	model.eval()

	a      :		Dict	=				"https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"

	a      :		Any	=				Image.open(requests.get(_lowercase    ,       stream=_lowercase					).raw					)
	a      :		int	=				ViTMAEImageProcessor(size=config.image_size					)
	a      :		List[str]	=				image_processor(images=_lowercase    ,       return_tensors="pt"					)

	# forward pass
	torch.manual_seed(2					)
	a      :		List[Any]	=				model(**_lowercase					)
	a      :		Optional[int]	=				outputs.logits

	if "large" in checkpoint_url:
		a      :		List[Any]	=				torch.tensor(
		    [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]					)
	elif "huge" in checkpoint_url:
		a      :		int	=				torch.tensor(
		    [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]					)
	else:
		a      :		Optional[int]	=				torch.tensor(
		    [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]					)

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

	print(F"""Saving model to {pytorch_dump_folder_path}"""					)
	model.save_pretrained(_lowercase					)

	print(F"""Saving image processor to {pytorch_dump_folder_path}"""					)
	image_processor.save_pretrained(_lowercase					)


if __name__ == "__main__":
		a						: Tuple     	=			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						: Dict     	=			parser.parse_args()
		convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)



import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_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 ASTForAudioClassification, ASTModel
      from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
          AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      )


if is_torchaudio_available():
      import torchaudio

      from transformers import ASTFeatureExtractor


class      lowerCamelCase    :
  """simple docstring"""







  def __init__(							self       :				List[Any]     ,   __magic_name__       :				Any     ,   __magic_name__       :				List[Any]=13     ,   __magic_name__       :				List[Any]=2     ,   __magic_name__       :				Tuple=24     ,   __magic_name__       :				List[str]=16     ,   __magic_name__       :				Dict=True     ,   __magic_name__       :				List[Any]=True     ,   __magic_name__       :				Optional[int]=32     ,   __magic_name__       :				Tuple=5     ,   __magic_name__       :				int=4     ,   __magic_name__       :				Tuple=37     ,   __magic_name__       :				List[str]="gelu"     ,   __magic_name__       :				Tuple=0.1     ,   __magic_name__       :				Tuple=0.1     ,   __magic_name__       :				Union[str, Any]=10     ,   __magic_name__       :				Tuple=0.02     ,   __magic_name__       :				Tuple=None     ,   __magic_name__       :				Any=2     ,   __magic_name__       :				Dict=2     ,   )     ->		int:
        SCREAMING_SNAKE_CASE_      =     parent
        SCREAMING_SNAKE_CASE_      =     batch_size
        SCREAMING_SNAKE_CASE_      =     patch_size
        SCREAMING_SNAKE_CASE_      =     max_length
        SCREAMING_SNAKE_CASE_      =     num_mel_bins
        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_      =     scope
        SCREAMING_SNAKE_CASE_      =     frequency_stride
        SCREAMING_SNAKE_CASE_      =     time_stride

        # in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
        SCREAMING_SNAKE_CASE_      =     (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
        SCREAMING_SNAKE_CASE_      =     (self.max_length - self.patch_size) // self.time_stride + 1
        SCREAMING_SNAKE_CASE_      =     frequency_out_dimension * time_out_dimension
        SCREAMING_SNAKE_CASE_      =     num_patches + 2







  def       __A   (							self       :				Any      )     ->		Any:
        SCREAMING_SNAKE_CASE_      =     floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]      )

        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, input_values, labels







  def       __A   (							self       :				Any      )     ->		Dict:
        return ASTConfig(
            patch_size=self.patch_size     ,   max_length=self.max_length     ,   num_mel_bins=self.num_mel_bins     ,   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=__magic_name__     ,   initializer_range=self.initializer_range     ,   frequency_stride=self.frequency_stride     ,   time_stride=self.time_stride     ,   )







  def       __A   (							self       :				List[Any]     ,   __magic_name__       :				List[str]     ,   __magic_name__       :				str     ,   __magic_name__       :				List[str]      )     ->		Union[str, Any]:
        SCREAMING_SNAKE_CASE_      =     ASTModel(config=__magic_name__      )
        model.to(__magic_name__      )
        model.eval()
        SCREAMING_SNAKE_CASE_      =     model(__magic_name__      )
        self.parent.assertEqual(result.last_hidden_state.shape     ,   (self.batch_size, self.seq_length, self.hidden_size)      )







  def       __A   (							self       :				Tuple      )     ->		str:
        SCREAMING_SNAKE_CASE_      =     self.prepare_config_and_inputs()
        (
            (
            SCREAMING_SNAKE_CASE_
        )   ,       (
            SCREAMING_SNAKE_CASE_
        )   ,       (
            SCREAMING_SNAKE_CASE_
        )   ,       
        )      =     config_and_inputs
        SCREAMING_SNAKE_CASE_      =     {"input_values": input_values}
        return config, inputs_dict



@require_torch
class      lowerCamelCase    (SCREAMING_SNAKE_CASE__     ,					SCREAMING_SNAKE_CASE__     ,					unittest.TestCase    ):
  """simple docstring"""







  lowerCamelCase__   =  (
      (
          ASTModel,
          ASTForAudioClassification,
      )
      if is_torch_available()
      else ()
  )
  lowerCamelCase__   =  (
      {'''audio-classification''': ASTForAudioClassification, '''feature-extraction''': ASTModel}
      if is_torch_available()
      else {}
  )
  lowerCamelCase__   =  False
  lowerCamelCase__   =  False
  lowerCamelCase__   =  False
  lowerCamelCase__   =  False







  def       __A   (							self       :				Tuple     ,   __magic_name__       :				Union[str, Any]     ,   __magic_name__       :				Optional[Any]     ,   __magic_name__       :				Dict     ,   __magic_name__       :				Optional[int]     ,   __magic_name__       :				Tuple      )     ->		Tuple:
        if pipeline_test_casse_name == "AudioClassificationPipelineTests":
              return True

        return False







  def       __A   (							self       :				Optional[Any]      )     ->		Optional[Any]:
        SCREAMING_SNAKE_CASE_      =     ASTModelTester(self      )
        SCREAMING_SNAKE_CASE_      =     ConfigTester(self     ,   config_class=__magic_name__     ,   has_text_modality=__magic_name__     ,   hidden_size=37      )







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







  @unittest.skip(reason="AST does not use inputs_embeds"      )
  def       __A   (							self       :				Optional[Any]      )     ->		Tuple:
        pass







  def       __A   (							self       :				int      )     ->		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(__magic_name__      )
              self.assertIsInstance(model.get_input_embeddings()     ,   (nn.Module)      )
              SCREAMING_SNAKE_CASE_      =     model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(__magic_name__     ,   nn.Linear      )      )







  def       __A   (							self       :				List[Any]      )     ->		str:
        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(__magic_name__      )
              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_      =     ["input_values"]
              self.assertListEqual(arg_names[:1]     ,   __magic_name__      )







  def       __A   (							self       :				int      )     ->		Any:
        SCREAMING_SNAKE_CASE_      =     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__magic_name__      )







  @slow
  def       __A   (							self       :				int      )     ->		int:
        for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              SCREAMING_SNAKE_CASE_      =     ASTModel.from_pretrained(__magic_name__      )
              self.assertIsNotNone(__magic_name__      )



def    a__       (						):
      SCREAMING_SNAKE_CASE_      =     hf_hub_download(
          repo_id="nielsr/audio-spectogram-transformer-checkpoint"   ,					filename="sample_audio.flac"   ,					repo_type="dataset" )

      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     torchaudio.load(__UpperCamelCase )

      return audio, sampling_rate





@require_torch
@require_torchaudio
class      lowerCamelCase    (unittest.TestCase    ):
  """simple docstring"""







  @cached_property
  def       __A   (							self       :				List[Any]      )     ->		List[Any]:
        return (
            ASTFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"      )
            if is_torchaudio_available()
            else None
        )







  @slow
  def       __A   (							self       :				Union[str, Any]      )     ->		Optional[int]:
        SCREAMING_SNAKE_CASE_      =     self.default_feature_extractor
        SCREAMING_SNAKE_CASE_      =     ASTForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"      ).to(__magic_name__      )

        SCREAMING_SNAKE_CASE_      =     self.default_feature_extractor
        SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     prepare_audio()
        SCREAMING_SNAKE_CASE_      =     audio.squeeze().numpy()
        SCREAMING_SNAKE_CASE_      =     feature_extractor(__magic_name__     ,   sampling_rate=__magic_name__     ,   return_tensors="pt"      ).to(__magic_name__      )

        # forward pass
        with torch.no_grad():
              SCREAMING_SNAKE_CASE_      =     model(**__magic_name__      )

        # verify the logits
        SCREAMING_SNAKE_CASE_      =     torch.Size((1, 527)      )
        self.assertEqual(outputs.logits.shape     ,   __magic_name__      )

        SCREAMING_SNAKE_CASE_      =     torch.tensor([-0.8760, -7.0042, -8.6602]      ).to(__magic_name__      )

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


