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







def 						__a   (	_SCREAMING_SNAKE_CASE   )	->Union[str, Any]:
  a__:     List[str]										=	ASTConfig()

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

  a__:     int										=	'huggingface/label-files'
  if "speech-commands" in model_name:
    a__:     Tuple										=	35
    a__:     str										=	'speech-commands-v2-id2label.json'
  else:
    a__:     Union[str, Any]										=	527
    a__:     int										=	'audioset-id2label.json'

  a__:     Optional[Any]										=	json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE		,     repo_type='dataset'   )		,     'r'   )   )
  a__:     int										=	{int(_SCREAMING_SNAKE_CASE   ): v for k, v in idalabel.items()}
  a__:     List[str]										=	idalabel
  a__:     List[Any]										=	{v: k for k, v in idalabel.items()}

  return config







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

  return name







def 						__a   (	_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE   )	->Any:
  for key in orig_state_dict.copy().keys():
    a__:     Optional[Any]										=	orig_state_dict.pop(_SCREAMING_SNAKE_CASE   )

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

  return orig_state_dict







def 						__a   (	_SCREAMING_SNAKE_CASE   )	->Optional[int]:
  a__:     Optional[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(_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE   )







@torch.no_grad()
def 						__a   (	_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE=False   )	->int:
  a__:     Tuple										=	get_audio_spectrogram_transformer_config(_SCREAMING_SNAKE_CASE   )

  a__:     List[Any]										=	{
      '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__:     int										=	model_name_to_url[model_name]
  a__:     str										=	torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE		,     map_location='cpu'   )
  # remove some keys
  remove_keys(_SCREAMING_SNAKE_CASE   )
  # rename some keys
  a__:     Optional[int]										=	convert_state_dict(_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE   )

  # load 🤗 model
  a__:     Optional[int]										=	ASTForAudioClassification(_SCREAMING_SNAKE_CASE   )
  model.eval()

  model.load_state_dict(_SCREAMING_SNAKE_CASE   )

  # verify outputs on dummy input
  # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
  a__:     Any										=	-4.2_677_393 if 'speech-commands' not in model_name else -6.845_978
  a__:     Dict										=	4.5_689_974 if 'speech-commands' not in model_name else 5.5_654_526
  a__:     List[Any]										=	1024 if 'speech-commands' not in model_name else 128
  a__:     str										=	ASTFeatureExtractor(mean=_SCREAMING_SNAKE_CASE		,     std=_SCREAMING_SNAKE_CASE		,     max_length=_SCREAMING_SNAKE_CASE   )

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

    a__							,    a__:     List[str]										=	torchaudio.load(_SCREAMING_SNAKE_CASE   )
    a__:     Union[str, Any]										=	waveform.squeeze().numpy()

  a__:     Tuple										=	feature_extractor(_SCREAMING_SNAKE_CASE		,     sampling_rate=16000		,     return_tensors='pt'   )

  # forward pass
  a__:     Optional[int]										=	model(**_SCREAMING_SNAKE_CASE   )
  a__:     Tuple										=	outputs.logits

  if model_name == "ast-finetuned-audioset-10-10-0.4593":
    a__:     str										=	torch.tensor([-0.8_760, -7.0_042, -8.6_602]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.450":
    a__:     str										=	torch.tensor([-1.1_986, -7.0_903, -8.2_718]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.448":
    a__:     Tuple										=	torch.tensor([-2.6_128, -8.0_080, -9.4_344]   )
  elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
    a__:     Tuple										=	torch.tensor([-1.5_080, -7.4_534, -8.8_917]   )
  elif model_name == "ast-finetuned-audioset-12-12-0.447":
    a__:     Dict										=	torch.tensor([-0.5_050, -6.5_833, -8.0_843]   )
  elif model_name == "ast-finetuned-audioset-14-14-0.443":
    a__:     Optional[Any]										=	torch.tensor([-0.3_826, -7.0_336, -8.2_413]   )
  elif model_name == "ast-finetuned-audioset-16-16-0.442":
    a__:     Optional[Any]										=	torch.tensor([-1.2_113, -6.9_101, -8.3_470]   )
  elif model_name == "ast-finetuned-speech-commands-v2":
    a__:     Tuple										=	torch.tensor([6.1_589, -8.0_566, -8.7_984]   )
  else:
    raise ValueError('Unknown model name'   )
  if not torch.allclose(logits[0, :3]		,     _SCREAMING_SNAKE_CASE		,     atol=1e-4   ):
    raise ValueError('Logits don\'t match'   )
  print('Looks ok!'   )

  if pytorch_dump_folder_path is not None:
    Path(_SCREAMING_SNAKE_CASE   ).mkdir(exist_ok=_SCREAMING_SNAKE_CASE   )
    print(F'Saving model {model_name} to {pytorch_dump_folder_path}'   )
    model.save_pretrained(_SCREAMING_SNAKE_CASE   )
    print(F'Saving feature extractor to {pytorch_dump_folder_path}'   )
    feature_extractor.save_pretrained(_SCREAMING_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__":
 lowercase__					       =	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.'
 )

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



"""simple docstring"""
import itertools
import random
import unittest

import numpy as np

from transformers import WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST, WavaVecaConfig, WavaVecaFeatureExtractor
from transformers.testing_utils import require_torch, slow

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


lowercase__					       =	random.Random()







def 						__a   (	_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE=1.0		,     _SCREAMING_SNAKE_CASE=None		,     _SCREAMING_SNAKE_CASE=None   )	->Optional[int]:
  if rng is None:
    a__:     Any										=	global_rng

  a__:     int										=	[]
  for batch_idx in range(shape[0]   ):
    values.append([]   )
    for _ in range(shape[1]   ):
      values[-1].append(rng.random() * scale   )

  return values




class    __snake_case       (     unittest.TestCase				):



      def __init__(						self    ,			lowercase    ,			lowercase=7    ,			lowercase=4_00    ,			lowercase=20_00    ,			lowercase=1    ,			lowercase=0.0    ,			lowercase=1_60_00    ,			lowercase=True    ,			lowercase=True    ,			)	->			Union[str, Any]:





        '''simple docstring'''





        a__:     Tuple										=	parent
        a__:     Optional[int]										=	batch_size
        a__:     Optional[Any]										=	min_seq_length
        a__:     Optional[int]										=	max_seq_length
        a__:     Tuple										=	(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
        a__:     Dict										=	feature_size
        a__:     Any										=	padding_value
        a__:     Optional[Any]										=	sampling_rate
        a__:     Optional[Any]										=	return_attention_mask
        a__:     str										=	do_normalize



      def   lowerCamelCase_   (						self)	->			Dict:





        '''simple docstring'''





        return {
            "feature_size": self.feature_size,
            "padding_value": self.padding_value,
            "sampling_rate": self.sampling_rate,
            "return_attention_mask": self.return_attention_mask,
            "do_normalize": self.do_normalize,
        }



      def   lowerCamelCase_   (						self    ,			lowercase=False    ,			lowercase=False)	->			Tuple:





        '''simple docstring'''





        def _flatten(lowercase):
          return list(itertools.chain(*lowercase))

        if equal_length:
          a__:     Dict										=	floats_list((self.batch_size, self.max_seq_length))
        else:
          # make sure that inputs increase in size
          a__:     List[Any]										=	[
              _flatten(floats_list((x, self.feature_size)))
              for x in range(self.min_seq_length    ,			self.max_seq_length    ,			self.seq_length_diff)
          ]

        if numpify:
          a__:     str										=	[np.asarray(lowercase) for x in speech_inputs]

        return speech_inputs







class    __snake_case       (     __lowerCAmelCase       ,      unittest.TestCase				):
      a__        =		WavaVecaFeatureExtractor



      def   lowerCamelCase_   (						self)	->			Tuple:





        '''simple docstring'''





        a__:     Optional[int]										=	WavaVecaFeatureExtractionTester(self)



      def   lowerCamelCase_   (						self    ,			lowercase)	->			List[Any]:





        '''simple docstring'''





        self.assertTrue(np.all(np.mean(lowercase    ,			axis=0) < 1e-3))
        self.assertTrue(np.all(np.abs(np.var(lowercase    ,			axis=0) - 1) < 1e-3))



      def   lowerCamelCase_   (						self)	->			List[str]:





        '''simple docstring'''





        a__:     List[Any]										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        # create three inputs of length 800, 1000, and 1200
        a__:     Optional[Any]										=	[floats_list((1, x))[0] for x in range(8_00    ,			14_00    ,			2_00)]
        a__:     List[str]										=	[np.asarray(lowercase) for speech_input in speech_inputs]

        # Test not batched input
        a__:     Optional[Any]										=	feat_extract(speech_inputs[0]    ,			return_tensors='np').input_values
        a__:     Dict										=	feat_extract(np_speech_inputs[0]    ,			return_tensors='np').input_values
        self.assertTrue(np.allclose(lowercase    ,			lowercase    ,			atol=1e-3))

        # Test batched
        a__:     Dict										=	feat_extract(lowercase    ,			return_tensors='np').input_values
        a__:     int										=	feat_extract(lowercase    ,			return_tensors='np').input_values
        for enc_seq_a, enc_seq_a in zip(lowercase    ,			lowercase):
          self.assertTrue(np.allclose(lowercase    ,			lowercase    ,			atol=1e-3))

        # Test 2-D numpy arrays are batched.
        a__:     int										=	[floats_list((1, x))[0] for x in (8_00, 8_00, 8_00)]
        a__:     Union[str, Any]										=	np.asarray(lowercase)
        a__:     int										=	feat_extract(lowercase    ,			return_tensors='np').input_values
        a__:     Any										=	feat_extract(lowercase    ,			return_tensors='np').input_values
        for enc_seq_a, enc_seq_a in zip(lowercase    ,			lowercase):
          self.assertTrue(np.allclose(lowercase    ,			lowercase    ,			atol=1e-3))



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





        '''simple docstring'''





        a__:     Tuple										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        a__:     List[Any]										=	[floats_list((1, x))[0] for x in range(8_00    ,			14_00    ,			2_00)]

        a__:     Optional[int]										=	['longest', 'max_length', 'do_not_pad']
        a__:     List[Any]										=	[None, 16_00, None]
        for max_length, padding in zip(lowercase    ,			lowercase):
          a__:     Dict										=	feat_extract(lowercase    ,			padding=lowercase    ,			max_length=lowercase    ,			return_tensors='np')
          a__:     Union[str, Any]										=	processed.input_values

          self._check_zero_mean_unit_variance(input_values[0][:8_00])
          self.assertTrue(input_values[0][8_00:].sum() < 1e-6)
          self._check_zero_mean_unit_variance(input_values[1][:10_00])
          self.assertTrue(input_values[0][10_00:].sum() < 1e-6)
          self._check_zero_mean_unit_variance(input_values[2][:12_00])



      def   lowerCamelCase_   (						self)	->			Optional[int]:





        '''simple docstring'''





        a__:     str										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        a__:     Optional[int]										=	range(8_00    ,			14_00    ,			2_00)
        a__:     List[str]										=	[floats_list((1, x))[0] for x in lengths]

        a__:     Tuple										=	['longest', 'max_length', 'do_not_pad']
        a__:     Dict										=	[None, 16_00, None]

        for max_length, padding in zip(lowercase    ,			lowercase):
          a__:     int										=	feat_extract(lowercase    ,			max_length=lowercase    ,			padding=lowercase)
          a__:     Any										=	processed.input_values

          self._check_zero_mean_unit_variance(input_values[0][:8_00])
          self._check_zero_mean_unit_variance(input_values[1][:10_00])
          self._check_zero_mean_unit_variance(input_values[2][:12_00])



      def   lowerCamelCase_   (						self)	->			List[str]:





        '''simple docstring'''





        a__:     Any										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        a__:     Any										=	[floats_list((1, x))[0] for x in range(8_00    ,			14_00    ,			2_00)]
        a__:     Dict										=	feat_extract(
            lowercase    ,			truncation=lowercase    ,			max_length=10_00    ,			padding='max_length'    ,			return_tensors='np')
        a__:     int										=	processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :8_00])
        self._check_zero_mean_unit_variance(input_values[1])
        self._check_zero_mean_unit_variance(input_values[2])



      def   lowerCamelCase_   (						self)	->			Dict:





        '''simple docstring'''





        a__:     Tuple										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        a__:     int										=	[floats_list((1, x))[0] for x in range(8_00    ,			14_00    ,			2_00)]
        a__:     str										=	feat_extract(
            lowercase    ,			truncation=lowercase    ,			max_length=10_00    ,			padding='longest'    ,			return_tensors='np')
        a__:     Any										=	processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :8_00])
        self._check_zero_mean_unit_variance(input_values[1, :10_00])
        self._check_zero_mean_unit_variance(input_values[2])

        # make sure that if max_length < longest -> then pad to max_length
        self.assertTrue(input_values.shape == (3, 10_00))

        a__:     Dict										=	[floats_list((1, x))[0] for x in range(8_00    ,			14_00    ,			2_00)]
        a__:     Tuple										=	feat_extract(
            lowercase    ,			truncation=lowercase    ,			max_length=20_00    ,			padding='longest'    ,			return_tensors='np')
        a__:     str										=	processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :8_00])
        self._check_zero_mean_unit_variance(input_values[1, :10_00])
        self._check_zero_mean_unit_variance(input_values[2])

        # make sure that if max_length > longest -> then pad to longest
        self.assertTrue(input_values.shape == (3, 12_00))



      @require_torch
      def   lowerCamelCase_   (						self)	->			Tuple:





        '''simple docstring'''





        import torch

        a__:     Optional[int]										=	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict())
        a__:     Tuple										=	np.random.rand(1_00).astype(np.floataa)
        a__:     Tuple										=	np_speech_inputs.tolist()

        for inputs in [py_speech_inputs, np_speech_inputs]:
          a__:     Any										=	feature_extractor.pad([{'input_values': inputs}]    ,			return_tensors='np')
          self.assertTrue(np_processed.input_values.dtype == np.floataa)
          a__:     Optional[Any]										=	feature_extractor.pad([{'input_values': inputs}]    ,			return_tensors='pt')
          self.assertTrue(pt_processed.input_values.dtype == torch.floataa)



      @slow
      @require_torch
      def   lowerCamelCase_   (						self)	->			Optional[Any]:





        '''simple docstring'''





        for model_id in WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST:
          a__:     str										=	WavaVecaConfig.from_pretrained(lowercase)
          a__:     str										=	WavaVecaFeatureExtractor.from_pretrained(lowercase)

          # only "layer" feature extraction norm should make use of
          # attention_mask
          self.assertEqual(feat_extract.return_attention_mask    ,			config.feat_extract_norm == 'layer')


