




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

def  UpperCAmelCase_							(				__UpperCAmelCase						:     List[str]							) ->			List[Any]:
      SCREAMING_SNAKE_CASE_     						=      ASTConfig()

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

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

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

      return config

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

      return name

def  UpperCAmelCase_							(				__UpperCAmelCase						:     str      ,      __UpperCAmelCase						:     Optional[int]							) ->			Dict:
      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_     						=      orig_state_dict.pop(__UpperCAmelCase							)

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

      return orig_state_dict

def  UpperCAmelCase_							(				__UpperCAmelCase						:     int							) ->			str:
      SCREAMING_SNAKE_CASE_     						=      [
          '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(__UpperCAmelCase      ,      __UpperCAmelCase							)

@torch.no_grad()
def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]      ,      __UpperCAmelCase						:     Any      ,      __UpperCAmelCase						:     List[Any]=False							) ->			Dict:
      SCREAMING_SNAKE_CASE_     						=      get_audio_spectrogram_transformer_config(__UpperCAmelCase							)

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

      # load 🤗 model
      SCREAMING_SNAKE_CASE_     						=      ASTForAudioClassification(__UpperCAmelCase							)
      model.eval()

      model.load_state_dict(__UpperCAmelCase							)

      # verify outputs on dummy input
      # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
      SCREAMING_SNAKE_CASE_     						=      -4.2_6_7_7_3_9_3 if 'speech-commands' not in model_name else -6.8_4_5_9_7_8
      SCREAMING_SNAKE_CASE_     						=      4.5_6_8_9_9_7_4 if 'speech-commands' not in model_name else 5.5_6_5_4_5_2_6
      SCREAMING_SNAKE_CASE_     						=      10_24 if 'speech-commands' not in model_name else 1_28
      SCREAMING_SNAKE_CASE_     						=      ASTFeatureExtractor(mean=__UpperCAmelCase      ,      std=__UpperCAmelCase      ,      max_length=__UpperCAmelCase							)

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

            SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      torchaudio.load(__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_     						=      waveform.squeeze().numpy()

      SCREAMING_SNAKE_CASE_     						=      feature_extractor(__UpperCAmelCase      ,      sampling_rate=1_60_00      ,      return_tensors='pt'							)

      # forward pass
      SCREAMING_SNAKE_CASE_     						=      model(**__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_     						=      outputs.logits

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

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

      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__			:       Union[str, 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.'
   )

   lowerCamelCase__			:       Tuple         =							parser.parse_args()
   convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402


lowerCamelCase__			:       str         =							get_tests_dir('fixtures/dummy-config.json')





class 	lowerCamelCase_							(				unittest.TestCase ):


 '''simple docstring'''







 def 				lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_     						=      0



 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       self.assertIsNotNone(transformers.models.auto.__spec__						)
       self.assertIsNotNone(importlib.util.find_spec('transformers.models.auto'						)						)



 def 				lowerCAmelCase_  (							self  :		Any						):
       SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('bert-base-uncased'						)
       self.assertIsInstance(_lowerCAmelCase    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase						)
       self.assertIsInstance(_lowerCAmelCase    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase						)
       self.assertIsInstance(_lowerCAmelCase    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_     						=      AutoConfig.for_model('roberta'						)
       self.assertIsInstance(_lowerCAmelCase    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Dict						):
       with tempfile.TemporaryDirectory() as tmp_dir:
             # This model name contains bert and roberta, but roberta ends up being picked.
             SCREAMING_SNAKE_CASE_     						=      os.path.join(_lowerCAmelCase    ,    'fake-roberta'						)
             os.makedirs(_lowerCAmelCase    ,    exist_ok=_lowerCAmelCase						)
             with open(os.path.join(_lowerCAmelCase    ,    'config.json'						)    ,    'w'						) as f:
                   f.write(json.dumps({}						)						)
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase						)
             self.assertEqual(type(_lowerCAmelCase						)    ,    _lowerCAmelCase						)



 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       try:
             AutoConfig.register('custom'    ,    _lowerCAmelCase						)
             # Wrong model type will raise an error
             with self.assertRaises(_lowerCAmelCase						):
                   AutoConfig.register('model'    ,    _lowerCAmelCase						)
             # Trying to register something existing in the Transformers library will raise an error
             with self.assertRaises(_lowerCAmelCase						):
                   AutoConfig.register('bert'    ,    _lowerCAmelCase						)

             # Now that the config is registered, it can be used as any other config with the auto-API
             SCREAMING_SNAKE_CASE_     						=      CustomConfig()
             with tempfile.TemporaryDirectory() as tmp_dir:
                   config.save_pretrained(_lowerCAmelCase						)
                   SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase						)
                   self.assertIsInstance(_lowerCAmelCase    ,    _lowerCAmelCase						)

       finally:
             if "custom" in CONFIG_MAPPING._extra_content:
                   del CONFIG_MAPPING._extra_content["custom"]



 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       with self.assertRaisesRegex(
           _lowerCAmelCase    ,    'bert-base is not a local folder and is not a valid model identifier'						):
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('bert-base'						)



 def 				lowerCAmelCase_  (							self  :		int						):
       with self.assertRaisesRegex(
           _lowerCAmelCase    ,    R'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'						):
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase    ,    revision='aaaaaa'						)



 def 				lowerCAmelCase_  (							self  :		Tuple						):
       with self.assertRaisesRegex(
           _lowerCAmelCase    ,    'hf-internal-testing/no-config-test-repo does not appear to have a file named config.json.'    ,    ):
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/no-config-test-repo'						)



 def 				lowerCAmelCase_  (							self  :		Union[str, Any]						):
       # If remote code is not set, we will time out when asking whether to load the model.
       with self.assertRaises(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'						)
       # If remote code is disabled, we can't load this config.
       with self.assertRaises(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'    ,    trust_remote_code=_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'    ,    trust_remote_code=_lowerCAmelCase						)
       self.assertEqual(config.__class__.__name__    ,    'NewModelConfig'						)

       # Test config can be reloaded.
       with tempfile.TemporaryDirectory() as tmp_dir:
             config.save_pretrained(_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained(_lowerCAmelCase    ,    trust_remote_code=_lowerCAmelCase						)
       self.assertEqual(reloaded_config.__class__.__name__    ,    'NewModelConfig'						)




 def 				lowerCAmelCase_  (							self  :		Any						):





       class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE ):


        '''simple docstring'''







        lowercase_       =   "new-model"

       try:
             AutoConfig.register('new-model'    ,    _lowerCAmelCase						)
             # If remote code is not set, the default is to use local
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'						)
             self.assertEqual(config.__class__.__name__    ,    'NewModelConfigLocal'						)

             # If remote code is disabled, we load the local one.
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'    ,    trust_remote_code=_lowerCAmelCase						)
             self.assertEqual(config.__class__.__name__    ,    'NewModelConfigLocal'						)

             # If remote is enabled, we load from the Hub
             SCREAMING_SNAKE_CASE_     						=      AutoConfig.from_pretrained('hf-internal-testing/test_dynamic_model'    ,    trust_remote_code=_lowerCAmelCase						)
             self.assertEqual(config.__class__.__name__    ,    'NewModelConfig'						)

       finally:
             if "new-model" in CONFIG_MAPPING._extra_content:
                   del CONFIG_MAPPING._extra_content["new-model"]