



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_:List[Any]					=						logging.get_logger(__name__)




def      __UpperCamelCase     (  _lowerCAmelCase		)	->  List[str]:



    """simple docstring"""






    A			:	Optional[int]										=       ASTConfig()

    if "10-10" in model_name:
        pass
    elif "speech-commands" in model_name:
        A			:	Tuple										=       128
    elif "12-12" in model_name:
        A			:	Tuple										=       12
        A			:	str										=       12
    elif "14-14" in model_name:
        A			:	Tuple										=       14
        A			:	Tuple										=       14
    elif "16-16" in model_name:
        A			:	Any										=       16
        A			:	Dict										=       16
    else:
        raise ValueError("""Model not supported"""		)

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

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

    return config




def      __UpperCamelCase     (  _lowerCAmelCase		)	->  int:



    """simple docstring"""






    if "module.v" in name:
        A			:	List[str]										=       name.replace("""module.v"""				,					"""audio_spectrogram_transformer"""		)
    if "cls_token" in name:
        A			:	int										=       name.replace("""cls_token"""				,					"""embeddings.cls_token"""		)
    if "dist_token" in name:
        A			:	Any										=       name.replace("""dist_token"""				,					"""embeddings.distillation_token"""		)
    if "pos_embed" in name:
        A			:	Any										=       name.replace("""pos_embed"""				,					"""embeddings.position_embeddings"""		)
    if "patch_embed.proj" in name:
        A			:	Optional[int]										=       name.replace("""patch_embed.proj"""				,					"""embeddings.patch_embeddings.projection"""		)
    # transformer blocks
    if "blocks" in name:
        A			:	List[Any]										=       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			:	str										=       name.replace("""attn"""				,					"""attention.self"""		)
    if "norm1" in name:
        A			:	Optional[Any]										=       name.replace("""norm1"""				,					"""layernorm_before"""		)
    if "norm2" in name:
        A			:	str										=       name.replace("""norm2"""				,					"""layernorm_after"""		)
    if "mlp.fc1" in name:
        A			:	str										=       name.replace("""mlp.fc1"""				,					"""intermediate.dense"""		)
    if "mlp.fc2" in name:
        A			:	Optional[int]										=       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[int]										=       name.replace("""module.mlp_head.0"""				,					"""classifier.layernorm"""		)
    if "module.mlp_head.1" in name:
        A			:	Union[str, Any]										=       name.replace("""module.mlp_head.1"""				,					"""classifier.dense"""		)

    return name




def      __UpperCamelCase     (  _lowerCAmelCase				,					_lowerCAmelCase		)	->  Any:



    """simple docstring"""






    for key in orig_state_dict.copy().keys():
        A			:	Dict										=       orig_state_dict.pop(_lowerCAmelCase		)

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

    return orig_state_dict




def      __UpperCamelCase     (  _lowerCAmelCase		)	->  str:



    """simple docstring"""






    A			:	int										=       [
        """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(_lowerCAmelCase				,					_lowerCAmelCase		)




@torch.no_grad()
def      __UpperCamelCase     (  _lowerCAmelCase				,					_lowerCAmelCase				,					_lowerCAmelCase=False		)	->  Dict:



    """simple docstring"""






    A			:	str										=       get_audio_spectrogram_transformer_config(_lowerCAmelCase		)

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

    # load 🤗 model
    A			:	Tuple										=       ASTForAudioClassification(_lowerCAmelCase		)
    model.eval()

    model.load_state_dict(_lowerCAmelCase		)

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

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

        A					,							A			:	Tuple										=       torchaudio.load(_lowerCAmelCase		)
        A			:	Any										=       waveform.squeeze().numpy()

    A			:	List[str]										=       feature_extractor(_lowerCAmelCase				,					sampling_rate=1_6000				,					return_tensors="""pt"""		)

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

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

    if pytorch_dump_folder_path is not None:
        Path(_lowerCAmelCase		).mkdir(exist_ok=_lowerCAmelCase		)
        print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''		)
        model.save_pretrained(_lowerCAmelCase		)
        print(f'''Saving feature extractor to {pytorch_dump_folder_path}'''		)
        feature_extractor.save_pretrained(_lowerCAmelCase		)

    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_:List[Any]					=						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_:Any					=						parser.parse_args()
     convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_:int					=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_:Any					=						{
    """roberta-base""": """https://huggingface.co/roberta-base/resolve/main/config.json""",
    """roberta-large""": """https://huggingface.co/roberta-large/resolve/main/config.json""",
    """roberta-large-mnli""": """https://huggingface.co/roberta-large-mnli/resolve/main/config.json""",
    """distilroberta-base""": """https://huggingface.co/distilroberta-base/resolve/main/config.json""",
    """roberta-base-openai-detector""": """https://huggingface.co/roberta-base-openai-detector/resolve/main/config.json""",
    """roberta-large-openai-detector""": """https://huggingface.co/roberta-large-openai-detector/resolve/main/config.json""",
}






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    __lowerCamelCase      :			Any =				"roberta"






    def __init__(	self,		lowerCamelCase__=5_0265,		lowerCamelCase__=768,		lowerCamelCase__=12,		lowerCamelCase__=12,		lowerCamelCase__=3072,		lowerCamelCase__="gelu",		lowerCamelCase__=0.1,		lowerCamelCase__=0.1,		lowerCamelCase__=512,		lowerCamelCase__=2,		lowerCamelCase__=0.02,		lowerCamelCase__=1e-12,		lowerCamelCase__=1,		lowerCamelCase__=0,		lowerCamelCase__=2,		lowerCamelCase__="absolute",		lowerCamelCase__=True,		lowerCamelCase__=None,		**lowerCamelCase__,		):
        super().__init__(pad_token_id=lowerCamelCase__,		bos_token_id=lowerCamelCase__,		eos_token_id=lowerCamelCase__,		**lowerCamelCase__		)

        A			:	List[str]										=       vocab_size
        A			:	List[Any]										=       hidden_size
        A			:	List[Any]										=       num_hidden_layers
        A			:	Optional[Any]										=       num_attention_heads
        A			:	Union[str, Any]										=       hidden_act
        A			:	Union[str, Any]										=       intermediate_size
        A			:	Optional[int]										=       hidden_dropout_prob
        A			:	List[str]										=       attention_probs_dropout_prob
        A			:	int										=       max_position_embeddings
        A			:	List[str]										=       type_vocab_size
        A			:	List[str]										=       initializer_range
        A			:	Dict										=       layer_norm_eps
        A			:	Optional[int]										=       position_embedding_type
        A			:	Any										=       use_cache
        A			:	Optional[int]										=       classifier_dropout






class SCREAMING_SNAKE_CASE__					(		SCREAMING_SNAKE_CASE__   ):




    '''simple docstring'''






    @property
    def       _lowerCAmelCase (	self		):
        if self.task == "multiple-choice":
            A			:	Optional[Any]										=       {0: """batch""", 1: """choice""", 2: """sequence"""}
        else:
            A			:	int										=       {0: """batch""", 1: """sequence"""}
        return OrderedDict(
            [
                ("""input_ids""", dynamic_axis),
                ("""attention_mask""", dynamic_axis),
            ]		)


