






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 		_a       ( a					:str					)				->     List[Any]:
 a						    = ASTConfig()

 if "10-10" in model_name:
  pass
 elif "speech-commands" in model_name:
  a						    = 128
 elif "12-12" in model_name:
  a						    = 12
  a						    = 12
 elif "14-14" in model_name:
  a						    = 14
  a						    = 14
 elif "16-16" in model_name:
  a						    = 16
  a						    = 16
 else:
  raise ValueError('''Model not supported'''					)

 a						    = '''huggingface/label-files'''
 if "speech-commands" in model_name:
  a						    = 35
  a						    = '''speech-commands-v2-id2label.json'''
 else:
  a						    = 527
  a						    = '''audioset-id2label.json'''

 a						    = json.load(open(hf_hub_download(a      ,				a      ,				repo_type='''dataset'''					)      ,				'''r'''					)					)
 a						    = {int(a					): v for k, v in idalabel.items()}
 a						    = idalabel
 a						    = {v: k for k, v in idalabel.items()}

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

 return name
def 		_a       ( a					:Any      ,				a					:Any					)				->     Dict:
 for key in orig_state_dict.copy().keys():
  a						    = orig_state_dict.pop(a					)

  if "qkv" in key:
   a						    = key.split('''.'''					)
   a						    = int(key_split[3]					)
   a						    = config.hidden_size
   if "weight" in key:
    a						    = val[:dim, :]
    a						    = val[dim : dim * 2, :]
    a						    = val[-dim:, :]
   else:
    a						    = val[:dim]
    a						    = val[dim : dim * 2]
    a						    = val[-dim:]
  else:
   a						    = val

 return orig_state_dict
def 		_a       ( a					:Optional[Any]					)				->     Union[str, Any]:
 a						    = [
     '''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 		_a       ( a					:Union[str, Any]      ,				a					:Dict      ,				a					:Optional[Any]=False					)				->     Union[str, Any]:
 a						    = get_audio_spectrogram_transformer_config(a					)

 a						    = {
     '''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						    = model_name_to_url[model_name]
 a						    = torch.hub.load_state_dict_from_url(a      ,				map_location='''cpu'''					)
 # remove some keys
 remove_keys(a					)
 # rename some keys
 a						    = convert_state_dict(a      ,				a					)

 # load 🤗 model
 a						    = 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
 a						    = -4.2_677_393 if '''speech-commands''' not in model_name else -6.845_978
 a						    = 4.5_689_974 if '''speech-commands''' not in model_name else 5.5_654_526
 a						    = 1_024 if '''speech-commands''' not in model_name else 128
 a						    = ASTFeatureExtractor(mean=a      ,				std=a      ,				max_length=a					)

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

  a    ,	a						    = torchaudio.load(a					)
  a						    = waveform.squeeze().numpy()

 a						    = feature_extractor(a      ,				sampling_rate=16_000      ,				return_tensors='''pt'''					)

 # forward pass
 a						    = model(**a					)
 a						    = outputs.logits

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










from ... import PretrainedConfig


UpperCAmelCase__							  =						{
    "sijunhe/nezha-cn-base": "https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json",
}






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
      __snake_case	    =    '''nezha'''







      def __init__( self :		Union[str, Any]					,   __UpperCAmelCase :		List[str]=21_128					,   __UpperCAmelCase :		Any=768					,   __UpperCAmelCase :		int=12					,   __UpperCAmelCase :		Optional[int]=12					,   __UpperCAmelCase :		Tuple=3_072					,   __UpperCAmelCase :		Union[str, Any]="gelu"					,   __UpperCAmelCase :		List[str]=0.1					,   __UpperCAmelCase :		Any=0.1					,   __UpperCAmelCase :		str=512					,   __UpperCAmelCase :		List[Any]=64					,   __UpperCAmelCase :		Tuple=2					,   __UpperCAmelCase :		List[Any]=0.02					,   __UpperCAmelCase :		List[str]=1e-1_2					,   __UpperCAmelCase :		int=0.1					,   __UpperCAmelCase :		Dict=0					,   __UpperCAmelCase :		Any=2					,   __UpperCAmelCase :		Optional[Any]=3					,   __UpperCAmelCase :		int=True					,   **__UpperCAmelCase :		Optional[Any]					,   )       ->Any:




       """simple docstring"""


       super().__init__(pad_token_id=__UpperCAmelCase					,   bos_token_id=__UpperCAmelCase					,   eos_token_id=__UpperCAmelCase					,   **__UpperCAmelCase    )

       a						    = vocab_size
       a						    = hidden_size
       a						    = num_hidden_layers
       a						    = num_attention_heads
       a						    = hidden_act
       a						    = intermediate_size
       a						    = hidden_dropout_prob
       a						    = attention_probs_dropout_prob
       a						    = max_position_embeddings
       a						    = max_relative_position
       a						    = type_vocab_size
       a						    = initializer_range
       a						    = layer_norm_eps
       a						    = classifier_dropout
       a						    = use_cache


