


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






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							str   			=       ASTConfig()

       if "10-10" in model_name:
              pass
       elif "speech-commands" in model_name:
              _lowerCAmelCase    :							List[Any]   			=       1_2_8
       elif "12-12" in model_name:
              _lowerCAmelCase    :							Dict   			=       1_2
              _lowerCAmelCase    :							str   			=       1_2
       elif "14-14" in model_name:
              _lowerCAmelCase    :							Tuple   			=       1_4
              _lowerCAmelCase    :							Tuple   			=       1_4
       elif "16-16" in model_name:
              _lowerCAmelCase    :							Optional[Any]   			=       1_6
              _lowerCAmelCase    :							Union[str, Any]   			=       1_6
       else:
              raise ValueError('Model not supported'					)

       _lowerCAmelCase    :							Dict   			=       'huggingface/label-files'
       if "speech-commands" in model_name:
              _lowerCAmelCase    :							Optional[int]   			=       3_5
              _lowerCAmelCase    :							str   			=       'speech-commands-v2-id2label.json'
       else:
              _lowerCAmelCase    :							Dict   			=       5_2_7
              _lowerCAmelCase    :							List[str]   			=       'audioset-id2label.json'

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

       return config






def 			lowercase				(_A					):







       """simple docstring"""







       if "module.v" in name:
              _lowerCAmelCase    :							Dict   			=       name.replace('module.v'   ,						'audio_spectrogram_transformer'					)
       if "cls_token" in name:
              _lowerCAmelCase    :							Dict   			=       name.replace('cls_token'   ,						'embeddings.cls_token'					)
       if "dist_token" in name:
              _lowerCAmelCase    :							Union[str, Any]   			=       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    :							Any   			=       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    :							Optional[Any]   			=       name.replace('attn.proj'   ,						'attention.output.dense'					)
       if "attn" in name:
              _lowerCAmelCase    :							int   			=       name.replace('attn'   ,						'attention.self'					)
       if "norm1" in name:
              _lowerCAmelCase    :							str   			=       name.replace('norm1'   ,						'layernorm_before'					)
       if "norm2" in name:
              _lowerCAmelCase    :							Optional[Any]   			=       name.replace('norm2'   ,						'layernorm_after'					)
       if "mlp.fc1" in name:
              _lowerCAmelCase    :							List[Any]   			=       name.replace('mlp.fc1'   ,						'intermediate.dense'					)
       if "mlp.fc2" in name:
              _lowerCAmelCase    :							Any   			=       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    :							Any   			=       name.replace('module.mlp_head.0'   ,						'classifier.layernorm'					)
       if "module.mlp_head.1" in name:
              _lowerCAmelCase    :							Union[str, Any]   			=       name.replace('module.mlp_head.1'   ,						'classifier.dense'					)

       return name






def 			lowercase				(_A   ,						_A					):







       """simple docstring"""







       for key in orig_state_dict.copy().keys():
              _lowerCAmelCase    :							Any   			=       orig_state_dict.pop(_A					)

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

       return orig_state_dict






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							str   			=       [
           '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 			lowercase				(_A   ,						_A   ,						_A=False					):







       """simple docstring"""







       _lowerCAmelCase    :							Any   			=       get_audio_spectrogram_transformer_config(_A					)

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

       # load 🤗 model
       _lowerCAmelCase    :							Optional[int]   			=       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    :							List[Any]   			=       -4.2_677_393 if 'speech-commands' not in model_name else -6.845_978
       _lowerCAmelCase    :							List[str]   			=       4.5_689_974 if 'speech-commands' not in model_name else 5.5_654_526
       _lowerCAmelCase    :							Any   			=       1_0_2_4 if 'speech-commands' not in model_name else 1_2_8
       _lowerCAmelCase    :							Optional[Any]   			=       ASTFeatureExtractor(mean=_A   ,						std=_A   ,						max_length=_A					)

       if "speech-commands" in model_name:
              _lowerCAmelCase    :							List[str]   			=       load_dataset('speech_commands'   ,						'v0.02'   ,						split='validation'					)
              _lowerCAmelCase    :							int   			=       dataset[0]['audio']['array']
       else:
              _lowerCAmelCase    :							Optional[int]   			=       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    :							Any   			=       waveform.squeeze().numpy()

       _lowerCAmelCase    :							Tuple   			=       feature_extractor(_A   ,						sampling_rate=1_6_0_0_0   ,						return_tensors='pt'					)

       # forward pass
       _lowerCAmelCase    :							List[Any]   			=       model(**_A					)
       _lowerCAmelCase    :							Tuple   			=       outputs.logits

       if model_name == "ast-finetuned-audioset-10-10-0.4593":
              _lowerCAmelCase    :							int   			=       torch.tensor([-0.8_760, -7.0_042, -8.6_602]					)
       elif model_name == "ast-finetuned-audioset-10-10-0.450":
              _lowerCAmelCase    :							Optional[Any]   			=       torch.tensor([-1.1_986, -7.0_903, -8.2_718]					)
       elif model_name == "ast-finetuned-audioset-10-10-0.448":
              _lowerCAmelCase    :							Optional[Any]   			=       torch.tensor([-2.6_128, -8.0_080, -9.4_344]					)
       elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
              _lowerCAmelCase    :							Dict   			=       torch.tensor([-1.5_080, -7.4_534, -8.8_917]					)
       elif model_name == "ast-finetuned-audioset-12-12-0.447":
              _lowerCAmelCase    :							int   			=       torch.tensor([-0.5_050, -6.5_833, -8.0_843]					)
       elif model_name == "ast-finetuned-audioset-14-14-0.443":
              _lowerCAmelCase    :							int   			=       torch.tensor([-0.3_826, -7.0_336, -8.2_413]					)
       elif model_name == "ast-finetuned-audioset-16-16-0.442":
              _lowerCAmelCase    :							int   			=       torch.tensor([-1.2_113, -6.9_101, -8.3_470]					)
       elif model_name == "ast-finetuned-speech-commands-v2":
              _lowerCAmelCase    :							List[str]   			=       torch.tensor([6.1_589, -8.0_566, -8.7_984]					)
       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__":
       lowerCAmelCase    :     Optional[int]					=						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    :     Union[str, Any]					=						parser.parse_args()
       convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





'''simple docstring'''







from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase    :     Optional[int]					=						{
    """configuration_mctct""": ["""MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MCTCTConfig"""],
    """feature_extraction_mctct""": ["""MCTCTFeatureExtractor"""],
    """processing_mctct""": ["""MCTCTProcessor"""],
}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase    :     List[str]					=						[
           """MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """MCTCTForCTC""",
           """MCTCTModel""",
           """MCTCTPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
       from .feature_extraction_mctct import MCTCTFeatureExtractor
       from .processing_mctct import MCTCTProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
       import sys

       lowerCAmelCase    :     Tuple					=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

