


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_      : Any  					=       logging.get_logger(__name__)







def        SCREAMING_SNAKE_CASE_						(			__A :			Union[str, Any]				)	->		str:

       """simple docstring"""

       a_					:		Optional[Any]							=  ASTConfig()

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

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

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

       return config







def        SCREAMING_SNAKE_CASE_						(			__A :			Union[str, Any]				)	->		Optional[int]:

       """simple docstring"""

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

       return name







def        SCREAMING_SNAKE_CASE_						(			__A :			List[str]  ,    __A :			List[Any]				)	->		Optional[int]:

       """simple docstring"""

       for key in orig_state_dict.copy().keys():
              a_					:		Dict							=  orig_state_dict.pop(__A				)

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

       return orig_state_dict







def        SCREAMING_SNAKE_CASE_						(			__A :			List[Any]				)	->		Union[str, Any]:

       """simple docstring"""

       a_					:		Optional[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(__A  ,    __A				)







@torch.no_grad()
def        SCREAMING_SNAKE_CASE_						(			__A :			Optional[Any]  ,    __A :			List[str]  ,    __A :			Union[str, Any]=False				)	->		Optional[int]:

       """simple docstring"""

       a_					:		Dict							=  get_audio_spectrogram_transformer_config(__A				)

       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_					:		Dict							=  model_name_to_url[model_name]
       a_					:		Union[str, Any]							=  torch.hub.load_state_dict_from_url(__A  ,    map_location='cpu'				)
       # remove some keys
       remove_keys(__A				)
       # rename some keys
       a_					:		List[str]							=  convert_state_dict(__A  ,    __A				)

       # load 🤗 model
       a_					:		List[str]							=  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_					:		Union[str, 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_					:		List[str]							=  10_24 if 'speech-commands' not in model_name else 1_28
       a_					:		Dict							=  ASTFeatureExtractor(mean=__A  ,    std=__A  ,    max_length=__A				)

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

              a_							,		a_					:		Any							=  torchaudio.load(__A				)
              a_					:		Dict							=  waveform.squeeze().numpy()

       a_					:		Tuple							=  feature_extractor(__A  ,    sampling_rate=1_60_00  ,    return_tensors='pt'				)

       # forward pass
       a_					:		Union[str, Any]							=  model(**__A				)
       a_					:		Optional[Any]							=  outputs.logits

       if model_name == "ast-finetuned-audioset-10-10-0.4593":
              a_					:		List[Any]							=  torch.tensor([-0.8760, -7.0042, -8.6602]				)
       elif model_name == "ast-finetuned-audioset-10-10-0.450":
              a_					:		Dict							=  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_					:		str							=  torch.tensor([-1.5080, -7.4534, -8.8917]				)
       elif model_name == "ast-finetuned-audioset-12-12-0.447":
              a_					:		Tuple							=  torch.tensor([-0.5050, -6.5833, -8.0843]				)
       elif model_name == "ast-finetuned-audioset-14-14-0.443":
              a_					:		int							=  torch.tensor([-0.3826, -7.0336, -8.2413]				)
       elif model_name == "ast-finetuned-audioset-16-16-0.442":
              a_					:		str							=  torch.tensor([-1.2113, -6.9101, -8.3470]				)
       elif model_name == "ast-finetuned-speech-commands-v2":
              a_					:		List[str]							=  torch.tensor([6.1589, -8.0566, -8.7984]				)
       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_      : str  					=       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_      : str  					=       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


UpperCAmelCase_      : str  					=       logging.get_logger(__name__)

UpperCAmelCase_      : List[str]  					=       {
    'YituTech/conv-bert-base': 'https://huggingface.co/YituTech/conv-bert-base/resolve/main/config.json',
    'YituTech/conv-bert-medium-small': (
        'https://huggingface.co/YituTech/conv-bert-medium-small/resolve/main/config.json'
    ),
    'YituTech/conv-bert-small': 'https://huggingface.co/YituTech/conv-bert-small/resolve/main/config.json',
    # See all ConvBERT models at https://huggingface.co/models?filter=convbert
}
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     Optional[int]        =						'''convbert'''
      def __init__(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=3_0_5_2_2     ,							SCREAMING_SNAKE_CASE__				:       Dict=7_6_8     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=1_2     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=1_2     ,							SCREAMING_SNAKE_CASE__				:       str=3_0_7_2     ,							SCREAMING_SNAKE_CASE__				:       Dict="gelu"     ,							SCREAMING_SNAKE_CASE__				:       Dict=0.1     ,							SCREAMING_SNAKE_CASE__				:       Tuple=0.1     ,							SCREAMING_SNAKE_CASE__				:       List[str]=5_1_2     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=2     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=0.02     ,							SCREAMING_SNAKE_CASE__				:       Any=1E-12     ,							SCREAMING_SNAKE_CASE__				:       int=1     ,							SCREAMING_SNAKE_CASE__				:       int=0     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=2     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=7_6_8     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=2     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=9     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=1     ,							SCREAMING_SNAKE_CASE__				:       Dict=None     ,							**SCREAMING_SNAKE_CASE__				:       List[Any]     ,							)		->			Any:
             super().__init__(
                 pad_token_id=SCREAMING_SNAKE_CASE__     ,							bos_token_id=SCREAMING_SNAKE_CASE__     ,							eos_token_id=SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     ,							)

             a_					:		Tuple							=  vocab_size
             a_					:		List[str]							=  hidden_size
             a_					:		List[str]							=  num_hidden_layers
             a_					:		Dict							=  num_attention_heads
             a_					:		Optional[int]							=  intermediate_size
             a_					:		int							=  hidden_act
             a_					:		Dict							=  hidden_dropout_prob
             a_					:		int							=  attention_probs_dropout_prob
             a_					:		str							=  max_position_embeddings
             a_					:		List[str]							=  type_vocab_size
             a_					:		List[str]							=  initializer_range
             a_					:		Tuple							=  layer_norm_eps
             a_					:		Optional[int]							=  embedding_size
             a_					:		List[Any]							=  head_ratio
             a_					:		List[Any]							=  conv_kernel_size
             a_					:		Tuple							=  num_groups
             a_					:		Tuple							=  classifier_dropout
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Mapping[str, Mapping[int, str]]:
             if self.task == "multiple-choice":
                    a_					:		Tuple							=  {0: 'batch', 1: 'choice', 2: 'sequence'}
             else:
                    a_					:		List[str]							=  {0: 'batch', 1: 'sequence'}
             return OrderedDict(
                 [
                     ('input_ids', dynamic_axis),
                     ('attention_mask', dynamic_axis),
                     ('token_type_ids', dynamic_axis),
                 ]     )


