




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







def        lowerCAmelCase_			(  snake_case_  ):
      _A							:  List[str]    =		ASTConfig()

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

      _A							:  Tuple    =		"""huggingface/label-files"""
      if "speech-commands" in model_name:
            _A							:  Optional[Any]    =		35
            _A							:  Optional[Any]    =		"""speech-commands-v2-id2label.json"""
      else:
            _A							:  Optional[int]    =		527
            _A							:  int    =		"""audioset-id2label.json"""

      _A							:  str    =		json.load(open(hf_hub_download(snake_case_,snake_case_,repo_type="""dataset"""  ),"""r"""  )  )
      _A							:  List[str]    =		{int(snake_case_  ): v for k, v in idalabel.items()}
      _A							:  Optional[int]    =		idalabel
      _A							:  Tuple    =		{v: k for k, v in idalabel.items()}

      return config







def        lowerCAmelCase_			(  snake_case_  ):
      if "module.v" in name:
            _A							:  Union[str, Any]    =		name.replace("""module.v""","""audio_spectrogram_transformer"""  )
      if "cls_token" in name:
            _A							:  Dict    =		name.replace("""cls_token""","""embeddings.cls_token"""  )
      if "dist_token" in name:
            _A							:  List[str]    =		name.replace("""dist_token""","""embeddings.distillation_token"""  )
      if "pos_embed" in name:
            _A							:  List[str]    =		name.replace("""pos_embed""","""embeddings.position_embeddings"""  )
      if "patch_embed.proj" in name:
            _A							:  Tuple    =		name.replace("""patch_embed.proj""","""embeddings.patch_embeddings.projection"""  )
      # transformer blocks
      if "blocks" in name:
            _A							:  Tuple    =		name.replace("""blocks""","""encoder.layer"""  )
      if "attn.proj" in name:
            _A							:  List[Any]    =		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							:  List[str]    =		name.replace("""norm1""","""layernorm_before"""  )
      if "norm2" in name:
            _A							:  Optional[Any]    =		name.replace("""norm2""","""layernorm_after"""  )
      if "mlp.fc1" in name:
            _A							:  str    =		name.replace("""mlp.fc1""","""intermediate.dense"""  )
      if "mlp.fc2" in name:
            _A							:  Dict    =		name.replace("""mlp.fc2""","""output.dense"""  )
      # final layernorm
      if "audio_spectrogram_transformer.norm" in name:
            _A							:  Tuple    =		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							:  str    =		name.replace("""module.mlp_head.1""","""classifier.dense"""  )

      return name







def        lowerCAmelCase_			(  snake_case_,snake_case_  ):
      for key in orig_state_dict.copy().keys():
            _A							:  int    =		orig_state_dict.pop(snake_case_  )

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

      return orig_state_dict







def        lowerCAmelCase_			(  snake_case_  ):
      _A							:  Union[str, 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(snake_case_,snake_case_  )







@torch.no_grad()
def        lowerCAmelCase_			(  snake_case_,snake_case_,snake_case_=False  ):
      _A							:  Any    =		get_audio_spectrogram_transformer_config(snake_case_  )

      _A							:  str    =		{
          """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[Any]    =		model_name_to_url[model_name]
      _A							:  Dict    =		torch.hub.load_state_dict_from_url(snake_case_,map_location="""cpu"""  )
      # remove some keys
      remove_keys(snake_case_  )
      # rename some keys
      _A							:  Tuple    =		convert_state_dict(snake_case_,snake_case_  )

      # load 🤗 model
      _A							:  str    =		ASTForAudioClassification(snake_case_  )
      model.eval()

      model.load_state_dict(snake_case_  )

      # verify outputs on dummy input
      # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
      _A							:  Any    =		-4.2_67_73_93 if """speech-commands""" not in model_name else -6.84_59_78
      _A							:  Tuple    =		4.5_68_99_74 if """speech-commands""" not in model_name else 5.5_65_45_26
      _A							:  Any    =		1024 if """speech-commands""" not in model_name else 128
      _A							:  Union[str, Any]    =		ASTFeatureExtractor(mean=snake_case_,std=snake_case_,max_length=snake_case_  )

      if "speech-commands" in model_name:
            _A							:  int    =		load_dataset("""speech_commands""","""v0.02""",split="""validation"""  )
            _A							:  Optional[int]    =		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							:  List[Any]    =		torchaudio.load(snake_case_  )
            _A							:  Union[str, Any]    =		waveform.squeeze().numpy()

      _A							:  List[str]    =		feature_extractor(snake_case_,sampling_rate=16000,return_tensors="""pt"""  )

      # forward pass
      _A							:  str    =		model(**snake_case_  )
      _A							:  Optional[Any]    =		outputs.logits

      if model_name == "ast-finetuned-audioset-10-10-0.4593":
            _A							:  Optional[int]    =		torch.tensor([-0.87_60, -7.00_42, -8.66_02]  )
      elif model_name == "ast-finetuned-audioset-10-10-0.450":
            _A							:  Any    =		torch.tensor([-1.19_86, -7.09_03, -8.27_18]  )
      elif model_name == "ast-finetuned-audioset-10-10-0.448":
            _A							:  Any    =		torch.tensor([-2.61_28, -8.00_80, -9.43_44]  )
      elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
            _A							:  Any    =		torch.tensor([-1.50_80, -7.45_34, -8.89_17]  )
      elif model_name == "ast-finetuned-audioset-12-12-0.447":
            _A							:  str    =		torch.tensor([-0.50_50, -6.58_33, -8.08_43]  )
      elif model_name == "ast-finetuned-audioset-14-14-0.443":
            _A							:  int    =		torch.tensor([-0.38_26, -7.03_36, -8.24_13]  )
      elif model_name == "ast-finetuned-audioset-16-16-0.442":
            _A							:  Union[str, Any]    =		torch.tensor([-1.21_13, -6.91_01, -8.34_70]  )
      elif model_name == "ast-finetuned-speech-commands-v2":
            _A							:  int    =		torch.tensor([6.15_89, -8.05_66, -8.79_84]  )
      else:
            raise ValueError("""Unknown model name"""  )
      if not torch.allclose(logits[0, :3],snake_case_,atol=1e-4  ):
            raise ValueError("""Logits don't match"""  )
      print("""Looks ok!"""  )

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

      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__":
  _snake_case            = 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."
  )

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







import unittest

import numpy as np

from transformers.file_utils import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import DPTImageProcessor







class 		lowercase							(     unittest.TestCase	):

     def __init__(						self	,  _a	,  _a=7	,  _a=3	,  _a=18	,  _a=30	,  _a=400	,  _a=True	,  _a=None	,  _a=True	,  _a=[0.5, 0.5, 0.5]	,  _a=[0.5, 0.5, 0.5]	,  )		->							Tuple:
           _A							:  Any    =		size if size is not None else {"""height""": 18, """width""": 18}
           _A							:  Optional[Any]    =		parent
           _A							:  Union[str, Any]    =		batch_size
           _A							:  List[Any]    =		num_channels
           _A							:  List[str]    =		image_size
           _A							:  Optional[Any]    =		min_resolution
           _A							:  List[Any]    =		max_resolution
           _A							:  Optional[Any]    =		do_resize
           _A							:  str    =		size
           _A							:  List[str]    =		do_normalize
           _A							:  Dict    =		image_mean
           _A							:  int    =		image_std



     def       a__			(						self       )		->							Any:
           return {
               "image_mean": self.image_mean,
               "image_std": self.image_std,
               "do_normalize": self.do_normalize,
               "do_resize": self.do_resize,
               "size": self.size,
           }







@require_torch
@require_vision
class 		lowercase							(     UpperCamelCase__,unittest.TestCase	):
     _a											=		DPTImageProcessor if is_vision_available() else None

     def       a__			(						self       )		->							Optional[int]:
           _A							:  Optional[Any]    =		DPTImageProcessingTester(self       )

     @property
     def       a__			(						self       )		->							Tuple:
           return self.image_processor_tester.prepare_image_processor_dict()

     def       a__			(						self       )		->							Any:
           _A							:  List[str]    =		self.image_processing_class(**self.image_processor_dict       )
           self.assertTrue(hasattr(_a	,  """image_mean"""       )       )
           self.assertTrue(hasattr(_a	,  """image_std"""       )       )
           self.assertTrue(hasattr(_a	,  """do_normalize"""       )       )
           self.assertTrue(hasattr(_a	,  """do_resize"""       )       )
           self.assertTrue(hasattr(_a	,  """size"""       )       )

     def       a__			(						self       )		->							Any:
           _A							:  List[str]    =		self.image_processing_class.from_dict(self.image_processor_dict       )
           self.assertEqual(image_processor.size	,  {"""height""": 18, """width""": 18}       )

           _A							:  str    =		self.image_processing_class.from_dict(self.image_processor_dict	,  size=42       )
           self.assertEqual(image_processor.size	,  {"""height""": 42, """width""": 42}       )

     def       a__			(						self       )		->							List[Any]:
           # Initialize image_processing
           _A							:  Optional[int]    =		self.image_processing_class(**self.image_processor_dict       )
           # create random PIL images
           _A							:  Any    =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=_a       )
           for image in image_inputs:
                 self.assertIsInstance(_a	,  Image.Image       )

           # Test not batched input
           _A							:  Union[str, Any]    =		image_processing(image_inputs[0]	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )

           # Test batched
           _A							:  int    =		image_processing(_a	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )

     def       a__			(						self       )		->							Union[str, Any]:
           # Initialize image_processing
           _A							:  List[Any]    =		self.image_processing_class(**self.image_processor_dict       )
           # create random numpy tensors
           _A							:  Optional[Any]    =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=_a	,  numpify=_a       )
           for image in image_inputs:
                 self.assertIsInstance(_a	,  np.ndarray       )

           # Test not batched input
           _A							:  Dict    =		image_processing(image_inputs[0]	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )

           # Test batched
           _A							:  Any    =		image_processing(_a	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )



     def       a__			(						self       )		->							List[str]:
           # Initialize image_processing
           _A							:  Dict    =		self.image_processing_class(**self.image_processor_dict       )
           # create random PyTorch tensors
           _A							:  str    =		prepare_image_inputs(self.image_processor_tester	,  equal_resolution=_a	,  torchify=_a       )
           for image in image_inputs:
                 self.assertIsInstance(_a	,  torch.Tensor       )

           # Test not batched input
           _A							:  Optional[Any]    =		image_processing(image_inputs[0]	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )

           # Test batched
           _A							:  int    =		image_processing(_a	,  return_tensors="""pt"""       ).pixel_values
           self.assertEqual(
               encoded_images.shape	,  (
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )	,  )

