





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







def    a__							(							lowercase   :	List[Any] )						->						Dict:





     """simple docstring"""

     _UpperCamelCase							       =			ASTConfig()

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

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

     _UpperCamelCase							       =			json.load(open(hf_hub_download(lowercase,   lowercase,   repo_type='''dataset''' ),   '''r''' ) )
     _UpperCamelCase							       =			{int(lowercase ): v for k, v in idalabel.items()}
     _UpperCamelCase							       =			idalabel
     _UpperCamelCase							       =			{v: k for k, v in idalabel.items()}

     return config







def    a__							(							lowercase   :	Union[str, Any] )						->						str:





     """simple docstring"""

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

     return name







def    a__							(							lowercase   :	Tuple,   lowercase   :	int )						->						str:





     """simple docstring"""

     for key in orig_state_dict.copy().keys():
          _UpperCamelCase							       =			orig_state_dict.pop(lowercase )

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

     return orig_state_dict







def    a__							(							lowercase   :	Union[str, Any] )						->						List[str]:





     """simple docstring"""

     _UpperCamelCase							       =			[
         '''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(lowercase,   lowercase )







@torch.no_grad()
def    a__							(							lowercase   :	List[str],   lowercase   :	int,   lowercase   :	str=False )						->						int:





     """simple docstring"""

     _UpperCamelCase							       =			get_audio_spectrogram_transformer_config(lowercase )

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

     # load 🤗 model
     _UpperCamelCase							       =			ASTForAudioClassification(lowercase )
     model.eval()

     model.load_state_dict(lowercase )

     # verify outputs on dummy input
     # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
     _UpperCamelCase							       =			-4.2_6_7_7_3_9_3 if '''speech-commands''' not in model_name else -6.8_4_5_9_7_8
     _UpperCamelCase							       =			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							       =			1024 if '''speech-commands''' not in model_name else 128
     _UpperCamelCase							       =			ASTFeatureExtractor(mean=lowercase,   std=lowercase,   max_length=lowercase )

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

          _UpperCamelCase       ,							_UpperCamelCase							       =			torchaudio.load(lowercase )
          _UpperCamelCase							       =			waveform.squeeze().numpy()

     _UpperCamelCase							       =			feature_extractor(lowercase,   sampling_rate=16000,   return_tensors='''pt''' )

     # forward pass
     _UpperCamelCase							       =			model(**lowercase )
     _UpperCamelCase							       =			outputs.logits

     if model_name == "ast-finetuned-audioset-10-10-0.4593":
          _UpperCamelCase							       =			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							       =			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							       =			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							       =			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							       =			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							       =			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							       =			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							       =			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],   lowercase,   atol=1e-4 ):
          raise ValueError('''Logits don\'t match''' )
     print('''Looks ok!''' )

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

     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__":
   lowercase__		:     List[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.'
   )

   lowercase__		:     int			=							parser.parse_args()
   convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''

from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

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


if TYPE_CHECKING:
   from ...feature_extraction_utils import FeatureExtractionMixin
   from ...tokenization_utils_base import PreTrainedTokenizerBase
   from ...utils import TensorType

lowercase__		:     Optional[Any]			=							logging.get_logger(__name__)

lowercase__		:     Union[str, Any]			=							{
    'openai/whisper-base': 'https://huggingface.co/openai/whisper-base/resolve/main/config.json',
}

# fmt: off
lowercase__		:     Any			=							[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 3_57, 3_66, 4_38, 5_32, 6_85,
    7_05, 7_96, 9_30, 10_58, 12_20, 12_67, 12_79, 13_03, 13_43, 13_77,
    13_91, 16_35, 17_82, 18_75, 21_62, 23_61, 24_88, 34_67, 40_08, 42_11,
    46_00, 48_08, 52_99, 58_55, 63_29, 72_03, 96_09, 99_59, 1_05_63, 1_07_86,
    1_14_20, 1_17_09, 1_19_07, 1_31_63, 1_36_97, 1_37_00, 1_48_08, 1_53_06, 1_64_10, 1_67_91,
    1_79_92, 1_92_03, 1_95_10, 2_07_24, 2_23_05, 2_29_35, 2_70_07, 3_01_09, 3_04_20, 3_34_09,
    3_49_49, 4_02_83, 4_04_93, 4_05_49, 4_72_82, 4_91_46, 5_02_57, 5_03_59, 5_03_60, 5_03_61
]
lowercase__		:     str			=							[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 3_59, 5_03, 5_22, 5_42, 8_73,
    8_93, 9_02, 9_18, 9_22, 9_31, 13_50, 18_53, 19_82, 24_60, 26_27,
    32_46, 32_53, 32_68, 35_36, 38_46, 39_61, 41_83, 46_67, 65_85, 66_47,
    72_73, 90_61, 93_83, 1_04_28, 1_09_29, 1_19_38, 1_20_33, 1_23_31, 1_25_62, 1_37_93,
    1_41_57, 1_46_35, 1_52_65, 1_56_18, 1_65_53, 1_66_04, 1_83_62, 1_89_56, 2_00_75, 2_16_75,
    2_25_20, 2_61_30, 2_61_61, 2_64_35, 2_82_79, 2_94_64, 3_16_50, 3_23_02, 3_24_70, 3_68_65,
    4_28_63, 4_74_25, 4_98_70, 5_02_54, 5_02_58, 5_03_60, 5_03_61, 5_03_62
]
class 							__lowerCAmelCase	(       __magic_name__					):






       """simple docstring"""
       _snake_case      :	Optional[int]									=       'whisper'
       _snake_case      :	Union[str, Any]									=       ['past_key_values']
       _snake_case      :	Union[str, Any]									=       {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}





       def __init__(				self  :						Tuple			,					lowerCAmelCase__  :						List[str]=51865			,					lowerCAmelCase__  :						Optional[Any]=80			,					lowerCAmelCase__  :						Optional[int]=6			,					lowerCAmelCase__  :						List[Any]=4			,					lowerCAmelCase__  :						Any=6			,					lowerCAmelCase__  :						List[Any]=4			,					lowerCAmelCase__  :						List[Any]=1536			,					lowerCAmelCase__  :						Dict=1536			,					lowerCAmelCase__  :						int=0.0			,					lowerCAmelCase__  :						Optional[Any]=0.0			,					lowerCAmelCase__  :						int=50257			,					lowerCAmelCase__  :						Tuple=True			,					lowerCAmelCase__  :						Optional[Any]=True			,					lowerCAmelCase__  :						Optional[int]="gelu"			,					lowerCAmelCase__  :						Optional[int]=256			,					lowerCAmelCase__  :						Any=0.0			,					lowerCAmelCase__  :						Tuple=0.0			,					lowerCAmelCase__  :						Optional[Any]=0.0			,					lowerCAmelCase__  :						str=0.02			,					lowerCAmelCase__  :						Dict=False			,					lowerCAmelCase__  :						int=1500			,					lowerCAmelCase__  :						int=448			,					lowerCAmelCase__  :						Union[str, Any]=50256			,					lowerCAmelCase__  :						str=50256			,					lowerCAmelCase__  :						Optional[int]=50256			,					lowerCAmelCase__  :						List[str]=None			,					lowerCAmelCase__  :						Optional[int]=[220, 50256]			,					lowerCAmelCase__  :						List[Any]=False			,					lowerCAmelCase__  :						str=256			,					lowerCAmelCase__  :						Optional[Any]=False			,					lowerCAmelCase__  :						Any=0.05			,					lowerCAmelCase__  :						Optional[int]=10			,					lowerCAmelCase__  :						Dict=2			,					lowerCAmelCase__  :						int=0.0			,					lowerCAmelCase__  :						Dict=10			,					lowerCAmelCase__  :						List[str]=0			,					lowerCAmelCase__  :						List[Any]=7			,					**lowerCAmelCase__  :						List[str]			,					)       ->	Any:

            '''simple docstring'''




            _UpperCamelCase							       =			vocab_size
            _UpperCamelCase							       =			num_mel_bins
            _UpperCamelCase							       =			d_model
            _UpperCamelCase							       =			encoder_layers
            _UpperCamelCase							       =			encoder_attention_heads
            _UpperCamelCase							       =			decoder_layers
            _UpperCamelCase							       =			decoder_attention_heads
            _UpperCamelCase							       =			decoder_ffn_dim
            _UpperCamelCase							       =			encoder_ffn_dim
            _UpperCamelCase							       =			dropout
            _UpperCamelCase							       =			attention_dropout
            _UpperCamelCase							       =			activation_dropout
            _UpperCamelCase							       =			activation_function
            _UpperCamelCase							       =			init_std
            _UpperCamelCase							       =			encoder_layerdrop
            _UpperCamelCase							       =			decoder_layerdrop
            _UpperCamelCase							       =			use_cache
            _UpperCamelCase							       =			encoder_layers
            _UpperCamelCase							       =			scale_embedding  # scale factor will be sqrt(d_model) if True
            _UpperCamelCase							       =			max_source_positions
            _UpperCamelCase							       =			max_target_positions

            # Audio Classification-specific parameters. Feel free to ignore for other classes.
            _UpperCamelCase							       =			classifier_proj_size
            _UpperCamelCase							       =			use_weighted_layer_sum

            # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
            _UpperCamelCase							       =			apply_spec_augment
            _UpperCamelCase							       =			mask_time_prob
            _UpperCamelCase							       =			mask_time_length
            _UpperCamelCase							       =			mask_time_min_masks
            _UpperCamelCase							       =			mask_feature_prob
            _UpperCamelCase							       =			mask_feature_length
            _UpperCamelCase							       =			mask_feature_min_masks

            _UpperCamelCase							       =			median_filter_width

            super().__init__(
                pad_token_id=lowerCAmelCase__			,					bos_token_id=lowerCAmelCase__			,					eos_token_id=lowerCAmelCase__			,					is_encoder_decoder=lowerCAmelCase__			,					decoder_start_token_id=lowerCAmelCase__			,					suppress_tokens=lowerCAmelCase__			,					begin_suppress_tokens=lowerCAmelCase__			,					**lowerCAmelCase__			,					)

class 							__lowerCAmelCase	(       __magic_name__					):






       """simple docstring"""





       @property
       def snake_case__		(				self  :						Optional[Any]	)       ->	Mapping[str, Mapping[int, str]]:

            '''simple docstring'''




            _UpperCamelCase							       =			OrderedDict(
                [
                    ('''input_features''', {0: '''batch''', 1: '''feature_size''', 2: '''encoder_sequence'''}),
                ]	)
            if self.use_past:
                 _UpperCamelCase							       =			{0: '''batch'''}
            else:
                 _UpperCamelCase							       =			{0: '''batch''', 1: '''decoder_sequence'''}

            if self.use_past:
                 self.fill_with_past_key_values_(lowerCAmelCase__			,					direction='''inputs'''	)

            return common_inputs





       def snake_case__		(				self  :						Dict			,					lowerCAmelCase__  :						Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]			,					lowerCAmelCase__  :						int = -1			,					lowerCAmelCase__  :						int = -1			,					lowerCAmelCase__  :						bool = False			,					lowerCAmelCase__  :						Optional["TensorType"] = None			,					lowerCAmelCase__  :						int = 22050			,					lowerCAmelCase__  :						float = 5.0			,					lowerCAmelCase__  :						int = 220			,					)       ->	Mapping[str, Any]:

            '''simple docstring'''




            _UpperCamelCase							       =			OrderedDict()
            _UpperCamelCase							       =			OnnxConfig.generate_dummy_inputs(
                self			,					preprocessor=preprocessor.feature_extractor			,					batch_size=lowerCAmelCase__			,					framework=lowerCAmelCase__			,					sampling_rate=lowerCAmelCase__			,					time_duration=lowerCAmelCase__			,					frequency=lowerCAmelCase__			,					)
            _UpperCamelCase							       =			encoder_inputs['''input_features'''].shape[2]
            _UpperCamelCase							       =			encoder_sequence_length // 2 if self.use_past else seq_length

            _UpperCamelCase							       =			super().generate_dummy_inputs(
                preprocessor.tokenizer			,					lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__	)

            _UpperCamelCase							       =			encoder_inputs.pop('''input_features'''	)
            _UpperCamelCase							       =			decoder_inputs.pop('''decoder_input_ids'''	)

            if "past_key_values" in decoder_inputs:
                 _UpperCamelCase							       =			decoder_inputs.pop('''past_key_values'''	)

            return dummy_inputs







       @property
       def snake_case__		(				self  :						List[Any]	)       ->	float:

            '''simple docstring'''




            return 1e-3



