





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





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				str:







			'''simple docstring'''

			UpperCAmelCase   		=				ASTConfig()

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

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

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

			return config





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				List[str]:







			'''simple docstring'''

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

			return name





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__ )  ->				Union[str, Any]:







			'''simple docstring'''

			for key in orig_state_dict.copy().keys():
						UpperCAmelCase   		=				orig_state_dict.pop(UpperCamelCase__ )

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

			return orig_state_dict





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				List[Any]:







			'''simple docstring'''

			UpperCAmelCase   		=				[
			    '''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(UpperCamelCase__			,			UpperCamelCase__ )





@torch.no_grad()
def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__=False )  ->				Optional[int]:







			'''simple docstring'''

			UpperCAmelCase   		=				get_audio_spectrogram_transformer_config(UpperCamelCase__ )

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

			# load 🤗 model
			UpperCAmelCase   		=				ASTForAudioClassification(UpperCamelCase__ )
			model.eval()

			model.load_state_dict(UpperCamelCase__ )

			# verify outputs on dummy input
			# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
			UpperCAmelCase   		=				-4.267_7393 if '''speech-commands''' not in model_name else -6.84_5978
			UpperCAmelCase   		=				4.568_9974 if '''speech-commands''' not in model_name else 5.565_4526
			UpperCAmelCase   		=				1024 if '''speech-commands''' not in model_name else 128
			UpperCAmelCase   		=				ASTFeatureExtractor(mean=UpperCamelCase__			,			std=UpperCamelCase__			,			max_length=UpperCamelCase__ )

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

						UpperCAmelCase		,			UpperCAmelCase   		=				torchaudio.load(UpperCamelCase__ )
						UpperCAmelCase   		=				waveform.squeeze().numpy()

			UpperCAmelCase   		=				feature_extractor(UpperCamelCase__			,			sampling_rate=1_6000			,			return_tensors='''pt''' )

			# forward pass
			UpperCAmelCase   		=				model(**UpperCamelCase__ )
			UpperCAmelCase   		=				outputs.logits

			if model_name == "ast-finetuned-audioset-10-10-0.4593":
						UpperCAmelCase   		=				torch.tensor([-0.8760, -7.0042, -8.6602] )
			elif model_name == "ast-finetuned-audioset-10-10-0.450":
						UpperCAmelCase   		=				torch.tensor([-1.1986, -7.0903, -8.2718] )
			elif model_name == "ast-finetuned-audioset-10-10-0.448":
						UpperCAmelCase   		=				torch.tensor([-2.6128, -8.0080, -9.4344] )
			elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
						UpperCAmelCase   		=				torch.tensor([-1.5080, -7.4534, -8.8917] )
			elif model_name == "ast-finetuned-audioset-12-12-0.447":
						UpperCAmelCase   		=				torch.tensor([-0.5050, -6.5833, -8.0843] )
			elif model_name == "ast-finetuned-audioset-14-14-0.443":
						UpperCAmelCase   		=				torch.tensor([-0.3826, -7.0336, -8.2413] )
			elif model_name == "ast-finetuned-audioset-16-16-0.442":
						UpperCAmelCase   		=				torch.tensor([-1.2113, -6.9101, -8.3470] )
			elif model_name == "ast-finetuned-speech-commands-v2":
						UpperCAmelCase   		=				torch.tensor([6.1589, -8.0566, -8.7984] )
			else:
						raise ValueError('''Unknown model name''' )
			if not torch.allclose(logits[0, :3]			,			UpperCamelCase__			,			atol=1E-4 ):
						raise ValueError('''Logits don\'t match''' )
			print('''Looks ok!''' )

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

			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				:    Union[str, Any]       =   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				:    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				:    int       =   logging.get_logger(__name__)

__A				:    Tuple       =   {
    "google/bigbird-roberta-base": "https://huggingface.co/google/bigbird-roberta-base/resolve/main/config.json",
    "google/bigbird-roberta-large": "https://huggingface.co/google/bigbird-roberta-large/resolve/main/config.json",
    "google/bigbird-base-trivia-itc": "https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/config.json",
    # See all BigBird models at https://huggingface.co/models?filter=big_bird
}

class 	A_       (a_			):
	UpperCAmelCase__					=			'''big_bird'''
	def __init__(    self							,					_A=5_0_3_5_8							,					_A=7_6_8							,					_A=1_2							,					_A=1_2							,					_A=3_0_7_2							,					_A="gelu_new"							,					_A=0.1							,					_A=0.1							,					_A=4_0_9_6							,					_A=2							,					_A=0.02							,					_A=1E-12							,					_A=True							,					_A=0							,					_A=1							,					_A=2							,					_A=6_6							,					_A="block_sparse"							,					_A=True							,					_A=False							,					_A=6_4							,					_A=3							,					_A=None							,					**_A							,					):






				'''simple docstring'''







				super().__init__(
				    pad_token_id=_A							,					bos_token_id=_A							,					eos_token_id=_A							,					sep_token_id=_A							,					**_A							,					)

				UpperCAmelCase   		=				vocab_size
				UpperCAmelCase   		=				max_position_embeddings
				UpperCAmelCase   		=				hidden_size
				UpperCAmelCase   		=				num_hidden_layers
				UpperCAmelCase   		=				num_attention_heads
				UpperCAmelCase   		=				intermediate_size
				UpperCAmelCase   		=				hidden_act
				UpperCAmelCase   		=				hidden_dropout_prob
				UpperCAmelCase   		=				attention_probs_dropout_prob
				UpperCAmelCase   		=				initializer_range
				UpperCAmelCase   		=				type_vocab_size
				UpperCAmelCase   		=				layer_norm_eps
				UpperCAmelCase   		=				use_cache

				UpperCAmelCase   		=				rescale_embeddings
				UpperCAmelCase   		=				attention_type
				UpperCAmelCase   		=				use_bias
				UpperCAmelCase   		=				block_size
				UpperCAmelCase   		=				num_random_blocks
				UpperCAmelCase   		=				classifier_dropout

class 	A_       (a_			):
	@property
	def      _lowercase				(    self		):






				'''simple docstring'''







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



