




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





def 							lowercase					(snake_case__   : int )     ->  Tuple:


				'''simple docstring'''






				lowerCAmelCase  =				ASTConfig()

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

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

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

				return config





def 							lowercase					(snake_case__   : List[str] )     ->  List[Any]:


				'''simple docstring'''






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

				return name





def 							lowercase					(snake_case__   : Dict      ,       snake_case__   : Any )     ->  str:


				'''simple docstring'''






				for key in orig_state_dict.copy().keys():
								lowerCAmelCase  =				orig_state_dict.pop(snake_case__ )

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

				return orig_state_dict





def 							lowercase					(snake_case__   : Optional[int] )     ->  Optional[Any]:


				'''simple docstring'''






				lowerCAmelCase  =				[
				    """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(snake_case__      ,       snake_case__ )





@torch.no_grad()
def 							lowercase					(snake_case__   : Optional[int]      ,       snake_case__   : List[str]      ,       snake_case__   : Optional[Any]=False )     ->  Tuple:


				'''simple docstring'''






				lowerCAmelCase  =				get_audio_spectrogram_transformer_config(snake_case__ )

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

				# load 🤗 model
				lowerCAmelCase  =				ASTForAudioClassification(snake_case__ )
				model.eval()

				model.load_state_dict(snake_case__ )

				# verify outputs on dummy input
				# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
				lowerCAmelCase  =				-4.2_677_393 if """speech-commands""" not in model_name else -6.845_978
				lowerCAmelCase  =				4.5_689_974 if """speech-commands""" not in model_name else 5.5_654_526
				lowerCAmelCase  =				1_024 if """speech-commands""" not in model_name else 128
				lowerCAmelCase  =				ASTFeatureExtractor(mean=snake_case__      ,       std=snake_case__      ,       max_length=snake_case__ )

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

								lowerCAmelCase ,						lowerCAmelCase  =				torchaudio.load(snake_case__ )
								lowerCAmelCase  =				waveform.squeeze().numpy()

				lowerCAmelCase  =				feature_extractor(snake_case__      ,       sampling_rate=16_000      ,       return_tensors="""pt""" )

				# forward pass
				lowerCAmelCase  =				model(**snake_case__ )
				lowerCAmelCase  =				outputs.logits

				if model_name == "ast-finetuned-audioset-10-10-0.4593":
								lowerCAmelCase  =				torch.tensor([-0.8_760, -7.0_042, -8.6_602] )
				elif model_name == "ast-finetuned-audioset-10-10-0.450":
								lowerCAmelCase  =				torch.tensor([-1.1_986, -7.0_903, -8.2_718] )
				elif model_name == "ast-finetuned-audioset-10-10-0.448":
								lowerCAmelCase  =				torch.tensor([-2.6_128, -8.0_080, -9.4_344] )
				elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
								lowerCAmelCase  =				torch.tensor([-1.5_080, -7.4_534, -8.8_917] )
				elif model_name == "ast-finetuned-audioset-12-12-0.447":
								lowerCAmelCase  =				torch.tensor([-0.5_050, -6.5_833, -8.0_843] )
				elif model_name == "ast-finetuned-audioset-14-14-0.443":
								lowerCAmelCase  =				torch.tensor([-0.3_826, -7.0_336, -8.2_413] )
				elif model_name == "ast-finetuned-audioset-16-16-0.442":
								lowerCAmelCase  =				torch.tensor([-1.2_113, -6.9_101, -8.3_470] )
				elif model_name == "ast-finetuned-speech-commands-v2":
								lowerCAmelCase  =				torch.tensor([6.1_589, -8.0_566, -8.7_984] )
				else:
								raise ValueError("""Unknown model name""" )
				if not torch.allclose(logits[0, :3]      ,       snake_case__      ,       atol=1e-4 ):
								raise ValueError("""Logits don't match""" )
				print("""Looks ok!""" )

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

				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)











"""simple docstring"""







import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class       SCREAMING_SNAKE_CASE__					(					nn.Module					):
						_a    						=						42
						_a    						=						42
						_a    						=						0.0
						_a    						=						1
						_a    						=						1
						_a    						=						True
						_a    						=						False
						_a    						=						False
						_a    						=						False
						_a    						=						jnp.floataa


						def   __lowercase (							self				:			List[str]					):
										lowerCAmelCase  =				[]
										lowerCAmelCase  =				[]

										for i in range(self.num_layers					):
														lowerCAmelCase  =				self.in_channels if i == 0 else self.out_channels

														lowerCAmelCase  =				FlaxResnetBlockaD(
														    in_channels=lowerCAmelCase						,						out_channels=self.out_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
														resnets.append(lowerCAmelCase					)

														lowerCAmelCase  =				FlaxTransformeraDModel(
														    in_channels=self.out_channels						,						n_heads=self.num_attention_heads						,						d_head=self.out_channels // self.num_attention_heads						,						depth=1						,						use_linear_projection=self.use_linear_projection						,						only_cross_attention=self.only_cross_attention						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)
														attentions.append(lowerCAmelCase					)

										lowerCAmelCase  =				resnets
										lowerCAmelCase  =				attentions

										if self.add_downsample:
														lowerCAmelCase  =				FlaxDownsampleaD(self.out_channels						,						dtype=self.dtype					)




						def __call__(							self				:			Tuple						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			Union[str, Any]						,						lowerCAmelCase				:			Tuple						,						lowerCAmelCase				:			int=True					):
										lowerCAmelCase  =				()

										for resnet, attn in zip(self.resnets						,						self.attentions					):
														lowerCAmelCase  =				resnet(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)
														lowerCAmelCase  =				attn(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)
														output_states += (hidden_states,)

										if self.add_downsample:
														lowerCAmelCase  =				self.downsamplers_a(lowerCAmelCase					)
														output_states += (hidden_states,)

										return hidden_states, output_states






class       SCREAMING_SNAKE_CASE__					(					nn.Module					):
						_a    						=						42
						_a    						=						42
						_a    						=						0.0
						_a    						=						1
						_a    						=						True
						_a    						=						jnp.floataa


						def   __lowercase (							self				:			Union[str, Any]					):
										lowerCAmelCase  =				[]

										for i in range(self.num_layers					):
														lowerCAmelCase  =				self.in_channels if i == 0 else self.out_channels

														lowerCAmelCase  =				FlaxResnetBlockaD(
														    in_channels=lowerCAmelCase						,						out_channels=self.out_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
														resnets.append(lowerCAmelCase					)
										lowerCAmelCase  =				resnets

										if self.add_downsample:
														lowerCAmelCase  =				FlaxDownsampleaD(self.out_channels						,						dtype=self.dtype					)




						def __call__(							self				:			int						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			List[str]=True					):
										lowerCAmelCase  =				()

										for resnet in self.resnets:
														lowerCAmelCase  =				resnet(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)
														output_states += (hidden_states,)

										if self.add_downsample:
														lowerCAmelCase  =				self.downsamplers_a(lowerCAmelCase					)
														output_states += (hidden_states,)

										return hidden_states, output_states






class       SCREAMING_SNAKE_CASE__					(					nn.Module					):
						_a    						=						42
						_a    						=						42
						_a    						=						42
						_a    						=						0.0
						_a    						=						1
						_a    						=						1
						_a    						=						True
						_a    						=						False
						_a    						=						False
						_a    						=						False
						_a    						=						jnp.floataa


						def   __lowercase (							self				:			List[str]					):
										lowerCAmelCase  =				[]
										lowerCAmelCase  =				[]

										for i in range(self.num_layers					):
														lowerCAmelCase  =				self.in_channels if (i == self.num_layers - 1) else self.out_channels
														lowerCAmelCase  =				self.prev_output_channel if i == 0 else self.out_channels

														lowerCAmelCase  =				FlaxResnetBlockaD(
														    in_channels=resnet_in_channels + res_skip_channels						,						out_channels=self.out_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
														resnets.append(lowerCAmelCase					)

														lowerCAmelCase  =				FlaxTransformeraDModel(
														    in_channels=self.out_channels						,						n_heads=self.num_attention_heads						,						d_head=self.out_channels // self.num_attention_heads						,						depth=1						,						use_linear_projection=self.use_linear_projection						,						only_cross_attention=self.only_cross_attention						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)
														attentions.append(lowerCAmelCase					)

										lowerCAmelCase  =				resnets
										lowerCAmelCase  =				attentions

										if self.add_upsample:
														lowerCAmelCase  =				FlaxUpsampleaD(self.out_channels						,						dtype=self.dtype					)




						def __call__(							self				:			Optional[Any]						,						lowerCAmelCase				:			Tuple						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			Dict=True					):
										for resnet, attn in zip(self.resnets						,						self.attentions					):
														# pop res hidden states
														lowerCAmelCase  =				res_hidden_states_tuple[-1]
														lowerCAmelCase  =				res_hidden_states_tuple[:-1]
														lowerCAmelCase  =				jnp.concatenate((hidden_states, res_hidden_states)						,						axis=-1					)

														lowerCAmelCase  =				resnet(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)
														lowerCAmelCase  =				attn(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)

										if self.add_upsample:
														lowerCAmelCase  =				self.upsamplers_a(lowerCAmelCase					)

										return hidden_states






class       SCREAMING_SNAKE_CASE__					(					nn.Module					):
						_a    						=						42
						_a    						=						42
						_a    						=						42
						_a    						=						0.0
						_a    						=						1
						_a    						=						True
						_a    						=						jnp.floataa


						def   __lowercase (							self				:			Tuple					):
										lowerCAmelCase  =				[]

										for i in range(self.num_layers					):
														lowerCAmelCase  =				self.in_channels if (i == self.num_layers - 1) else self.out_channels
														lowerCAmelCase  =				self.prev_output_channel if i == 0 else self.out_channels

														lowerCAmelCase  =				FlaxResnetBlockaD(
														    in_channels=resnet_in_channels + res_skip_channels						,						out_channels=self.out_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
														resnets.append(lowerCAmelCase					)

										lowerCAmelCase  =				resnets

										if self.add_upsample:
														lowerCAmelCase  =				FlaxUpsampleaD(self.out_channels						,						dtype=self.dtype					)




						def __call__(							self				:			Optional[Any]						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			List[str]						,						lowerCAmelCase				:			List[str]						,						lowerCAmelCase				:			Dict=True					):
										for resnet in self.resnets:
														# pop res hidden states
														lowerCAmelCase  =				res_hidden_states_tuple[-1]
														lowerCAmelCase  =				res_hidden_states_tuple[:-1]
														lowerCAmelCase  =				jnp.concatenate((hidden_states, res_hidden_states)						,						axis=-1					)

														lowerCAmelCase  =				resnet(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)

										if self.add_upsample:
														lowerCAmelCase  =				self.upsamplers_a(lowerCAmelCase					)

										return hidden_states






class       SCREAMING_SNAKE_CASE__					(					nn.Module					):
						_a    						=						42
						_a    						=						0.0
						_a    						=						1
						_a    						=						1
						_a    						=						False
						_a    						=						False
						_a    						=						jnp.floataa


						def   __lowercase (							self				:			List[Any]					):
										# there is always at least one resnet
										lowerCAmelCase  =				[
										    FlaxResnetBlockaD(
										        in_channels=self.in_channels						,						out_channels=self.in_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
										]

										lowerCAmelCase  =				[]

										for _ in range(self.num_layers					):
														lowerCAmelCase  =				FlaxTransformeraDModel(
														    in_channels=self.in_channels						,						n_heads=self.num_attention_heads						,						d_head=self.in_channels // self.num_attention_heads						,						depth=1						,						use_linear_projection=self.use_linear_projection						,						use_memory_efficient_attention=self.use_memory_efficient_attention						,						dtype=self.dtype						,						)
														attentions.append(lowerCAmelCase					)

														lowerCAmelCase  =				FlaxResnetBlockaD(
														    in_channels=self.in_channels						,						out_channels=self.in_channels						,						dropout_prob=self.dropout						,						dtype=self.dtype						,						)
														resnets.append(lowerCAmelCase					)

										lowerCAmelCase  =				resnets
										lowerCAmelCase  =				attentions




						def __call__(							self				:			Optional[int]						,						lowerCAmelCase				:			List[str]						,						lowerCAmelCase				:			int						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			List[Any]=True					):
										lowerCAmelCase  =				self.resnets[0](lowerCAmelCase						,						lowerCAmelCase					)
										for attn, resnet in zip(self.attentions						,						self.resnets[1:]					):
														lowerCAmelCase  =				attn(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)
														lowerCAmelCase  =				resnet(lowerCAmelCase						,						lowerCAmelCase						,						deterministic=lowerCAmelCase					)

										return hidden_states





