


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:						Any   				=       logging.get_logger(__name__)






def   _lowercase(					__a				:						Tuple     ):
							a__      =ASTConfig()

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

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

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

							return config






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

							return name






def   _lowercase(					__a				:						Any    ,     __a				:						Tuple     ):
							for key in orig_state_dict.copy().keys():
														a__      =orig_state_dict.pop(__a     )

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

							return orig_state_dict






def   _lowercase(					__a				:						Dict     ):
							a__      =[
							    '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   _lowercase(					__a				:						str    ,     __a				:						List[Any]    ,     __a				:						Any=False     ):
							a__      =get_audio_spectrogram_transformer_config(__a     )

							a__      ={
							    '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
							a__      =model_name_to_url[model_name]
							a__      =torch.hub.load_state_dict_from_url(__a    ,     map_location='cpu'     )
							# remove some keys
							remove_keys(__a     )
							# rename some keys
							a__      =convert_state_dict(__a    ,     __a     )

							# load 🤗 model
							a__      =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
							a__      =-4.2_67_73_93 if 'speech-commands' not in model_name else -6.84_59_78
							a__      =4.5_68_99_74 if 'speech-commands' not in model_name else 5.5_65_45_26
							a__      =1024 if 'speech-commands' not in model_name else 128
							a__      =ASTFeatureExtractor(mean=__a    ,     std=__a    ,     max_length=__a     )

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

														a__    ,     a__      =torchaudio.load(__a     )
														a__      =waveform.squeeze().numpy()

							a__      =feature_extractor(__a    ,     sampling_rate=1_6000    ,     return_tensors='pt'     )

							# forward pass
							a__      =model(**__a     )
							a__      =outputs.logits

							if model_name == "ast-finetuned-audioset-10-10-0.4593":
														a__      =torch.tensor([-0.87_60, -7.00_42, -8.66_02]     )
							elif model_name == "ast-finetuned-audioset-10-10-0.450":
														a__      =torch.tensor([-1.19_86, -7.09_03, -8.27_18]     )
							elif model_name == "ast-finetuned-audioset-10-10-0.448":
														a__      =torch.tensor([-2.61_28, -8.00_80, -9.43_44]     )
							elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
														a__      =torch.tensor([-1.50_80, -7.45_34, -8.89_17]     )
							elif model_name == "ast-finetuned-audioset-12-12-0.447":
														a__      =torch.tensor([-0.50_50, -6.58_33, -8.08_43]     )
							elif model_name == "ast-finetuned-audioset-14-14-0.443":
														a__      =torch.tensor([-0.38_26, -7.03_36, -8.24_13]     )
							elif model_name == "ast-finetuned-audioset-16-16-0.442":
														a__      =torch.tensor([-1.21_13, -6.91_01, -8.34_70]     )
							elif model_name == "ast-finetuned-speech-commands-v2":
														a__      =torch.tensor([6.15_89, -8.05_66, -8.79_84]     )
							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:						Union[str, Any]   				=       parser.parse_args()
	convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





from typing import List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

_lowerCAmelCase:						Any   				=       {
    'huggingface/autoformer-tourism-monthly': 'https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json',
}




class       lowercase_					(lowercase__   ):
				snake_case	    ='autoformer'
				snake_case	    ={
				    'hidden_size': 'd_model',
				    'num_attention_heads': 'encoder_attention_heads',
				    'num_hidden_layers': 'encoder_layers',
				}





				def __init__(				self       ,  lowercase_ = None       ,  lowercase_ = None       ,  lowercase_ = "student_t"       ,  lowercase_ = "nll"       ,  lowercase_ = 1       ,  lowercase_ = [1, 2, 3, 4, 5, 6, 7]       ,  lowercase_ = True       ,  lowercase_ = 0       ,  lowercase_ = 0       ,  lowercase_ = 0       ,  lowercase_ = 0       ,  lowercase_ = None       ,  lowercase_ = None       ,  lowercase_ = 64       ,  lowercase_ = 2       ,  lowercase_ = 2       ,  lowercase_ = 2       ,  lowercase_ = 2       ,  lowercase_ = 32       ,  lowercase_ = 32       ,  lowercase_ = "gelu"       ,  lowercase_ = 0.1       ,  lowercase_ = 0.1       ,  lowercase_ = 0.1       ,  lowercase_ = 0.1       ,  lowercase_ = 0.1       ,  lowercase_ = 100       ,  lowercase_ = 0.02       ,  lowercase_ = True       ,  lowercase_=True       ,  lowercase_ = 10       ,  lowercase_ = 25       ,  lowercase_ = 3       ,  **lowercase_       ,  )  -> Union[str, Any]:
											# time series specific configuration
											a__      =prediction_length
											a__      =context_length if context_length is not None else prediction_length
											a__      =distribution_output
											a__      =loss
											a__      =input_size
											a__      =num_time_features
											a__      =lags_sequence
											a__      =scaling
											a__      =num_dynamic_real_features
											a__      =num_static_real_features
											a__      =num_static_categorical_features
											if cardinality is not None and num_static_categorical_features > 0:
																		if len(lowercase_) != num_static_categorical_features:
																									raise ValueError(
																									    'The cardinality should be a list of the same length as `num_static_categorical_features`')
																		a__      =cardinality
											else:
																		a__      =[0]
											if embedding_dimension is not None and num_static_categorical_features > 0:
																		if len(lowercase_) != num_static_categorical_features:
																									raise ValueError(
																									    'The embedding dimension should be a list of the same length as `num_static_categorical_features`')
																		a__      =embedding_dimension
											else:
																		a__      =[min(50       ,  (cat + 1) // 2) for cat in self.cardinality]
											a__      =num_parallel_samples

											# Transformer architecture configuration
											a__      =input_size * len(self.lags_sequence) + self._number_of_features
											a__      =d_model
											a__      =encoder_attention_heads
											a__      =decoder_attention_heads
											a__      =encoder_ffn_dim
											a__      =decoder_ffn_dim
											a__      =encoder_layers
											a__      =decoder_layers

											a__      =dropout
											a__      =attention_dropout
											a__      =activation_dropout
											a__      =encoder_layerdrop
											a__      =decoder_layerdrop

											a__      =activation_function
											a__      =init_std

											a__      =use_cache

											# Autoformer
											a__      =label_length
											a__      =moving_average
											a__      =autocorrelation_factor

											super().__init__(is_encoder_decoder=lowercase_       ,  **lowercase_)





				@property
				def 						__UpperCamelCase    (				self)  -> int:
											return (
											    sum(self.embedding_dimension)
											    + self.num_dynamic_real_features
											    + self.num_time_features
											    + self.num_static_real_features
											    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
											)

