



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

def 			A_      (	__a     :				Any		):

				"""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 			A_      (	__a     :				Union[str, Any]		):

				"""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 			A_      (	__a     :				int ,		__a     :				str		):

				"""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 			A_      (	__a     :				Dict		):

				"""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 			A_      (	__a     :				int ,		__a     :				Union[str, Any] ,		__a     :				List[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_6_7_7_3_9_3 if """speech-commands""" not in model_name else -6.8_4_5_9_7_8
				a__				      =		4.5_6_8_9_9_7_4 if """speech-commands""" not in model_name else 5.5_6_5_4_5_2_6
				a__				      =		1_024 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=16_000 ,		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.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]		)
				elif model_name == "ast-finetuned-audioset-10-10-0.450":
								a__				      =		torch.tensor([-1.1_9_8_6, -7.0_9_0_3, -8.2_7_1_8]		)
				elif model_name == "ast-finetuned-audioset-10-10-0.448":
								a__				      =		torch.tensor([-2.6_1_2_8, -8.0_0_8_0, -9.4_3_4_4]		)
				elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
								a__				      =		torch.tensor([-1.5_0_8_0, -7.4_5_3_4, -8.8_9_1_7]		)
				elif model_name == "ast-finetuned-audioset-12-12-0.447":
								a__				      =		torch.tensor([-0.5_0_5_0, -6.5_8_3_3, -8.0_8_4_3]		)
				elif model_name == "ast-finetuned-audioset-14-14-0.443":
								a__				      =		torch.tensor([-0.3_8_2_6, -7.0_3_3_6, -8.2_4_1_3]		)
				elif model_name == "ast-finetuned-audioset-16-16-0.442":
								a__				      =		torch.tensor([-1.2_1_1_3, -6.9_1_0_1, -8.3_4_7_0]		)
				elif model_name == "ast-finetuned-speech-commands-v2":
								a__				      =		torch.tensor([6.1_5_8_9, -8.0_5_6_6, -8.7_9_8_4]		)
				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__":
	UpperCAmelCase							 =					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."""
	)

	UpperCAmelCase							 =					parser.parse_args()
	convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin

class  __snake_case      (   SCREAMING_SNAKE_CASE			,SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							@register_to_config
							def __init__(							self  ,  a_ = 128  ,  a_ = 256  ,  a_ = 2_000.0  ,  a_ = 768  ,  a_ = 12  ,  a_ = 12  ,  a_ = 64  ,  a_ = 2_048  ,  a_ = 0.1  ,  ):
											super().__init__()

											a__				      =		nn.Sequential(
											    nn.Linear(a_  ,  d_model * 4  ,  bias=a_		)  ,  nn.SiLU()  ,  nn.Linear(d_model * 4  ,  d_model * 4  ,  bias=a_		)  ,  nn.SiLU()  ,  )

											a__				      =		nn.Embedding(a_  ,  a_		)
											a__				      =		False

											a__				      =		nn.Linear(a_  ,  a_  ,  bias=a_		)

											a__				      =		nn.Dropout(p=a_		)

											a__				      =		nn.ModuleList()
											for lyr_num in range(a_		):
															# FiLM conditional T5 decoder
															a__				      =		DecoderLayer(d_model=a_  ,  d_kv=a_  ,  num_heads=a_  ,  d_ff=a_  ,  dropout_rate=a_		)
															self.decoders.append(a_		)

											a__				      =		TaLayerNorm(a_		)

											a__				      =		nn.Dropout(p=a_		)
											a__				      =		nn.Linear(a_  ,  a_  ,  bias=a_		)


							def _a      (							self  ,  a_  ,  a_		):
											a__				      =		torch.mul(query_input.unsqueeze(-1		)  ,  key_input.unsqueeze(-2		)		)
											return mask.unsqueeze(-3		)







							def _a      (							self  ,  a_  ,  a_  ,  a_		):
											a__      ,			a__      ,			a__				      =		decoder_input_tokens.shape
											assert decoder_noise_time.shape == (batch,)

											# decoder_noise_time is in [0, 1), so rescale to expected timing range.
											a__				      =		get_timestep_embedding(
											    decoder_noise_time * self.config.max_decoder_noise_time  ,  embedding_dim=self.config.d_model  ,  max_period=self.config.max_decoder_noise_time  ,  ).to(dtype=self.dtype		)

											a__				      =		self.conditioning_emb(a_		).unsqueeze(1		)

											assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

											a__				      =		decoder_input_tokens.shape[1]

											# If we want to use relative positions for audio context, we can just offset
											# this sequence by the length of encodings_and_masks.
											a__				      =		torch.broadcast_to(
											    torch.arange(a_  ,  device=decoder_input_tokens.device		)  ,  (batch, seq_length)  ,  )

											a__				      =		self.position_encoding(a_		)

											a__				      =		self.continuous_inputs_projection(a_		)
											inputs += position_encodings
											a__				      =		self.dropout(a_		)

											# decoder: No padding present.
											a__				      =		torch.ones(
											    decoder_input_tokens.shape[:2]  ,  device=decoder_input_tokens.device  ,  dtype=inputs.dtype		)

											# Translate encoding masks to encoder-decoder masks.
											a__				      =		[(x, self.encoder_decoder_mask(a_  ,  a_		)) for x, y in encodings_and_masks]

											# cross attend style: concat encodings
											a__				      =		torch.cat([x[0] for x in encodings_and_encdec_masks]  ,  dim=1		)
											a__				      =		torch.cat([x[1] for x in encodings_and_encdec_masks]  ,  dim=-1		)

											for lyr in self.decoders:
															a__				      =		lyr(
															    a_  ,  conditioning_emb=a_  ,  encoder_hidden_states=a_  ,  encoder_attention_mask=a_  ,  )[0]

											a__				      =		self.decoder_norm(a_		)
											a__				      =		self.post_dropout(a_		)

											a__				      =		self.spec_out(a_		)
											return spec_out




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_=1E-6		):
											super().__init__()
											a__				      =		nn.ModuleList()

											# cond self attention: layer 0
											self.layer.append(
											    TaLayerSelfAttentionCond(d_model=a_  ,  d_kv=a_  ,  num_heads=a_  ,  dropout_rate=a_		)		)

											# cross attention: layer 1
											self.layer.append(
											    TaLayerCrossAttention(
											        d_model=a_  ,  d_kv=a_  ,  num_heads=a_  ,  dropout_rate=a_  ,  layer_norm_epsilon=a_  ,  )		)

											# Film Cond MLP + dropout: last layer
											self.layer.append(
											    TaLayerFFCond(d_model=a_  ,  d_ff=a_  ,  dropout_rate=a_  ,  layer_norm_epsilon=a_		)		)







							def _a      (							self  ,  a_  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=None  ,  ):
											a__				      =		self.layer[0](
											    a_  ,  conditioning_emb=a_  ,  attention_mask=a_  ,  )

											if encoder_hidden_states is not None:
															a__				      =		torch.where(encoder_attention_mask > 0  ,  0  ,  -1E1_0		).to(
															    encoder_hidden_states.dtype		)

															a__				      =		self.layer[1](
															    a_  ,  key_value_states=a_  ,  attention_mask=a_  ,  )

											# Apply Film Conditional Feed Forward layer
											a__				      =		self.layer[-1](a_  ,  a_		)

											return (hidden_states,)




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_  ,  a_  ,  a_		):
											super().__init__()
											a__				      =		TaLayerNorm(a_		)
											a__				      =		TaFiLMLayer(in_features=d_model * 4  ,  out_features=a_		)
											a__				      =		Attention(query_dim=a_  ,  heads=a_  ,  dim_head=a_  ,  out_bias=a_  ,  scale_qk=a_		)
											a__				      =		nn.Dropout(a_		)







							def _a      (							self  ,  a_  ,  a_=None  ,  a_=None  ,  ):
											# pre_self_attention_layer_norm
											a__				      =		self.layer_norm(a_		)

											if conditioning_emb is not None:
															a__				      =		self.FiLMLayer(a_  ,  a_		)

											# Self-attention block
											a__				      =		self.attention(a_		)

											a__				      =		hidden_states + self.dropout(a_		)

											return hidden_states




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_  ,  a_  ,  a_  ,  a_		):
											super().__init__()
											a__				      =		Attention(query_dim=a_  ,  heads=a_  ,  dim_head=a_  ,  out_bias=a_  ,  scale_qk=a_		)
											a__				      =		TaLayerNorm(a_  ,  eps=a_		)
											a__				      =		nn.Dropout(a_		)







							def _a      (							self  ,  a_  ,  a_=None  ,  a_=None  ,  ):
											a__				      =		self.layer_norm(a_		)
											a__				      =		self.attention(
											    a_  ,  encoder_hidden_states=a_  ,  attention_mask=attention_mask.squeeze(1		)  ,  )
											a__				      =		hidden_states + self.dropout(a_		)
											return layer_output




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_  ,  a_  ,  a_		):
											super().__init__()
											a__				      =		TaDenseGatedActDense(d_model=a_  ,  d_ff=a_  ,  dropout_rate=a_		)
											a__				      =		TaFiLMLayer(in_features=d_model * 4  ,  out_features=a_		)
											a__				      =		TaLayerNorm(a_  ,  eps=a_		)
											a__				      =		nn.Dropout(a_		)







							def _a      (							self  ,  a_  ,  a_=None		):
											a__				      =		self.layer_norm(a_		)
											if conditioning_emb is not None:
															a__				      =		self.film(a_  ,  a_		)

											a__				      =		self.DenseReluDense(a_		)
											a__				      =		hidden_states + self.dropout(a_		)
											return hidden_states




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_  ,  a_		):
											super().__init__()
											a__				      =		nn.Linear(a_  ,  a_  ,  bias=a_		)
											a__				      =		nn.Linear(a_  ,  a_  ,  bias=a_		)
											a__				      =		nn.Linear(a_  ,  a_  ,  bias=a_		)
											a__				      =		nn.Dropout(a_		)
											a__				      =		NewGELUActivation()







							def _a      (							self  ,  a_		):
											a__				      =		self.act(self.wi_a(a_		)		)
											a__				      =		self.wi_a(a_		)
											a__				      =		hidden_gelu * hidden_linear
											a__				      =		self.dropout(a_		)

											a__				      =		self.wo(a_		)
											return hidden_states




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_=1E-6		):
											super().__init__()
											a__				      =		nn.Parameter(torch.ones(a_		)		)
											a__				      =		eps







							def _a      (							self  ,  a_		):
											# T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
											# Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated
											# w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
											# half-precision inputs is done in fp32

											a__				      =		hidden_states.to(torch.floataa		).pow(2		).mean(-1  ,  keepdim=a_		)
											a__				      =		hidden_states * torch.rsqrt(variance + self.variance_epsilon		)

											# convert into half-precision if necessary
											if self.weight.dtype in [torch.floataa, torch.bfloataa]:
															a__				      =		hidden_states.to(self.weight.dtype		)

											return self.weight * hidden_states




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def _a      (							self  ,  a_		):
											return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi		) * (input + 0.044_715 * torch.pow(a_  ,  3.0		))		))




class  __snake_case      (   nn.Module):







							'''simple docstring'''






							def __init__(							self  ,  a_  ,  a_		):
											super().__init__()
											a__				      =		nn.Linear(a_  ,  out_features * 2  ,  bias=a_		)







							def _a      (							self  ,  a_  ,  a_		):
											a__				      =		self.scale_bias(a_		)
											a__      ,			a__				      =		torch.chunk(a_  ,  2  ,  -1		)
											a__				      =		x * (1 + scale) + shift
											return x

