

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


def 					_UpperCAmelCase				(						a__):

			'''simple docstring'''



			a_						:						Optional[int] 							= ASTConfig()

			if "10-10" in model_name:
						pass
			elif "speech-commands" in model_name:
						a_						:						Tuple 							= 1_2_8
			elif "12-12" in model_name:
						a_						:						Optional[Any] 							= 1_2
						a_						:						Any 							= 1_2
			elif "14-14" in model_name:
						a_						:						str 							= 1_4
						a_						:						Optional[Any] 							= 1_4
			elif "16-16" in model_name:
						a_						:						Union[str, Any] 							= 1_6
						a_						:						List[str] 							= 1_6
			else:
						raise ValueError("""Model not supported""")

			a_						:						Optional[int] 							= """huggingface/label-files"""
			if "speech-commands" in model_name:
						a_						:						Optional[Any] 							= 3_5
						a_						:						List[Any] 							= """speech-commands-v2-id2label.json"""
			else:
						a_						:						str 							= 5_2_7
						a_						:						str 							= """audioset-id2label.json"""

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

			return config


def 					_UpperCAmelCase				(						a__):

			'''simple docstring'''



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

			return name


def 					_UpperCAmelCase				(						a__  ,	a__):

			'''simple docstring'''



			for key in orig_state_dict.copy().keys():
						a_						:						Tuple 							= orig_state_dict.pop(a__)

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

			return orig_state_dict


def 					_UpperCAmelCase				(						a__):

			'''simple docstring'''



			a_						:						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 					_UpperCAmelCase				(						a__  ,	a__  ,	a__=False):

			'''simple docstring'''



			a_						:						List[str] 							= get_audio_spectrogram_transformer_config(a__)

			a_						:						str 							= {
			    """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_						:						Union[str, Any] 							= model_name_to_url[model_name]
			a_						:						str 							= torch.hub.load_state_dict_from_url(a__  ,	map_location="""cpu""")
			# remove some keys
			remove_keys(a__)
			# rename some keys
			a_						:						List[Any] 							= convert_state_dict(a__  ,	a__)

			# load 🤗 model
			a_						:						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
			a_						:						Dict 							= -4.267_7393 if """speech-commands""" not in model_name else -6.84_5978
			a_						:						int 							= 4.568_9974 if """speech-commands""" not in model_name else 5.565_4526
			a_						:						Tuple 							= 1_0_2_4 if """speech-commands""" not in model_name else 1_2_8
			a_						:						Any 							= ASTFeatureExtractor(mean=a__  ,	std=a__  ,	max_length=a__)

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

						a_							,					a_						:						Dict 							= torchaudio.load(a__)
						a_						:						Tuple 							= waveform.squeeze().numpy()

			a_						:						int 							= feature_extractor(a__  ,	sampling_rate=1_6_0_0_0  ,	return_tensors="""pt""")

			# forward pass
			a_						:						Dict 							= model(**a__)
			a_						:						Any 							= outputs.logits

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

	__snake_case				:      str   		=			parser.parse_args()
	convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__snake_case				:      int   		=			{
    """configuration_llama""": ["""LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LlamaConfig"""],
}

try:
	if not is_sentencepiece_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      int   		=			["""LlamaTokenizer"""]

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      int   		=			["""LlamaTokenizerFast"""]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      Union[str, Any]   		=			[
	    """LlamaForCausalLM""",
	    """LlamaModel""",
	    """LlamaPreTrainedModel""",
	    """LlamaForSequenceClassification""",
	]


if TYPE_CHECKING:
	from .configuration_llama import LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, LlamaConfig

	try:
		if not is_sentencepiece_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_llama import LlamaTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_llama_fast import LlamaTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_llama import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaPreTrainedModel


else:
	import sys

	__snake_case				:      Dict   		=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


