
"""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()
lowerCamelCase__ :       List[str]						     =			logging.get_logger(__name__)





def __A						(   a_ :							Dict     )->							Union[str, Any]:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     List[str]				= ASTConfig()

		if "10-10" in model_name:
				pass
		elif "speech-commands" in model_name:
				SCREAMING_SNAKE_CASE       :     Dict				= 1_28
		elif "12-12" in model_name:
				SCREAMING_SNAKE_CASE       :     Tuple				= 12
				SCREAMING_SNAKE_CASE       :     int				= 12
		elif "14-14" in model_name:
				SCREAMING_SNAKE_CASE       :     List[str]				= 14
				SCREAMING_SNAKE_CASE       :     Any				= 14
		elif "16-16" in model_name:
				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= 16
				SCREAMING_SNAKE_CASE       :     str				= 16
		else:
				raise ValueError('''Model not supported'''     )

		SCREAMING_SNAKE_CASE       :     Optional[Any]				= '''huggingface/label-files'''
		if "speech-commands" in model_name:
				SCREAMING_SNAKE_CASE       :     Tuple				= 35
				SCREAMING_SNAKE_CASE       :     str				= '''speech-commands-v2-id2label.json'''
		else:
				SCREAMING_SNAKE_CASE       :     int				= 5_27
				SCREAMING_SNAKE_CASE       :     List[str]				= '''audioset-id2label.json'''

		SCREAMING_SNAKE_CASE       :     int				= json.load(open(hf_hub_download(a_		,		a_		,		repo_type='''dataset'''     )		,		'''r'''     )     )
		SCREAMING_SNAKE_CASE       :     Optional[int]				= {int(a_     ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE       :     int				= idalabel
		SCREAMING_SNAKE_CASE       :     Any				= {v: k for k, v in idalabel.items()}

		return config





def __A						(   a_ :							List[str]     )->							Any:




		'''simple docstring'''






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

		return name





def __A						(   a_ :							Tuple		,		a_ :							List[Any]     )->							Dict:




		'''simple docstring'''






		for key in orig_state_dict.copy().keys():
				SCREAMING_SNAKE_CASE       :     Optional[Any]				= orig_state_dict.pop(a_     )

				if "qkv" in key:
						SCREAMING_SNAKE_CASE       :     int				= key.split('''.'''     )
						SCREAMING_SNAKE_CASE       :     int				= int(key_split[3]     )
						SCREAMING_SNAKE_CASE       :     Dict				= config.hidden_size
						if "weight" in key:
								SCREAMING_SNAKE_CASE       :     Any				= val[:dim, :]
								SCREAMING_SNAKE_CASE       :     int				= val[dim : dim * 2, :]
								SCREAMING_SNAKE_CASE       :     str				= val[-dim:, :]
						else:
								SCREAMING_SNAKE_CASE       :     List[Any]				= val[:dim]
								SCREAMING_SNAKE_CASE       :     Any				= val[dim : dim * 2]
								SCREAMING_SNAKE_CASE       :     str				= val[-dim:]
				else:
						SCREAMING_SNAKE_CASE       :     Any				= val

		return orig_state_dict





def __A						(   a_ :							str     )->							Optional[int]:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     Optional[Any]				= [
		    '''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 __A						(   a_ :							Optional[int]		,		a_ :							str		,		a_ :							Union[str, Any]=False     )->							Tuple:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= get_audio_spectrogram_transformer_config(a_     )

		SCREAMING_SNAKE_CASE       :     Dict				= {
		    '''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
		SCREAMING_SNAKE_CASE       :     str				= model_name_to_url[model_name]
		SCREAMING_SNAKE_CASE       :     Optional[int]				= torch.hub.load_state_dict_from_url(a_		,		map_location='''cpu'''     )
		# remove some keys
		remove_keys(a_     )
		# rename some keys
		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= convert_state_dict(a_		,		a_     )

		# load 🤗 model
		SCREAMING_SNAKE_CASE       :     Union[str, Any]				= 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
		SCREAMING_SNAKE_CASE       :     str				= -4.267_7393 if '''speech-commands''' not in model_name else -6.84_5978
		SCREAMING_SNAKE_CASE       :     Dict				= 4.568_9974 if '''speech-commands''' not in model_name else 5.565_4526
		SCREAMING_SNAKE_CASE       :     int				= 10_24 if '''speech-commands''' not in model_name else 1_28
		SCREAMING_SNAKE_CASE       :     List[Any]				= ASTFeatureExtractor(mean=a_		,		std=a_		,		max_length=a_     )

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

				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= torchaudio.load(a_     )
				SCREAMING_SNAKE_CASE       :     List[str]				= waveform.squeeze().numpy()

		SCREAMING_SNAKE_CASE       :     Tuple				= feature_extractor(a_		,		sampling_rate=1_60_00		,		return_tensors='''pt'''     )

		# forward pass
		SCREAMING_SNAKE_CASE       :     List[Any]				= model(**a_     )
		SCREAMING_SNAKE_CASE       :     Optional[Any]				= outputs.logits

		if model_name == "ast-finetuned-audioset-10-10-0.4593":
				SCREAMING_SNAKE_CASE       :     Tuple				= torch.tensor([-0.8760, -7.0042, -8.6602]     )
		elif model_name == "ast-finetuned-audioset-10-10-0.450":
				SCREAMING_SNAKE_CASE       :     Any				= torch.tensor([-1.1986, -7.0903, -8.2718]     )
		elif model_name == "ast-finetuned-audioset-10-10-0.448":
				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= torch.tensor([-2.6128, -8.0080, -9.4344]     )
		elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
				SCREAMING_SNAKE_CASE       :     int				= torch.tensor([-1.5080, -7.4534, -8.8917]     )
		elif model_name == "ast-finetuned-audioset-12-12-0.447":
				SCREAMING_SNAKE_CASE       :     Any				= torch.tensor([-0.5050, -6.5833, -8.0843]     )
		elif model_name == "ast-finetuned-audioset-14-14-0.443":
				SCREAMING_SNAKE_CASE       :     Optional[Any]				= torch.tensor([-0.3826, -7.0336, -8.2413]     )
		elif model_name == "ast-finetuned-audioset-16-16-0.442":
				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= torch.tensor([-1.2113, -6.9101, -8.3470]     )
		elif model_name == "ast-finetuned-speech-commands-v2":
				SCREAMING_SNAKE_CASE       :     List[Any]				= 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__":
					lowerCamelCase__ :       List[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."
					)

					lowerCamelCase__ :       Tuple						     =			parser.parse_args()
					convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


lowerCamelCase__ :       int						     =			logging.get_logger(__name__)




class 				lowercase__(				_UpperCAmelCase		,						_UpperCAmelCase   ):
	'''simple docstring'''

	UpperCamelCase				=  """maskformer-swin"""

	UpperCamelCase				=  {
	    """num_attention_heads""": """num_heads""",
	    """num_hidden_layers""": """num_layers""",
	}





	def __init__(						self    :Optional[int]		, lowerCamelCase_    :List[Any]=2_24		, lowerCamelCase_    :Tuple=4		, lowerCamelCase_    :Optional[Any]=3		, lowerCamelCase_    :List[str]=96		, lowerCamelCase_    :int=[2, 2, 6, 2]		, lowerCamelCase_    :Union[str, Any]=[3, 6, 12, 24]		, lowerCamelCase_    :Optional[int]=7		, lowerCamelCase_    :Tuple=4.0		, lowerCamelCase_    :Tuple=True		, lowerCamelCase_    :Dict=0.0		, lowerCamelCase_    :Any=0.0		, lowerCamelCase_    :List[Any]=0.1		, lowerCamelCase_    :Dict="gelu"		, lowerCamelCase_    :Optional[int]=False		, lowerCamelCase_    :List[str]=0.0_2		, lowerCamelCase_    :Any=1E-5		, lowerCamelCase_    :Union[str, Any]=None		, lowerCamelCase_    :List[str]=None		, **lowerCamelCase_    :Union[str, Any]		, )							->			Dict:







			'''simple docstring'''




			super().__init__(**lowerCamelCase_					)

			SCREAMING_SNAKE_CASE       :     int				= image_size
			SCREAMING_SNAKE_CASE       :     Optional[Any]				= patch_size
			SCREAMING_SNAKE_CASE       :     str				= num_channels
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= embed_dim
			SCREAMING_SNAKE_CASE       :     List[Any]				= depths
			SCREAMING_SNAKE_CASE       :     List[str]				= len(lowerCamelCase_					)
			SCREAMING_SNAKE_CASE       :     str				= num_heads
			SCREAMING_SNAKE_CASE       :     Any				= window_size
			SCREAMING_SNAKE_CASE       :     List[str]				= mlp_ratio
			SCREAMING_SNAKE_CASE       :     str				= qkv_bias
			SCREAMING_SNAKE_CASE       :     Tuple				= hidden_dropout_prob
			SCREAMING_SNAKE_CASE       :     int				= attention_probs_dropout_prob
			SCREAMING_SNAKE_CASE       :     int				= drop_path_rate
			SCREAMING_SNAKE_CASE       :     Tuple				= hidden_act
			SCREAMING_SNAKE_CASE       :     Any				= use_absolute_embeddings
			SCREAMING_SNAKE_CASE       :     List[str]				= layer_norm_eps
			SCREAMING_SNAKE_CASE       :     List[str]				= initializer_range
			# we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
			# this indicates the channel dimension after the last stage of the model
			SCREAMING_SNAKE_CASE       :     int				= int(embed_dim * 2 ** (len(lowerCamelCase_					) - 1)					)
			SCREAMING_SNAKE_CASE       :     Dict				= ['''stem'''] + [f"stage{idx}" for idx in range(1		, len(lowerCamelCase_					) + 1					)]
			SCREAMING_SNAKE_CASE,					SCREAMING_SNAKE_CASE       :     int				= get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_		, out_indices=lowerCamelCase_		, stage_names=self.stage_names					)





