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  :     Optional[int]       =  logging.get_logger(__name__)





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



				__lowercase 			=							ASTConfig()

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

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

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

				return config





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



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

				return name





def 				UpperCAmelCase						(    lowercase							,			lowercase							):







				"""simple docstring"""



				for key in orig_state_dict.copy().keys():
								__lowercase 			=							orig_state_dict.pop(lowercase							)

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

				return orig_state_dict





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



				__lowercase 			=							[
				    '''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 				UpperCAmelCase						(    lowercase							,			lowercase							,			lowercase=False							):







				"""simple docstring"""



				__lowercase 			=							get_audio_spectrogram_transformer_config(lowercase							)

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

				# load 🤗 model
				__lowercase 			=							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
				__lowercase 			=							-4.2677393 if '''speech-commands''' not in model_name else -6.845978
				__lowercase 			=							4.5689974 if '''speech-commands''' not in model_name else 5.5654526
				__lowercase 			=							1024 if '''speech-commands''' not in model_name else 128
				__lowercase 			=							ASTFeatureExtractor(mean=lowercase							,			std=lowercase							,			max_length=lowercase							)

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

								__lowercase			,     __lowercase 			=							torchaudio.load(lowercase							)
								__lowercase 			=							waveform.squeeze().numpy()

				__lowercase 			=							feature_extractor(lowercase							,			sampling_rate=16000							,			return_tensors='''pt'''							)

				# forward pass
				__lowercase 			=							model(**lowercase							)
				__lowercase 			=							outputs.logits

				if model_name == "ast-finetuned-audioset-10-10-0.4593":
								__lowercase 			=							torch.tensor([-0.8760, -7.0042, -8.6602]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.450":
								__lowercase 			=							torch.tensor([-1.1986, -7.0903, -8.2718]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.448":
								__lowercase 			=							torch.tensor([-2.6128, -8.0080, -9.4344]							)
				elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
								__lowercase 			=							torch.tensor([-1.5080, -7.4534, -8.8917]							)
				elif model_name == "ast-finetuned-audioset-12-12-0.447":
								__lowercase 			=							torch.tensor([-0.5050, -6.5833, -8.0843]							)
				elif model_name == "ast-finetuned-audioset-14-14-0.443":
								__lowercase 			=							torch.tensor([-0.3826, -7.0336, -8.2413]							)
				elif model_name == "ast-finetuned-audioset-16-16-0.442":
								__lowercase 			=							torch.tensor([-1.2113, -6.9101, -8.3470]							)
				elif model_name == "ast-finetuned-speech-commands-v2":
								__lowercase 			=							torch.tensor([6.1589, -8.0566, -8.7984]							)
				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__":
		__a  :     Dict       =  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[Any]       =  parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__a  :     Optional[int]       =  logging.get_logger(__name__)

__a  :     Tuple       =  {
    """facebook/data2vec-text-base""": """https://huggingface.co/data2vec/resolve/main/config.json""",
}




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			__a  : Dict			 =						'''data2vec-text'''





			def __init__(   self      ,   lowerCAmelCase__=3_05_22      ,   lowerCAmelCase__=7_68      ,   lowerCAmelCase__=12      ,   lowerCAmelCase__=12      ,   lowerCAmelCase__=30_72      ,   lowerCAmelCase__="gelu"      ,   lowerCAmelCase__=0.1      ,   lowerCAmelCase__=0.1      ,   lowerCAmelCase__=5_12      ,   lowerCAmelCase__=2      ,   lowerCAmelCase__=0.02      ,   lowerCAmelCase__=1E-12      ,   lowerCAmelCase__=1      ,   lowerCAmelCase__=0      ,   lowerCAmelCase__=2      ,   lowerCAmelCase__="absolute"      ,   lowerCAmelCase__=True      ,   lowerCAmelCase__=None      ,   **lowerCAmelCase__      ,   )     ->					Any:

							'''simple docstring'''


							super().__init__(pad_token_id=lowerCAmelCase__      ,   bos_token_id=lowerCAmelCase__      ,   eos_token_id=lowerCAmelCase__      ,   **lowerCAmelCase__						)

							__lowercase 			=							vocab_size
							__lowercase 			=							hidden_size
							__lowercase 			=							num_hidden_layers
							__lowercase 			=							num_attention_heads
							__lowercase 			=							hidden_act
							__lowercase 			=							intermediate_size
							__lowercase 			=							hidden_dropout_prob
							__lowercase 			=							attention_probs_dropout_prob
							__lowercase 			=							max_position_embeddings
							__lowercase 			=							type_vocab_size
							__lowercase 			=							initializer_range
							__lowercase 			=							layer_norm_eps
							__lowercase 			=							position_embedding_type
							__lowercase 			=							use_cache
							__lowercase 			=							classifier_dropout




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			@property
			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Mapping[str, Mapping[int, str]]:

							'''simple docstring'''


							if self.task == "multiple-choice":
											__lowercase 			=							{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
							else:
											__lowercase 			=							{0: '''batch''', 1: '''sequence'''}
							return OrderedDict(
							    [
							        ('''input_ids''', dynamic_axis),
							        ('''attention_mask''', dynamic_axis),
							    ]						)