




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







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[str]:
 A__	=    ASTConfig()

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

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

 A__	=    json.load(open(hf_hub_download(lowercase_      ,			lowercase_      ,			repo_type="dataset"						)      ,			"r"						)						)
 A__	=    {int(lowercase_						): v for k, v in idalabel.items()}
 A__	=    idalabel
 A__	=    {v: k for k, v in idalabel.items()}

 return config







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

 return name







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Optional[int]:
 for key in orig_state_dict.copy().keys():
  A__	=    orig_state_dict.pop(lowercase_						)

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

 return orig_state_dict







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[Any]:
 A__	=    [
     "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   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_=False						)						->		str:
 A__	=    get_audio_spectrogram_transformer_config(lowercase_						)

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

 # load 🤗 model
 A__	=    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
 A__	=    -4.2_6_7_7_3_9_3 if "speech-commands" not in model_name else -6.8_4_5_9_7_8
 A__	=    4.5_6_8_9_9_7_4 if "speech-commands" not in model_name else 5.5_6_5_4_5_2_6
 A__	=    10_24 if "speech-commands" not in model_name else 1_28
 A__	=    ASTFeatureExtractor(mean=lowercase_      ,			std=lowercase_      ,			max_length=lowercase_						)

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

  A__,				A__	=    torchaudio.load(lowercase_						)
  A__	=    waveform.squeeze().numpy()

 A__	=    feature_extractor(lowercase_      ,			sampling_rate=1_60_00      ,			return_tensors="pt"						)

 # forward pass
 A__	=    model(**lowercase_						)
 A__	=    outputs.logits

 if model_name == "ast-finetuned-audioset-10-10-0.4593":
  A__	=    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":
  A__	=    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":
  A__	=    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":
  A__	=    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":
  A__	=    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":
  A__	=    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":
  A__	=    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":
  A__	=    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]      ,			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__":
      SCREAMING_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."
      )

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







"""simple docstring"""



import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 return x + 2






class      UpperCAmelCase_  (						unittest.TestCase							):




    def 							__magic_name__							(   self						:	Any				)				->							Any:




     '''simple docstring'''

     A__	=    "x = 3"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3}				)

     A__	=    "x = y"
     A__	=    {"y": 5}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 5, "y": 5}				)




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Any:




     '''simple docstring'''

     A__	=    "y = add_two(x)"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)

     # Won't work without the tool
     with CaptureStdout() as out:
      A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result is None
     assert "tried to execute add_two" in out.out




    def 							__magic_name__							(   self						:	Dict				)				->							List[str]:




     '''simple docstring'''

     A__	=    "x = 3"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3}				)




    def 							__magic_name__							(   self						:	Dict				)				->							Optional[Any]:




     '''simple docstring'''

     A__	=    "test_dict = {'x': x, 'y': add_two(x)}"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_dict": {"x": 3, "y": 5}}				)




    def 							__magic_name__							(   self						:	int				)				->							str:




     '''simple docstring'''

     A__	=    "x = 3\ny = 5"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 5}				)




    def 							__magic_name__							(   self						:	Any				)				->							List[Any]:




     '''simple docstring'''

     A__	=    "text = f'This is x: {x}.'"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == "This is x: 3."
     self.assertDictEqual(snake_case_     ,			{"x": 3, "text": "This is x: 3."}				)




    def 							__magic_name__							(   self						:	str				)				->							Optional[int]:




     '''simple docstring'''

     A__	=    "if x <= 3:\n    y = 2\nelse:\n    y = 5"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 2
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 2}				)

     A__	=    {"x": 8}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     # evaluate returns the value of the last assignment.
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 8, "y": 5}				)




    def 							__magic_name__							(   self						:	List[str]				)				->							Tuple:




     '''simple docstring'''

     A__	=    "test_list = [x, add_two(x)]"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     self.assertListEqual(snake_case_     ,			[3, 5]				)
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_list": [3, 5]}				)




    def 							__magic_name__							(   self						:	Tuple				)				->							int:




     '''simple docstring'''

     A__	=    "y = x"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{}     ,			state=snake_case_				)
     assert result == 3
     self.assertDictEqual(snake_case_     ,			{"x": 3, "y": 3}				)




    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     A__	=    "test_list = [x, add_two(x)]\ntest_list[1]"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_list": [3, 5]}				)

     A__	=    "test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
     A__	=    {"x": 3}
     A__	=    evaluate(snake_case_     ,			{"add_two": add_two}     ,			state=snake_case_				)
     assert result == 5
     self.assertDictEqual(snake_case_     ,			{"x": 3, "test_dict": {"x": 3, "y": 5}}				)




    def 							__magic_name__							(   self						:	Optional[int]				)				->							Tuple:




     '''simple docstring'''

     A__	=    "x = 0\nfor i in range(3):\n    x = i"
     A__	=    {}
     A__	=    evaluate(snake_case_     ,			{"range": range}     ,			state=snake_case_				)
     assert result == 2
     self.assertDictEqual(snake_case_     ,			{"x": 2, "i": 2}				)





