

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

def        lowerCamelCase__       (    _a):
	SCREAMING_SNAKE_CASE      :					List[str]    =				ASTConfig()

	if "10-10" in model_name:
		pass
	elif "speech-commands" in model_name:
		SCREAMING_SNAKE_CASE      :					Dict    =				128
	elif "12-12" in model_name:
		SCREAMING_SNAKE_CASE      :					List[str]    =				12
		SCREAMING_SNAKE_CASE      :					Optional[int]    =				12
	elif "14-14" in model_name:
		SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				14
		SCREAMING_SNAKE_CASE      :					Tuple    =				14
	elif "16-16" in model_name:
		SCREAMING_SNAKE_CASE      :					Optional[Any]    =				16
		SCREAMING_SNAKE_CASE      :					str    =				16
	else:
		raise ValueError("Model not supported")

	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				"huggingface/label-files"
	if "speech-commands" in model_name:
		SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				35
		SCREAMING_SNAKE_CASE      :					Tuple    =				"speech-commands-v2-id2label.json"
	else:
		SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				527
		SCREAMING_SNAKE_CASE      :					str    =				"audioset-id2label.json"

	SCREAMING_SNAKE_CASE      :					List[str]    =				json.load(open(hf_hub_download(_a							,	_a							,	repo_type="dataset")							,	"r"))
	SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				{int(_a): v for k, v in idalabel.items()}
	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				idalabel
	SCREAMING_SNAKE_CASE      :					Tuple    =				{v: k for k, v in idalabel.items()}

	return config

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

	return name

def        lowerCamelCase__       (    _a							,	_a):
	for key in orig_state_dict.copy().keys():
		SCREAMING_SNAKE_CASE      :					List[Any]    =				orig_state_dict.pop(_a)

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

	return orig_state_dict

def        lowerCamelCase__       (    _a):
	SCREAMING_SNAKE_CASE      :					List[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        lowerCamelCase__       (    _a							,	_a							,	_a=False):
	SCREAMING_SNAKE_CASE      :					int    =				get_audio_spectrogram_transformer_config(_a)

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

	# load 🤗 model
	SCREAMING_SNAKE_CASE      :					Dict    =				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
	SCREAMING_SNAKE_CASE      :					List[str]    =				-4.267_7393 if "speech-commands" not in model_name else -6.84_5978
	SCREAMING_SNAKE_CASE      :					Tuple    =				4.568_9974 if "speech-commands" not in model_name else 5.565_4526
	SCREAMING_SNAKE_CASE      :					List[Any]    =				1024 if "speech-commands" not in model_name else 128
	SCREAMING_SNAKE_CASE      :					Optional[int]    =				ASTFeatureExtractor(mean=_a							,	std=_a							,	max_length=_a)

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

		SCREAMING_SNAKE_CASE   ,SCREAMING_SNAKE_CASE      :					List[str]    =				torchaudio.load(_a)
		SCREAMING_SNAKE_CASE      :					Any    =				waveform.squeeze().numpy()

	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				feature_extractor(_a							,	sampling_rate=16000							,	return_tensors="pt")

	# forward pass
	SCREAMING_SNAKE_CASE      :					Optional[Any]    =				model(**_a)
	SCREAMING_SNAKE_CASE      :					int    =				outputs.logits

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





import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
					import torch_xla.core.xla_model as xm
					import torch_xla.debug.metrics as met
class 			a__    (					snake_case__      ):






						def __init__(   self		,			*_A		,			_A=None		,			_A=None		,			**_A			):




											"""simple docstring"""



											super().__init__(*_A		,			**_A			)
											__lowerCAmelCase				   =       eval_examples
											__lowerCAmelCase				   =       post_process_function






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A = None		,			_A=None		,			_A = None		,			_A = "eval"		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       gen_kwargs.copy()
											__lowerCAmelCase				   =       (
											    gen_kwargs["max_length"] if gen_kwargs.get("max_length"			) is not None else self.args.generation_max_length
											)
											__lowerCAmelCase				   =       (
											    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"			) is not None else self.args.generation_num_beams
											)
											__lowerCAmelCase				   =       gen_kwargs

											__lowerCAmelCase				   =       self.eval_dataset if eval_dataset is None else eval_dataset
											__lowerCAmelCase				   =       self.get_eval_dataloader(_A			)
											__lowerCAmelCase				   =       self.eval_examples if eval_examples is None else eval_examples

											# Temporarily disable metric computation, we will do it in the loop here.
											__lowerCAmelCase				   =       self.compute_metrics
											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       time.time()
											__lowerCAmelCase				   =       self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__lowerCAmelCase				   =       eval_loop(
																    _A		,			description="Evaluation"		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=_A		,			metric_key_prefix=_A		,			)
											finally:
																__lowerCAmelCase				   =       compute_metrics
											__lowerCAmelCase				   =       self.args.eval_batch_size * self.args.world_size
											if f"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
																start_time += output.metrics[f"""{metric_key_prefix}_jit_compilation_time"""]
											output.metrics.update(
											    speed_metrics(
											        _A		,			_A		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size			)		,			)			)

											if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
																# Only the main node write the results by default
																__lowerCAmelCase				   =       self.post_process_function(_A		,			_A		,			_A			)
																__lowerCAmelCase				   =       self.compute_metrics(_A			)

																# Prefix all keys with metric_key_prefix + '_'
																for key in list(metrics.keys()			):
																					if not key.startswith(f"""{metric_key_prefix}_"""			):
																										__lowerCAmelCase				   =       metrics.pop(_A			)

																metrics.update(output.metrics			)
											else:
																__lowerCAmelCase				   =       output.metrics

											if self.args.should_log:
																# Only the main node log the results by default
																self.log(_A			)

											if self.args.tpu_metrics_debug or self.args.debug:
																# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
																xm.master_print(met.metrics_report()			)

											__lowerCAmelCase				   =       self.callback_handler.on_evaluate(self.args		,			self.state		,			self.control		,			_A			)
											return metrics






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A=None		,			_A = "test"		,			**_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       gen_kwargs.copy()

											__lowerCAmelCase				   =       self.get_test_dataloader(_A			)

											# Temporarily disable metric computation, we will do it in the loop here.
											__lowerCAmelCase				   =       self.compute_metrics
											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       time.time()
											__lowerCAmelCase				   =       self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__lowerCAmelCase				   =       eval_loop(
																    _A		,			description="Prediction"		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=_A		,			metric_key_prefix=_A		,			)
											finally:
																__lowerCAmelCase				   =       compute_metrics

											__lowerCAmelCase				   =       self.args.eval_batch_size * self.args.world_size
											if f"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
																start_time += output.metrics[f"""{metric_key_prefix}_jit_compilation_time"""]
											output.metrics.update(
											    speed_metrics(
											        _A		,			_A		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size			)		,			)			)
											if self.post_process_function is None or self.compute_metrics is None:
																return output

											__lowerCAmelCase				   =       self.post_process_function(_A		,			_A		,			_A		,			"predict"			)
											__lowerCAmelCase				   =       self.compute_metrics(_A			)

											# Prefix all keys with metric_key_prefix + '_'
											for key in list(metrics.keys()			):
																if not key.startswith(f"""{metric_key_prefix}_"""			):
																					__lowerCAmelCase				   =       metrics.pop(_A			)
											metrics.update(output.metrics			)
											return PredictionOutput(predictions=predictions.predictions		,			label_ids=predictions.label_ids		,			metrics=_A			)




