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





def 					UpperCAmelCase_		(__a      :    Tuple ):






   """simple docstring"""

   _a				:	Optional[Any]			=   ASTConfig()

   if "10-10" in model_name:
      pass
   elif "speech-commands" in model_name:
      _a				:	Any			=   1_2_8
   elif "12-12" in model_name:
      _a				:	Dict			=   1_2
      _a				:	Any			=   1_2
   elif "14-14" in model_name:
      _a				:	str			=   1_4
      _a				:	int			=   1_4
   elif "16-16" in model_name:
      _a				:	int			=   1_6
      _a				:	Optional[int]			=   1_6
   else:
      raise ValueError('Model not supported' )

   _a				:	List[Any]			=   '''huggingface/label-files'''
   if "speech-commands" in model_name:
      _a				:	Tuple			=   3_5
      _a				:	Any			=   '''speech-commands-v2-id2label.json'''
   else:
      _a				:	Tuple			=   5_2_7
      _a				:	Union[str, Any]			=   '''audioset-id2label.json'''

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

   return config





def 					UpperCAmelCase_		(__a      :    Optional[int] ):






   """simple docstring"""

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

   return name





def 					UpperCAmelCase_		(__a      :    Tuple ,	__a      :    Any ):






   """simple docstring"""

   for key in orig_state_dict.copy().keys():
      _a				:	Optional[Any]			=   orig_state_dict.pop(__lowerCAmelCase )

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

   return orig_state_dict





def 					UpperCAmelCase_		(__a      :    List[Any] ):






   """simple docstring"""

   _a				:	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(__lowerCAmelCase ,	__lowerCAmelCase )





@torch.no_grad()
def 					UpperCAmelCase_		(__a      :    Optional[int] ,	__a      :    Any ,	__a      :    Dict=False ):






   """simple docstring"""

   _a				:	Optional[Any]			=   get_audio_spectrogram_transformer_config(__lowerCAmelCase )

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

   # load 🤗 model
   _a				:	str			=   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
   _a				:	List[Any]			=   -4.2677393 if '''speech-commands''' not in model_name else -6.845978
   _a				:	Union[str, Any]			=   4.5689974 if '''speech-commands''' not in model_name else 5.5654526
   _a				:	Any			=   1_0_2_4 if '''speech-commands''' not in model_name else 1_2_8
   _a				:	Tuple			=   ASTFeatureExtractor(mean=__lowerCAmelCase ,	std=__lowerCAmelCase ,	max_length=__lowerCAmelCase )

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

      _a				:	List[Any]			=   torchaudio.load(__lowerCAmelCase )
      _a				:	List[Any]			=   waveform.squeeze().numpy()

   _a				:	int			=   feature_extractor(__lowerCAmelCase ,	sampling_rate=1_6_0_0_0 ,	return_tensors='pt' )

   # forward pass
   _a				:	List[str]			=   model(**__lowerCAmelCase )
   _a				:	Any			=   outputs.logits

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

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








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__								=						logging.get_logger(__name__)

A__								=						{
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}


class       a							(			__lowerCamelCase				):
		__lowerCAmelCase	:    str											=					"""distilbert"""
		__lowerCAmelCase	:    str											=					{
		    """hidden_size""": """dim""",
		    """num_attention_heads""": """n_heads""",
		    """num_hidden_layers""": """n_layers""",
		}



		def __init__(     self     :Dict   ,__lowercase     :Tuple=3_0_5_2_2   ,__lowercase     :Optional[Any]=5_1_2   ,__lowercase     :List[str]=False   ,__lowercase     :List[str]=6   ,__lowercase     :Optional[Any]=1_2   ,__lowercase     :Tuple=7_6_8   ,__lowercase     :int=4 * 7_6_8   ,__lowercase     :List[Any]=0.1   ,__lowercase     :List[str]=0.1   ,__lowercase     :Union[str, Any]="gelu"   ,__lowercase     :List[str]=0.02   ,__lowercase     :Optional[int]=0.1   ,__lowercase     :Dict=0.2   ,__lowercase     :Union[str, Any]=0   ,**__lowercase     :Optional[Any]   ,):
						snake_case__	:      List[str]     		=	vocab_size
						snake_case__	:      Tuple     		=	max_position_embeddings
						snake_case__	:      Optional[int]     		=	sinusoidal_pos_embds
						snake_case__	:      str     		=	n_layers
						snake_case__	:      List[Any]     		=	n_heads
						snake_case__	:      Tuple     		=	dim
						snake_case__	:      str     		=	hidden_dim
						snake_case__	:      int     		=	dropout
						snake_case__	:      Dict     		=	attention_dropout
						snake_case__	:      Tuple     		=	activation
						snake_case__	:      int     		=	initializer_range
						snake_case__	:      Optional[Any]     		=	qa_dropout
						snake_case__	:      Union[str, Any]     		=	seq_classif_dropout
						super().__init__(**__lowercase   ,pad_token_id=__lowercase  )


class       a							(			__lowerCamelCase				):



		@property
		def 	__lowerCamelCase			(     self     :Union[str, Any]  ):
						if self.task == "multiple-choice":
										snake_case__	:      str     		=	{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
						else:
										snake_case__	:      Optional[int]     		=	{0: '''batch''', 1: '''sequence'''}
						return OrderedDict(
						    [
						        ('''input_ids''', dynamic_axis),
						        ('''attention_mask''', dynamic_axis),
						    ]  )

