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

def 		UpperCamelCase	(      lowercase_   )		->	Dict:


  '''simple docstring'''






  lowercase__  :   str    		=      ASTConfig()

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

  lowercase__  :   Union[str, Any]    		=      """huggingface/label-files"""
  if "speech-commands" in model_name:
    lowercase__  :   List[Any]    		=      35
    lowercase__  :   Optional[Any]    		=      """speech-commands-v2-id2label.json"""
  else:
    lowercase__  :   Union[str, Any]    		=      5_27
    lowercase__  :   Tuple    		=      """audioset-id2label.json"""

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

  return config

def 		UpperCamelCase	(      lowercase_   )		->	str:


  '''simple docstring'''






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

  return name

def 		UpperCamelCase	(      lowercase_ , lowercase_   )		->	List[Any]:


  '''simple docstring'''






  for key in orig_state_dict.copy().keys():
    lowercase__  :   Tuple    		=      orig_state_dict.pop(lowercase_   )

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

  return orig_state_dict

def 		UpperCamelCase	(      lowercase_   )		->	Tuple:


  '''simple docstring'''






  lowercase__  :   Dict    		=      [
      """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 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_=False   )		->	Tuple:


  '''simple docstring'''






  lowercase__  :   str    		=      get_audio_spectrogram_transformer_config(lowercase_   )

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

  # load 🤗 model
  lowercase__  :   int    		=      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
  lowercase__  :   str    		=      -4.267_7393 if """speech-commands""" not in model_name else -6.84_5978
  lowercase__  :   Tuple    		=      4.568_9974 if """speech-commands""" not in model_name else 5.565_4526
  lowercase__  :   Dict    		=      10_24 if """speech-commands""" not in model_name else 1_28
  lowercase__  :   List[str]    		=      ASTFeatureExtractor(mean=lowercase_ , std=lowercase_ , max_length=lowercase_   )

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

    lowercase__		,    lowercase__  :   int    		=      torchaudio.load(lowercase_   )
    lowercase__  :   int    		=      waveform.squeeze().numpy()

  lowercase__  :   str    		=      feature_extractor(lowercase_ , sampling_rate=1_60_00 , return_tensors="""pt"""   )

  # forward pass
  lowercase__  :   Optional[Any]    		=      model(**lowercase_   )
  lowercase__  :   List[str]    		=      outputs.logits

  if model_name == "ast-finetuned-audioset-10-10-0.4593":
    lowercase__  :   List[str]    		=      torch.tensor([-0.8760, -7.0042, -8.6602]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.450":
    lowercase__  :   List[Any]    		=      torch.tensor([-1.1986, -7.0903, -8.2718]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.448":
    lowercase__  :   List[Any]    		=      torch.tensor([-2.6128, -8.0080, -9.4344]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
    lowercase__  :   Tuple    		=      torch.tensor([-1.5080, -7.4534, -8.8917]   )
  elif model_name == "ast-finetuned-audioset-12-12-0.447":
    lowercase__  :   Union[str, Any]    		=      torch.tensor([-0.5050, -6.5833, -8.0843]   )
  elif model_name == "ast-finetuned-audioset-14-14-0.443":
    lowercase__  :   List[Any]    		=      torch.tensor([-0.3826, -7.0336, -8.2413]   )
  elif model_name == "ast-finetuned-audioset-16-16-0.442":
    lowercase__  :   int    		=      torch.tensor([-1.2113, -6.9101, -8.3470]   )
  elif model_name == "ast-finetuned-speech-commands-v2":
    lowercase__  :   int    		=      torch.tensor([6.1589, -8.0566, -8.7984]   )
  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__":
     lowerCamelCase__    :		Dict          =   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__    :		Any          =   parser.parse_args()
     convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








from abc import ABC, abstractmethod
from argparse import ArgumentParser



class       _snake_case   (	UpperCAmelCase_							):


 @staticmethod
 @abstractmethod
 def      lowercase__     (					SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   raise NotImplementedError()



 @abstractmethod
 def      lowercase__     (					self):


   '''simple docstring'''




   raise NotImplementedError()






