

"""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()
__UpperCAmelCase =				logging.get_logger(__name__)
def   lowerCAmelCase						(					__UpperCamelCase      ):
	'''simple docstring'''





	UpperCAmelCase__					:  int					   =					ASTConfig()

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

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

	UpperCAmelCase__					:  Any					   =					json.load(open(hf_hub_download(_snake_case							,    _snake_case							,    repo_type="""dataset"""      )							,    """r"""      )      )
	UpperCAmelCase__					:  Any					   =					{int(_snake_case      ): v for k, v in idalabel.items()}
	UpperCAmelCase__					:  Tuple					   =					idalabel
	UpperCAmelCase__					:  List[Any]					   =					{v: k for k, v in idalabel.items()}

	return config
def   lowerCAmelCase						(					__UpperCamelCase      ):
	'''simple docstring'''





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

	return name
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase      ):
	'''simple docstring'''





	for key in orig_state_dict.copy().keys():
		UpperCAmelCase__					:  Dict					   =					orig_state_dict.pop(_snake_case      )

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

	return orig_state_dict
def   lowerCAmelCase						(					__UpperCamelCase      ):
	'''simple docstring'''





	UpperCAmelCase__					:  List[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(_snake_case							,    _snake_case      )
@torch.no_grad()
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase=False      ):
	'''simple docstring'''





	UpperCAmelCase__					:  Union[str, Any]					   =					get_audio_spectrogram_transformer_config(_snake_case      )

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

	# load 🤗 model
	UpperCAmelCase__					:  Any					   =					ASTForAudioClassification(_snake_case      )
	model.eval()

	model.load_state_dict(_snake_case      )

	# verify outputs on dummy input
	# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
	UpperCAmelCase__					:  Optional[int]					   =					-4.267_7393 if """speech-commands""" not in model_name else -6.84_5978
	UpperCAmelCase__					:  List[str]					   =					4.568_9974 if """speech-commands""" not in model_name else 5.565_4526
	UpperCAmelCase__					:  Optional[Any]					   =					1024 if """speech-commands""" not in model_name else 128
	UpperCAmelCase__					:  Any					   =					ASTFeatureExtractor(mean=_snake_case							,    std=_snake_case							,    max_length=_snake_case      )

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

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

	UpperCAmelCase__					:  Optional[int]					   =					feature_extractor(_snake_case							,    sampling_rate=16000							,    return_tensors="""pt"""      )

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

	if model_name == "ast-finetuned-audioset-10-10-0.4593":
		UpperCAmelCase__					:  int					   =					torch.tensor([-0.8760, -7.0042, -8.6602]      )
	elif model_name == "ast-finetuned-audioset-10-10-0.450":
		UpperCAmelCase__					:  Optional[int]					   =					torch.tensor([-1.1986, -7.0903, -8.2718]      )
	elif model_name == "ast-finetuned-audioset-10-10-0.448":
		UpperCAmelCase__					:  Dict					   =					torch.tensor([-2.6128, -8.0080, -9.4344]      )
	elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
		UpperCAmelCase__					:  Optional[Any]					   =					torch.tensor([-1.5080, -7.4534, -8.8917]      )
	elif model_name == "ast-finetuned-audioset-12-12-0.447":
		UpperCAmelCase__					:  Tuple					   =					torch.tensor([-0.5050, -6.5833, -8.0843]      )
	elif model_name == "ast-finetuned-audioset-14-14-0.443":
		UpperCAmelCase__					:  Dict					   =					torch.tensor([-0.3826, -7.0336, -8.2413]      )
	elif model_name == "ast-finetuned-audioset-16-16-0.442":
		UpperCAmelCase__					:  Optional[Any]					   =					torch.tensor([-1.2113, -6.9101, -8.3470]      )
	elif model_name == "ast-finetuned-speech-commands-v2":
		UpperCAmelCase__					:  int					   =					torch.tensor([6.1589, -8.0566, -8.7984]      )
	else:
		raise ValueError("""Unknown model name"""      )
	if not torch.allclose(logits[0, :3]							,    _snake_case							,    atol=1e-4      ):
		raise ValueError("""Logits don't match"""      )
	print("""Looks ok!"""      )

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

						__UpperCAmelCase =				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_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase =				{
    'configuration_roberta': ['ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RobertaConfig', 'RobertaOnnxConfig'],
    'tokenization_roberta': ['RobertaTokenizer'],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				['RobertaTokenizerFast']

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST',
          'RobertaForCausalLM',
          'RobertaForMaskedLM',
          'RobertaForMultipleChoice',
          'RobertaForQuestionAnswering',
          'RobertaForSequenceClassification',
          'RobertaForTokenClassification',
          'RobertaModel',
          'RobertaPreTrainedModel',
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFRobertaForCausalLM',
          'TFRobertaForMaskedLM',
          'TFRobertaForMultipleChoice',
          'TFRobertaForQuestionAnswering',
          'TFRobertaForSequenceClassification',
          'TFRobertaForTokenClassification',
          'TFRobertaMainLayer',
          'TFRobertaModel',
          'TFRobertaPreTrainedModel',
      ]

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'FlaxRobertaForCausalLM',
          'FlaxRobertaForMaskedLM',
          'FlaxRobertaForMultipleChoice',
          'FlaxRobertaForQuestionAnswering',
          'FlaxRobertaForSequenceClassification',
          'FlaxRobertaForTokenClassification',
          'FlaxRobertaModel',
          'FlaxRobertaPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaOnnxConfig
      from .tokenization_roberta import RobertaTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_roberta_fast import RobertaTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_roberta import (
                ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                RobertaForCausalLM,
                RobertaForMaskedLM,
                RobertaForMultipleChoice,
                RobertaForQuestionAnswering,
                RobertaForSequenceClassification,
                RobertaForTokenClassification,
                RobertaModel,
                RobertaPreTrainedModel,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_roberta import (
                TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRobertaForCausalLM,
                TFRobertaForMaskedLM,
                TFRobertaForMultipleChoice,
                TFRobertaForQuestionAnswering,
                TFRobertaForSequenceClassification,
                TFRobertaForTokenClassification,
                TFRobertaMainLayer,
                TFRobertaModel,
                TFRobertaPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_roberta import (
                FlaxRobertaForCausalLM,
                FlaxRobertaForMaskedLM,
                FlaxRobertaForMultipleChoice,
                FlaxRobertaForQuestionAnswering,
                FlaxRobertaForSequenceClassification,
                FlaxRobertaForTokenClassification,
                FlaxRobertaModel,
                FlaxRobertaPreTrainedModel,
            )

else:
      import sys

      __UpperCAmelCase =				_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





