





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

def 						__magic_name__       (							__UpperCAmelCase      )						->					Any:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       ASTConfig()

   if "10-10" in model_name:
      pass
   elif "speech-commands" in model_name:
      __SCREAMING_SNAKE_CASE      							=       128
   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      							=       527
      __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 						__magic_name__       (							__UpperCAmelCase      )						->					int:

   '''simple docstring'''







   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 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase      )						->					Dict:

   '''simple docstring'''







   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 						__magic_name__       (							__UpperCAmelCase      )						->					List[Any]:

   '''simple docstring'''







   __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 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase							,	__UpperCAmelCase=False      )						->					int:

   '''simple docstring'''







   __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      							=       1024 if """speech-commands""" not in model_name else 128
   __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=16000							,	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__":
 a        =						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."
 )

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












'''simple docstring'''

import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD

def 						__magic_name__       (							__UpperCAmelCase      )						->					Tuple:

   '''simple docstring'''







   def wrapper(*__UpperCAmelCase							,	**__UpperCAmelCase      ):
      __SCREAMING_SNAKE_CASE      							=       timeit.default_timer()
      __SCREAMING_SNAKE_CASE      							=       func(*__UpperCAmelCase							,	**__UpperCAmelCase      )
      __SCREAMING_SNAKE_CASE      							=       timeit.default_timer() - starttime
      return delta

   __SCREAMING_SNAKE_CASE      							=       func.__name__

   return wrapper

def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase=100							,	__UpperCAmelCase=None      )						->					Optional[Any]:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       []
   __SCREAMING_SNAKE_CASE      							=       seq_shapes or {}
   for i in range(__UpperCAmelCase      ):
      __SCREAMING_SNAKE_CASE      							=       {}
      for col_id, (k, v) in enumerate(features.items()      ):
         if isinstance(__UpperCAmelCase							,	_ArrayXD      ):
            __SCREAMING_SNAKE_CASE      							=       np.random.rand(*v.shape      ).astype(v.dtype      )
         elif isinstance(__UpperCAmelCase							,	datasets.Value      ):
            if v.dtype == "string":
               __SCREAMING_SNAKE_CASE      							=       """The small grey turtle was surprisingly fast when challenged."""
            else:
               __SCREAMING_SNAKE_CASE      							=       np.random.randint(10							,	size=1      ).astype(v.dtype      ).item()
         elif isinstance(__UpperCAmelCase							,	datasets.Sequence      ):
            while isinstance(__UpperCAmelCase							,	datasets.Sequence      ):
               __SCREAMING_SNAKE_CASE      							=       v.feature
            __SCREAMING_SNAKE_CASE      							=       seq_shapes[k]
            __SCREAMING_SNAKE_CASE      							=       np.random.rand(*__UpperCAmelCase      ).astype(v.dtype      )
         __SCREAMING_SNAKE_CASE      							=       data

      dummy_data.append((i, example)      )

   return dummy_data







def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase							,	__UpperCAmelCase=100							,	__UpperCAmelCase=None      )						->					str:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       generate_examples(__UpperCAmelCase							,	num_examples=__UpperCAmelCase							,	seq_shapes=__UpperCAmelCase      )

   with ArrowWriter(features=__UpperCAmelCase							,	path=__UpperCAmelCase      ) as writer:
      for key, record in dummy_data:
         __SCREAMING_SNAKE_CASE      							=       features.encode_example(__UpperCAmelCase      )
         writer.write(__UpperCAmelCase      )

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       writer.finalize()

   if not num_final_examples == num_examples:
      raise ValueError(
          f"""Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}."""      )

   __SCREAMING_SNAKE_CASE      							=       datasets.Dataset.from_file(filename=__UpperCAmelCase							,	info=datasets.DatasetInfo(features=__UpperCAmelCase      )      )

   return dataset





