





'''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   :	Optional[int]   				=logging.get_logger(__name__)
def 				UpperCAmelCase_    (		__lowerCamelCase   :   List[str]					):
   lowercase_     :Optional[int]            =					ASTConfig()

   if "10-10" in model_name:
      pass
   elif "speech-commands" in model_name:
      lowercase_     :Tuple            =					1_28
   elif "12-12" in model_name:
      lowercase_     :Optional[Any]            =					12
      lowercase_     :Any            =					12
   elif "14-14" in model_name:
      lowercase_     :str            =					14
      lowercase_     :Optional[Any]            =					14
   elif "16-16" in model_name:
      lowercase_     :Union[str, Any]            =					16
      lowercase_     :List[str]            =					16
   else:
      raise ValueError("Model not supported"					)

   lowercase_     :Optional[int]            =					"""huggingface/label-files"""
   if "speech-commands" in model_name:
      lowercase_     :Optional[Any]            =					35
      lowercase_     :List[Any]            =					"""speech-commands-v2-id2label.json"""
   else:
      lowercase_     :str            =					5_27
      lowercase_     :str            =					"""audioset-id2label.json"""

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

   return config
def 				UpperCAmelCase_    (		__lowerCamelCase   :   Tuple					):



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

   return name
def 				UpperCAmelCase_    (		__lowerCamelCase   :   Optional[int]      ,__lowerCamelCase   :   Union[str, Any]					):



   for key in orig_state_dict.copy().keys():
      lowercase_     :Tuple            =					orig_state_dict.pop(a__					)

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

   return orig_state_dict
def 				UpperCAmelCase_    (		__lowerCamelCase   :   int					):
   lowercase_     :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(a__      ,a__					)
@torch.no_grad()
def 				UpperCAmelCase_    (		__lowerCamelCase   :   Union[str, Any]      ,__lowerCamelCase   :   Dict      ,__lowerCamelCase   :   Dict=False					):
   lowercase_     :List[str]            =					get_audio_spectrogram_transformer_config(a__					)

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

   # load 🤗 model
   lowercase_     :Any            =					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
   lowercase_     :Dict            =					-4.2_677_393 if """speech-commands""" not in model_name else -6.845_978
   lowercase_     :int            =					4.5_689_974 if """speech-commands""" not in model_name else 5.5_654_526
   lowercase_     :Tuple            =					10_24 if """speech-commands""" not in model_name else 1_28
   lowercase_     :Any            =					ASTFeatureExtractor(mean=a__      ,std=a__      ,max_length=a__					)

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

      lowercase_     :Dict            =					torchaudio.load(a__					)
      lowercase_     :Tuple            =					waveform.squeeze().numpy()

   lowercase_     :int            =					feature_extractor(a__      ,sampling_rate=1_60_00      ,return_tensors="pt"					)

   # forward pass
   lowercase_     :Dict            =					model(**a__					)
   lowercase_     :Any            =					outputs.logits

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

 lowerCAmelCase   :	str   				=parser.parse_args()
 convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





from typing import TYPE_CHECKING

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


__snake_case				:      Optional[int]   		=			{
    """configuration_xlm""": ["""XLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XLMConfig""", """XLMOnnxConfig"""],
    """tokenization_xlm""": ["""XLMTokenizer"""],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      Optional[Any]   		=			[
	    """XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """XLMForMultipleChoice""",
	    """XLMForQuestionAnswering""",
	    """XLMForQuestionAnsweringSimple""",
	    """XLMForSequenceClassification""",
	    """XLMForTokenClassification""",
	    """XLMModel""",
	    """XLMPreTrainedModel""",
	    """XLMWithLMHeadModel""",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:      Optional[Any]   		=			[
	    """TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
	    """TFXLMForMultipleChoice""",
	    """TFXLMForQuestionAnsweringSimple""",
	    """TFXLMForSequenceClassification""",
	    """TFXLMForTokenClassification""",
	    """TFXLMMainLayer""",
	    """TFXLMModel""",
	    """TFXLMPreTrainedModel""",
	    """TFXLMWithLMHeadModel""",
	]


if TYPE_CHECKING:
	from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
	from .tokenization_xlm import XLMTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_xlm import (
		    XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    XLMForMultipleChoice,
		    XLMForQuestionAnswering,
		    XLMForQuestionAnsweringSimple,
		    XLMForSequenceClassification,
		    XLMForTokenClassification,
		    XLMModel,
		    XLMPreTrainedModel,
		    XLMWithLMHeadModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_xlm import (
		    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFXLMForMultipleChoice,
		    TFXLMForQuestionAnsweringSimple,
		    TFXLMForSequenceClassification,
		    TFXLMForTokenClassification,
		    TFXLMMainLayer,
		    TFXLMModel,
		    TFXLMPreTrainedModel,
		    TFXLMWithLMHeadModel,
		)

else:
	import sys

	__snake_case				:      Union[str, Any]   		=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


