


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





def    a__			(   lowerCAmelCase__		)       ->      str:
    UpperCAmelCase__     :						Union[str, Any]						=							ASTConfig()

    if "10-10" in model_name:
        pass
    elif "speech-commands" in model_name:
        UpperCAmelCase__     :						List[Any]						=							1_28
    elif "12-12" in model_name:
        UpperCAmelCase__     :						Dict						=							12
        UpperCAmelCase__     :						str						=							12
    elif "14-14" in model_name:
        UpperCAmelCase__     :						Dict						=							14
        UpperCAmelCase__     :						Tuple						=							14
    elif "16-16" in model_name:
        UpperCAmelCase__     :						List[str]						=							16
        UpperCAmelCase__     :						Optional[Any]						=							16
    else:
        raise ValueError('''Model not supported'''		)

    UpperCAmelCase__     :						Optional[int]						=							'''huggingface/label-files'''
    if "speech-commands" in model_name:
        UpperCAmelCase__     :						List[Any]						=							35
        UpperCAmelCase__     :						Optional[int]						=							'''speech-commands-v2-id2label.json'''
    else:
        UpperCAmelCase__     :						List[Any]						=							5_27
        UpperCAmelCase__     :						Union[str, Any]						=							'''audioset-id2label.json'''

    UpperCAmelCase__     :						Dict						=							json.load(open(hf_hub_download(lowerCAmelCase__    ,   lowerCAmelCase__    ,   repo_type='''dataset'''		)    ,   '''r'''		)		)
    UpperCAmelCase__     :						Optional[Any]						=							{int(lowerCAmelCase__		): v for k, v in idalabel.items()}
    UpperCAmelCase__     :						Optional[Any]						=							idalabel
    UpperCAmelCase__     :						str						=							{v: k for k, v in idalabel.items()}

    return config





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

    return name





def    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__		)       ->      List[Any]:
    for key in orig_state_dict.copy().keys():
        UpperCAmelCase__     :						Dict						=							orig_state_dict.pop(lowerCAmelCase__		)

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

    return orig_state_dict





def    a__			(   lowerCAmelCase__		)       ->      List[Any]:
    UpperCAmelCase__     :						int						=							[
        '''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    a__			(   lowerCAmelCase__    ,   lowerCAmelCase__    ,   lowerCAmelCase__=False		)       ->      List[Any]:
    UpperCAmelCase__     :						Tuple						=							get_audio_spectrogram_transformer_config(lowerCAmelCase__		)

    UpperCAmelCase__     :						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
    UpperCAmelCase__     :						Any						=							model_name_to_url[model_name]
    UpperCAmelCase__     :						str						=							torch.hub.load_state_dict_from_url(lowerCAmelCase__    ,   map_location='''cpu'''		)
    # remove some keys
    remove_keys(lowerCAmelCase__		)
    # rename some keys
    UpperCAmelCase__     :						int						=							convert_state_dict(lowerCAmelCase__    ,   lowerCAmelCase__		)

    # load 🤗 model
    UpperCAmelCase__     :						int						=							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
    UpperCAmelCase__     :						Any						=							-4.2_6_7_7_3_9_3 if '''speech-commands''' not in model_name else -6.8_4_5_9_7_8
    UpperCAmelCase__     :						Union[str, Any]						=							4.5_6_8_9_9_7_4 if '''speech-commands''' not in model_name else 5.5_6_5_4_5_2_6
    UpperCAmelCase__     :						int						=							10_24 if '''speech-commands''' not in model_name else 1_28
    UpperCAmelCase__     :						Optional[Any]						=							ASTFeatureExtractor(mean=lowerCAmelCase__    ,   std=lowerCAmelCase__    ,   max_length=lowerCAmelCase__		)

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

        UpperCAmelCase__		,		UpperCAmelCase__     :						Any						=							torchaudio.load(lowerCAmelCase__		)
        UpperCAmelCase__     :						Dict						=							waveform.squeeze().numpy()

    UpperCAmelCase__     :						Tuple						=							feature_extractor(lowerCAmelCase__    ,   sampling_rate=1_60_00    ,   return_tensors='''pt'''		)

    # forward pass
    UpperCAmelCase__     :						Tuple						=							model(**lowerCAmelCase__		)
    UpperCAmelCase__     :						Any						=							outputs.logits

    if model_name == "ast-finetuned-audioset-10-10-0.4593":
        UpperCAmelCase__     :						Optional[int]						=							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":
        UpperCAmelCase__     :						Union[str, Any]						=							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":
        UpperCAmelCase__     :						Tuple						=							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":
        UpperCAmelCase__     :						Tuple						=							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":
        UpperCAmelCase__     :						Union[str, Any]						=							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":
        UpperCAmelCase__     :						List[str]						=							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":
        UpperCAmelCase__     :						Optional[Any]						=							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":
        UpperCAmelCase__     :						str						=							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]    ,   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__":
  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)





'''simple docstring'''

from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCAmelCase		     =  logging.get_logger(__name__)

# TODO Update this
__lowerCAmelCase		     =  {
    '''facebook/esm-1b''': '''https://huggingface.co/facebook/esm-1b/resolve/main/config.json''',
    # See all ESM models at https://huggingface.co/models?filter=esm
}

class       __magic_name__		(       _UpperCamelCase  ):
							lowerCAmelCase			:       Any  							=   'esm'







							def __init__(				self   :			Optional[int]      ,_UpperCAmelCase   :			Optional[Any]=None      ,_UpperCAmelCase   :			Dict=None      ,_UpperCAmelCase   :			Union[str, Any]=None      ,_UpperCAmelCase   :			List[Any]=768      ,_UpperCAmelCase   :			Union[str, Any]=12      ,_UpperCAmelCase   :			List[str]=12      ,_UpperCAmelCase   :			Tuple=3072      ,_UpperCAmelCase   :			Dict=0.1      ,_UpperCAmelCase   :			Tuple=0.1      ,_UpperCAmelCase   :			List[str]=1026      ,_UpperCAmelCase   :			List[str]=0.02      ,_UpperCAmelCase   :			Optional[int]=1E-12      ,_UpperCAmelCase   :			List[str]="absolute"      ,_UpperCAmelCase   :			Tuple=True      ,_UpperCAmelCase   :			Tuple=None      ,_UpperCAmelCase   :			List[Any]=False      ,_UpperCAmelCase   :			int=False      ,_UpperCAmelCase   :			int=None      ,_UpperCAmelCase   :			Union[str, Any]=None      ,**_UpperCAmelCase   :			List[Any]      ,):
									super().__init__(pad_token_id=_UpperCAmelCase      ,mask_token_id=_UpperCAmelCase      ,**_UpperCAmelCase						)

									_a	: Optional[Any]								=				vocab_size
									_a	: Union[str, Any]								=				hidden_size
									_a	: Dict								=				num_hidden_layers
									_a	: int								=				num_attention_heads
									_a	: Dict								=				intermediate_size
									_a	: List[Any]								=				hidden_dropout_prob
									_a	: List[Any]								=				attention_probs_dropout_prob
									_a	: Optional[Any]								=				max_position_embeddings
									_a	: Optional[int]								=				initializer_range
									_a	: List[Any]								=				layer_norm_eps
									_a	: int								=				position_embedding_type
									_a	: Optional[int]								=				use_cache
									_a	: Any								=				emb_layer_norm_before
									_a	: List[str]								=				token_dropout
									_a	: List[str]								=				is_folding_model
									if is_folding_model:
											if esmfold_config is None:
													logger.info('No esmfold_config supplied for folding model, using default values.'						)
													_a	: Dict								=				EsmFoldConfig()
											elif isinstance(_UpperCAmelCase      ,_UpperCAmelCase						):
													_a	: Dict								=				EsmFoldConfig(**_UpperCAmelCase						)
											_a	: Optional[int]								=				esmfold_config
											if vocab_list is None:
													logger.warning('No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'						)
													_a	: Optional[int]								=				get_default_vocab_list()
											else:
													_a	: Optional[int]								=				vocab_list
									else:
											_a	: Optional[Any]								=				None
											_a	: Union[str, Any]								=				None
									if self.esmfold_config is not None and getattr(self.esmfold_config      ,'use_esm_attn_map'      ,_UpperCAmelCase						):
											raise ValueError('The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'						)







							def 				__lowercase     (				self   :			Any						):
									_a	: str								=				super().to_dict()
									if isinstance(self.esmfold_config      ,_UpperCAmelCase						):
											_a	: List[str]								=				self.esmfold_config.to_dict()
									return output

@dataclass
class       __magic_name__		:
							lowerCAmelCase			:       str  							=   None
							lowerCAmelCase			:       bool  							=   True
							lowerCAmelCase			:       bool  							=   False
							lowerCAmelCase			:       bool  							=   False
							lowerCAmelCase			:       bool  							=   False
							lowerCAmelCase			:       float  							=   0

							lowerCAmelCase			:       bool  							=   True
							lowerCAmelCase			:       bool  							=   False

							lowerCAmelCase			:       int  							=   1_2_8
							lowerCAmelCase			:       "TrunkConfig"  							=   None







							def 				__lowercase     (				self   :			List[str]						):
									if self.trunk is None:
											_a	: Dict								=				TrunkConfig()
									elif isinstance(self.trunk      ,_UpperCAmelCase						):
											_a	: str								=				TrunkConfig(**self.trunk						)







							def 				__lowercase     (				self   :			List[Any]						):
									_a	: List[str]								=				asdict(self						)
									_a	: List[str]								=				self.trunk.to_dict()
									return output

@dataclass
class       __magic_name__		:
							lowerCAmelCase			:       int  							=   4_8
							lowerCAmelCase			:       int  							=   1_0_2_4
							lowerCAmelCase			:       int  							=   1_2_8
							lowerCAmelCase			:       int  							=   3_2
							lowerCAmelCase			:       int  							=   3_2
							lowerCAmelCase			:       int  							=   3_2
							lowerCAmelCase			:       float  							=   0
							lowerCAmelCase			:       float  							=   0
							lowerCAmelCase			:       bool  							=   False
							lowerCAmelCase			:       int  							=   4
							lowerCAmelCase			:       Optional[int]  							=   1_2_8
							lowerCAmelCase			:       "StructureModuleConfig"  							=   None







							def 				__lowercase     (				self   :			str						):
									if self.structure_module is None:
											_a	: Tuple								=				StructureModuleConfig()
									elif isinstance(self.structure_module      ,_UpperCAmelCase						):
											_a	: List[str]								=				StructureModuleConfig(**self.structure_module						)

									if self.max_recycles <= 0:
											raise ValueError(F"""`max_recycles` should be positive, got {self.max_recycles}."""						)
									if self.sequence_state_dim % self.sequence_state_dim != 0:
											raise ValueError(
											    '`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got'
											    F""" {self.sequence_state_dim} and {self.sequence_state_dim}."""						)
									if self.pairwise_state_dim % self.pairwise_state_dim != 0:
											raise ValueError(
											    '`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got'
											    F""" {self.pairwise_state_dim} and {self.pairwise_state_dim}."""						)

									_a	: Optional[int]								=				self.sequence_state_dim // self.sequence_head_width
									_a	: int								=				self.pairwise_state_dim // self.pairwise_head_width

									if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
											raise ValueError(
											    '`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got'
											    F""" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."""						)
									if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
											raise ValueError(
											    '`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got'
											    F""" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."""						)
									if self.pairwise_state_dim % 2 != 0:
											raise ValueError(F"""`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."""						)

									if self.dropout >= 0.4:
											raise ValueError(F"""`dropout` should not be greater than 0.4, got {self.dropout}."""						)







							def 				__lowercase     (				self   :			Optional[int]						):
									_a	: Optional[Any]								=				asdict(self						)
									_a	: Optional[Any]								=				self.structure_module.to_dict()
									return output

@dataclass
class       __magic_name__		:
							lowerCAmelCase			:       int  							=   3_8_4
							lowerCAmelCase			:       int  							=   1_2_8
							lowerCAmelCase			:       int  							=   1_6
							lowerCAmelCase			:       int  							=   1_2_8
							lowerCAmelCase			:       int  							=   1_2
							lowerCAmelCase			:       int  							=   4
							lowerCAmelCase			:       int  							=   8
							lowerCAmelCase			:       float  							=   0.1
							lowerCAmelCase			:       int  							=   8
							lowerCAmelCase			:       int  							=   1
							lowerCAmelCase			:       int  							=   2
							lowerCAmelCase			:       int  							=   7
							lowerCAmelCase			:       int  							=   1_0
							lowerCAmelCase			:       float  							=   1e-8
							lowerCAmelCase			:       float  							=   1e5







							def 				__lowercase     (				self   :			str						):
									return asdict(self						)



def 	__lowerCamelCase	( )     -> Optional[int]:
		return (
		    "<cls>",
		    "<pad>",
		    "<eos>",
		    "<unk>",
		    "L",
		    "A",
		    "G",
		    "V",
		    "S",
		    "E",
		    "R",
		    "T",
		    "I",
		    "D",
		    "P",
		    "K",
		    "Q",
		    "N",
		    "F",
		    "Y",
		    "M",
		    "H",
		    "W",
		    "C",
		    "X",
		    "B",
		    "U",
		    "Z",
		    "O",
		    ".",
		    "-",
		    "<null_1>",
		    "<mask>",
		)




