


"""simple docstring"""







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




def      UpperCAmelCase__ (lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       ASTConfig()

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

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

					__SCREAMING_SNAKE_CASE		     =       json.load(open(hf_hub_download(lowerCAmelCase_						,		lowerCAmelCase_						,		repo_type="dataset"     )						,		"r"     )     )
					__SCREAMING_SNAKE_CASE		     =       {int(lowerCAmelCase_     ): v for k, v in idalabel.items()}
					__SCREAMING_SNAKE_CASE		     =       idalabel
					__SCREAMING_SNAKE_CASE		     =       {v: k for k, v in idalabel.items()}

					return config




def      UpperCAmelCase__ (lowerCAmelCase_     ):




					'''simple docstring'''


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

					return name




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					for key in orig_state_dict.copy().keys():
										__SCREAMING_SNAKE_CASE		     =       orig_state_dict.pop(lowerCAmelCase_     )

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

					return orig_state_dict




def      UpperCAmelCase__ (lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       [
					    "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(lowerCAmelCase_						,		lowerCAmelCase_     )




@torch.no_grad()
def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_=False     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       get_audio_spectrogram_transformer_config(lowerCAmelCase_     )

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

					# load 🤗 model
					__SCREAMING_SNAKE_CASE		     =       ASTForAudioClassification(lowerCAmelCase_     )
					model.eval()

					model.load_state_dict(lowerCAmelCase_     )

					# verify outputs on dummy input
					# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
					__SCREAMING_SNAKE_CASE		     =       -4.2677393 if "speech-commands" not in model_name else -6.845978
					__SCREAMING_SNAKE_CASE		     =       4.5689974 if "speech-commands" not in model_name else 5.5654526
					__SCREAMING_SNAKE_CASE		     =       1024 if "speech-commands" not in model_name else 128
					__SCREAMING_SNAKE_CASE		     =       ASTFeatureExtractor(mean=lowerCAmelCase_						,		std=lowerCAmelCase_						,		max_length=lowerCAmelCase_     )

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

										__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       torchaudio.load(lowerCAmelCase_     )
										__SCREAMING_SNAKE_CASE		     =       waveform.squeeze().numpy()

					__SCREAMING_SNAKE_CASE		     =       feature_extractor(lowerCAmelCase_						,		sampling_rate=1_6000						,		return_tensors="pt"     )

					# forward pass
					__SCREAMING_SNAKE_CASE		     =       model(**lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       outputs.logits

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

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

					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__       :							Union[str, Any]        =  parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




"""simple docstring"""







import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
		from sagemaker import Session, TrainingJobAnalytics
		from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"    ,						"False")) is not True    ,						reason="Skipping test because should only be run when releasing minor transformers version"    ,						)
@pytest.mark.usefixtures("sm_env")
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "pytorch",
            "script": "run_ddp.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.7, "eval_loss": 0.6},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf_dist.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.p3.16xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.6, "eval_loss": 0.7},
        },
    ])
class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			Any:
						if self.framework == "pytorch":
											subprocess.run(
											    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()    ,    encoding="utf-8"    ,    check=UpperCAmelCase__    ,    )
						assert hasattr(self    ,    "env"     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       int     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						__SCREAMING_SNAKE_CASE		     =       {"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

						# creates estimator
						return HuggingFace(
						    entry_point=self.script    ,    source_dir=self.env.test_path    ,    role=self.env.role    ,    image_uri=self.env.image_uri    ,    base_job_name=UpperCAmelCase__    ,    instance_count=UpperCAmelCase__    ,    instance_type=self.instance_type    ,    debugger_hook_config=UpperCAmelCase__    ,    hyperparameters={**self.env.distributed_hyperparameters, "model_name_or_path": self.model_name_or_path}    ,    metric_definitions=self.env.metric_definitions    ,    distribution=UpperCAmelCase__    ,    py_version="py36"    ,    )




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       Optional[int]     ) ->			Optional[Any]:
						TrainingJobAnalytics(UpperCAmelCase__     ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""     )




	@parameterized.expand([(2,)]     )
	def 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Optional[Any]     ) ->			List[str]:
						# create estimator
						__SCREAMING_SNAKE_CASE		     =       self.create_estimator(UpperCAmelCase__     )

						# run training
						estimator.fit()

						# result dataframe
						__SCREAMING_SNAKE_CASE		     =       TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

						# extract kpis
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]     )
						__SCREAMING_SNAKE_CASE		     =       list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]     )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						__SCREAMING_SNAKE_CASE		     =       (
						    Session().describe_training_job(estimator.latest_training_job.name     ).get("TrainingTimeInSeconds"    ,    9_9_9_9_9_9     )
						)

						# assert kpis
						assert train_runtime <= self.results["train_runtime"]
						assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy     )
						assert all(t <= self.results["eval_loss"] for t in eval_loss     )

						# dump tests result into json file to share in PR
						with open(F"""{estimator.latest_training_job.name}.json"""    ,    "w"     ) as outfile:
											json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}    ,    UpperCAmelCase__     )
