



"""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       :					Dict      					=					logging.get_logger(__name__)
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Dict    ):


      '''simple docstring'''






      snake_case_   : Dict						=			ASTConfig()

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

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

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

      return config
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]    ):


      '''simple docstring'''






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

      return name
def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[str]					, __UpperCamelCase      :			Optional[Any]    ):


      '''simple docstring'''






      for key in orig_state_dict.copy().keys():
            snake_case_   : Optional[int]						=			orig_state_dict.pop(__UpperCamelCase    )

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

      return orig_state_dict
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Any    ):


      '''simple docstring'''






      snake_case_   : 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(__UpperCamelCase					, __UpperCamelCase    )
@torch.no_grad()
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Union[str, Any]					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			Union[str, Any]=False    ):


      '''simple docstring'''






      snake_case_   : Any						=			get_audio_spectrogram_transformer_config(__UpperCamelCase    )

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

      # load 🤗 model
      snake_case_   : List[str]						=			ASTForAudioClassification(__UpperCamelCase    )
      model.eval()

      model.load_state_dict(__UpperCamelCase    )

      # verify outputs on dummy input
      # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
      snake_case_   : List[str]						=			-4.2_677_393 if """speech-commands""" not in model_name else -6.845_978
      snake_case_   : Optional[Any]						=			4.5_689_974 if """speech-commands""" not in model_name else 5.5_654_526
      snake_case_   : List[Any]						=			1_0_2_4 if """speech-commands""" not in model_name else 1_2_8
      snake_case_   : Dict						=			ASTFeatureExtractor(mean=__UpperCamelCase					, std=__UpperCamelCase					, max_length=__UpperCamelCase    )

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

            snake_case_			,		snake_case_   : Union[str, Any]						=			torchaudio.load(__UpperCamelCase    )
            snake_case_   : Union[str, Any]						=			waveform.squeeze().numpy()

      snake_case_   : Optional[int]						=			feature_extractor(__UpperCamelCase					, sampling_rate=1_6_0_0_0					, return_tensors="""pt"""    )

      # forward pass
      snake_case_   : List[str]						=			model(**__UpperCamelCase    )
      snake_case_   : Tuple						=			outputs.logits

      if model_name == "ast-finetuned-audioset-10-10-0.4593":
            snake_case_   : Dict						=			torch.tensor([-0.8_760, -7.0_042, -8.6_602]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.450":
            snake_case_   : List[Any]						=			torch.tensor([-1.1_986, -7.0_903, -8.2_718]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.448":
            snake_case_   : List[str]						=			torch.tensor([-2.6_128, -8.0_080, -9.4_344]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
            snake_case_   : Dict						=			torch.tensor([-1.5_080, -7.4_534, -8.8_917]    )
      elif model_name == "ast-finetuned-audioset-12-12-0.447":
            snake_case_   : Tuple						=			torch.tensor([-0.5_050, -6.5_833, -8.0_843]    )
      elif model_name == "ast-finetuned-audioset-14-14-0.443":
            snake_case_   : List[str]						=			torch.tensor([-0.3_826, -7.0_336, -8.2_413]    )
      elif model_name == "ast-finetuned-audioset-16-16-0.442":
            snake_case_   : Tuple						=			torch.tensor([-1.2_113, -6.9_101, -8.3_470]    )
      elif model_name == "ast-finetuned-speech-commands-v2":
            snake_case_   : Union[str, 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]					, __UpperCamelCase					, atol=1E-4    ):
            raise ValueError("""Logits don\'t match"""    )
      print("""Looks ok!"""    )

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

      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       :					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       :					Tuple      					=					parser.parse_args()
     convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






"""simple docstring"""





from .data_collator import (
    DataCollatorForLanguageModeling,
    DataCollatorForPermutationLanguageModeling,
    DataCollatorForSeqaSeq,
    DataCollatorForSOP,
    DataCollatorForTokenClassification,
    DataCollatorForWholeWordMask,
    DataCollatorWithPadding,
    DefaultDataCollator,
    default_data_collator,
)
from .metrics import glue_compute_metrics, xnli_compute_metrics
from .processors import (
    DataProcessor,
    InputExample,
    InputFeatures,
    SingleSentenceClassificationProcessor,
    SquadExample,
    SquadFeatures,
    SquadVaProcessor,
    SquadVaProcessor,
    glue_convert_examples_to_features,
    glue_output_modes,
    glue_processors,
    glue_tasks_num_labels,
    squad_convert_examples_to_features,
    xnli_output_modes,
    xnli_processors,
    xnli_tasks_num_labels,
)

