




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




def 			__magic_name__				( A		:					int   ):






			'''simple docstring'''







			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 			__magic_name__				( A		:					Dict   ):






			'''simple docstring'''







			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 			__magic_name__				( A		:					Tuple,  A		:					Optional[Any]   ):






			'''simple docstring'''







			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 			__magic_name__				( A		:					Optional[int]   ):






			'''simple docstring'''







			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 			__magic_name__				( A		:					List[str],  A		:					Optional[int],  A		:					Any=False   ):






			'''simple docstring'''







			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=16000,  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		:				Tuple       			=		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		:				Dict       			=		parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__lowerCAmelCase		:				str       			=		logging.get_logger(__name__)

__lowerCAmelCase		:				Dict       			=		{'vocab_file': 'spiece.model'}

__lowerCAmelCase		:				Optional[int]       			=		{
    'vocab_file': {
        'bert_for_seq_generation': (
            'https://huggingface.co/google/bert_for_seq_generation_L-24_bbc_encoder/resolve/main/spiece.model'
        ),
    }
}

__lowerCAmelCase		:				Dict       			=		{'bert_for_seq_generation': 512}




class     snake_case__							(_UpperCamelCase   ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							Any							 =					VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE_       :							Union[str, Any]							 =					PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE_       :							int							 =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE_       :							List[int]							 =					[]
				SCREAMING_SNAKE_CASE_       :							Optional[Any]							 =					["""input_ids""", """attention_mask"""]





				def __init__(      self				:     List[Any]					,      __lowerCamelCase				:     Tuple					,      __lowerCamelCase				:     str="<s>"					,      __lowerCamelCase				:     Optional[int]="</s>"					,      __lowerCamelCase				:     int="<unk>"					,      __lowerCamelCase				:     List[str]="<pad>"					,      __lowerCamelCase				:     Union[str, Any]="<::::>"					,      __lowerCamelCase				:     Optional[Dict[str, Any]] = None					,      **__lowerCamelCase				:     List[str]					,      )					->     None:
							a    =			{} if sp_model_kwargs is None else sp_model_kwargs

							# Add extra_ids to the special token list
							super().__init__(
							    bos_token=__lowerCamelCase					,      eos_token=__lowerCamelCase					,      unk_token=__lowerCamelCase					,      pad_token=__lowerCamelCase					,      sep_token=__lowerCamelCase					,      sp_model_kwargs=self.sp_model_kwargs					,      **__lowerCamelCase					,      )

							a    =			vocab_file

							a    =			spm.SentencePieceProcessor(**self.sp_model_kwargs							)
							self.sp_model.Load(__lowerCamelCase							)





				@property
				def 							__UpperCAmelCase							(      self				:     Dict							)					->     Dict:
							return self.sp_model.get_piece_size()





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     Optional[int]:
							a    =			{self.convert_ids_to_tokens(__lowerCamelCase							): i for i in range(self.vocab_size							)}
							vocab.update(self.added_tokens_encoder							)
							return vocab





				def __getstate__(      self				:     Dict							)					->     Optional[Any]:
							a    =			self.__dict__.copy()
							a    =			None
							return state





				def __setstate__(      self				:     Optional[Any]					,      __lowerCamelCase				:     Dict							)					->     Optional[Any]:
							a    =			d

							# for backward compatibility
							if not hasattr(self					,      "sp_model_kwargs"							):
										a    =			{}

							a    =			spm.SentencePieceProcessor(**self.sp_model_kwargs							)
							self.sp_model.Load(self.vocab_file							)





				def 							__UpperCAmelCase							(      self				:     Union[str, Any]					,      __lowerCamelCase				:     str							)					->     List[str]:
							return self.sp_model.encode(__lowerCamelCase					,      out_type=__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     str					,      __lowerCamelCase				:     Union[str, Any]							)					->     int:
							return self.sp_model.piece_to_id(__lowerCamelCase							)





				def 							__UpperCAmelCase							(      self				:     Tuple					,      __lowerCamelCase				:     List[str]							)					->     Any:
							a    =			self.sp_model.IdToPiece(__lowerCamelCase							)
							return token





				def 							__UpperCAmelCase							(      self				:     Any					,      __lowerCamelCase				:     Dict							)					->     Any:
							a    =			[]
							a    =			""
							for token in tokens:
										# make sure that special tokens are not decoded using sentencepiece model
										if token in self.all_special_tokens:
													out_string += self.sp_model.decode(__lowerCamelCase							) + token
													a    =			[]
										else:
													current_sub_tokens.append(__lowerCamelCase							)
							out_string += self.sp_model.decode(__lowerCamelCase							)
							return out_string.strip()





				def 							__UpperCAmelCase							(      self				:     Dict					,      __lowerCamelCase				:     str					,      __lowerCamelCase				:     Optional[str] = None							)					->     Tuple[str]:
							if not os.path.isdir(__lowerCamelCase							):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""							)
										return
							a    =			os.path.join(
							    __lowerCamelCase					,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]							)

							if os.path.abspath(self.vocab_file							) != os.path.abspath(__lowerCamelCase							) and os.path.isfile(self.vocab_file							):
										copyfile(self.vocab_file					,      __lowerCamelCase							)
							elif not os.path.isfile(self.vocab_file							):
										with open(__lowerCamelCase					,      "wb"							) as fi:
													a    =			self.sp_model.serialized_model_proto()
													fi.write(__lowerCamelCase							)

							return (out_vocab_file,)






