





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 					lowerCAmelCase_			(	__a		)					->		List[Any]:







  """simple docstring"""





  lowerCamelCase__:				Any        =ASTConfig()

  if "10-10" in model_name:
    pass
  elif "speech-commands" in model_name:
    lowerCamelCase__:				Optional[Any]        =128
  elif "12-12" in model_name:
    lowerCamelCase__:				Union[str, Any]        =12
    lowerCamelCase__:				Any        =12
  elif "14-14" in model_name:
    lowerCamelCase__:				Optional[Any]        =14
    lowerCamelCase__:				int        =14
  elif "16-16" in model_name:
    lowerCamelCase__:				str        =16
    lowerCamelCase__:				Dict        =16
  else:
    raise ValueError("Model not supported"		)

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

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

  return config





def 					lowerCAmelCase_			(	__a		)					->		str:







  """simple docstring"""





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

  return name





def 					lowerCAmelCase_			(	__a     ,   __a		)					->		str:







  """simple docstring"""





  for key in orig_state_dict.copy().keys():
    lowerCamelCase__:				Any        =orig_state_dict.pop(__a		)

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

  return orig_state_dict





def 					lowerCAmelCase_			(	__a		)					->		int:







  """simple docstring"""





  lowerCamelCase__:				Tuple        =[
      "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 					lowerCAmelCase_			(	__a     ,   __a     ,   __a=False		)					->		Union[str, Any]:







  """simple docstring"""





  lowerCamelCase__:				Optional[int]        =get_audio_spectrogram_transformer_config(__a		)

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

  # load 🤗 model
  lowerCamelCase__:				Union[str, Any]        =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
  lowerCamelCase__:				Any        =-4.2_6_7_7_3_9_3 if "speech-commands" not in model_name else -6.8_4_5_9_7_8
  lowerCamelCase__:				Tuple        =4.5_6_8_9_9_7_4 if "speech-commands" not in model_name else 5.5_6_5_4_5_2_6
  lowerCamelCase__:				Dict        =1024 if "speech-commands" not in model_name else 128
  lowerCamelCase__:				int        =ASTFeatureExtractor(mean=__a     ,   std=__a     ,   max_length=__a		)

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

    lowerCamelCase__  , lowerCamelCase__:				Optional[Any]        =torchaudio.load(__a		)
    lowerCamelCase__:				Optional[int]        =waveform.squeeze().numpy()

  lowerCamelCase__:				Optional[Any]        =feature_extractor(__a     ,   sampling_rate=16000     ,   return_tensors="pt"		)

  # forward pass
  lowerCamelCase__:				List[Any]        =model(**__a		)
  lowerCamelCase__:				Any        =outputs.logits

  if model_name == "ast-finetuned-audioset-10-10-0.4593":
    lowerCamelCase__:				Union[str, Any]        =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":
    lowerCamelCase__:				List[str]        =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":
    lowerCamelCase__:				Tuple        =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":
    lowerCamelCase__:				List[str]        =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":
    lowerCamelCase__:				Union[str, Any]        =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":
    lowerCamelCase__:				Tuple        =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":
    lowerCamelCase__:				List[Any]        =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":
    lowerCamelCase__:				Dict        =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]     ,   __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__":
     __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)








from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand





def 					lowerCAmelCase_			(	__a		)					->		int:







  """simple docstring"""





  return ConvertCommand(
      args.model_type     ,   args.tf_checkpoint     ,   args.pytorch_dump_output     ,   args.config     ,   args.finetuning_task_name		)


__A     =    "\ntransformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires\nTensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.\n"






class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''

 @staticmethod
 def 							SCREAMING_SNAKE_CASE_ (UpperCAmelCase_    :       ArgumentParser)	->str:



   '''simple docstring'''



   lowerCamelCase__:				Dict        =parser.add_parser(
       "convert"	,			help="CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints."	,			)
   train_parser.add_argument("--model_type"	,			type=UpperCAmelCase_	,			required=UpperCAmelCase_	,			help="Model's type.")
   train_parser.add_argument(
       "--tf_checkpoint"	,			type=UpperCAmelCase_	,			required=UpperCAmelCase_	,			help="TensorFlow checkpoint path or folder.")
   train_parser.add_argument(
       "--pytorch_dump_output"	,			type=UpperCAmelCase_	,			required=UpperCAmelCase_	,			help="Path to the PyTorch saved model output.")
   train_parser.add_argument("--config"	,			type=UpperCAmelCase_	,			default=""	,			help="Configuration file path or folder.")
   train_parser.add_argument(
       "--finetuning_task_name"	,			type=UpperCAmelCase_	,			default=UpperCAmelCase_	,			help="Optional fine-tuning task name if the TF model was a finetuned model."	,			)
   train_parser.set_defaults(func=UpperCAmelCase_)

 def __init__(self    :       List[Any]	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       str	,			UpperCAmelCase_    :       str	,			*UpperCAmelCase_    :       Optional[int]	,			)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Dict        =logging.get_logger("transformers-cli/converting")

   self._logger.info(F"""Loading model {model_type}""")
   lowerCamelCase__:				Any        =model_type
   lowerCamelCase__:				Optional[int]        =tf_checkpoint
   lowerCamelCase__:				Any        =pytorch_dump_output
   lowerCamelCase__:				Union[str, Any]        =config
   lowerCamelCase__:				str        =finetuning_task_name







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



   '''simple docstring'''



   if self._model_type == "albert":
     try:
       from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
           convert_tf_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_tf_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "bert":
     try:
       from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
           convert_tf_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_tf_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "funnel":
     try:
       from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
           convert_tf_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_tf_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "t5":
     try:
       from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_tf_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "gpt":
     from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
         convert_openai_checkpoint_to_pytorch,
     )

     convert_openai_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "transfo_xl":
     try:
       from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
           convert_transfo_xl_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     if "ckpt" in self._tf_checkpoint.lower():
       lowerCamelCase__:				Tuple        =self._tf_checkpoint
       lowerCamelCase__:				List[str]        =""
     else:
       lowerCamelCase__:				Any        =self._tf_checkpoint
       lowerCamelCase__:				Dict        =""
     convert_transfo_xl_checkpoint_to_pytorch(
         UpperCAmelCase_	,			self._config	,			self._pytorch_dump_output	,			UpperCAmelCase_)
   elif self._model_type == "gpt2":
     try:
       from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
           convert_gpta_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   elif self._model_type == "xlnet":
     try:
       from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
           convert_xlnet_checkpoint_to_pytorch,
       )
     except ImportError:
       raise ImportError(UpperCAmelCase_)

     convert_xlnet_checkpoint_to_pytorch(
         self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output	,			self._finetuning_task_name)
   elif self._model_type == "xlm":
     from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
         convert_xlm_checkpoint_to_pytorch,
     )

     convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint	,			self._pytorch_dump_output)
   elif self._model_type == "lxmert":
     from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
         convert_lxmert_checkpoint_to_pytorch,
     )

     convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint	,			self._pytorch_dump_output)
   elif self._model_type == "rembert":
     from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
         convert_rembert_tf_checkpoint_to_pytorch,
     )

     convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint	,			self._config	,			self._pytorch_dump_output)
   else:
     raise ValueError(
         "--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]")

