






"""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 _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	ASTConfig()

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

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

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

      return config

def _a			(		_snake_case							):







      """simple docstring"""






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

      return name

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






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

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

      return orig_state_dict

def _a			(		_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	[
          """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 _a			(		_snake_case				,	_snake_case				,	_snake_case=False							):







      """simple docstring"""






      UpperCAmelCase          =	get_audio_spectrogram_transformer_config(_snake_case							)

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

      # load 🤗 model
      UpperCAmelCase          =	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          =	-4.2677393 if """speech-commands""" not in model_name else -6.845978
      UpperCAmelCase          =	4.5689974 if """speech-commands""" not in model_name else 5.5654526
      UpperCAmelCase          =	1024 if """speech-commands""" not in model_name else 128
      UpperCAmelCase          =	ASTFeatureExtractor(mean=_snake_case				,	std=_snake_case				,	max_length=_snake_case							)

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

            UpperCAmelCase      , UpperCAmelCase          =	torchaudio.load(_snake_case							)
            UpperCAmelCase          =	waveform.squeeze().numpy()

      UpperCAmelCase          =	feature_extractor(_snake_case				,	sampling_rate=1_6000				,	return_tensors="""pt"""							)

      # forward pass
      UpperCAmelCase          =	model(**_snake_case							)
      UpperCAmelCase          =	outputs.logits

      if model_name == "ast-finetuned-audioset-10-10-0.4593":
            UpperCAmelCase          =	torch.tensor([-0.8760, -7.0042, -8.6602]							)
      elif model_name == "ast-finetuned-audioset-10-10-0.450":
            UpperCAmelCase          =	torch.tensor([-1.1986, -7.0903, -8.2718]							)
      elif model_name == "ast-finetuned-audioset-10-10-0.448":
            UpperCAmelCase          =	torch.tensor([-2.6128, -8.0080, -9.4344]							)
      elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
            UpperCAmelCase          =	torch.tensor([-1.5080, -7.4534, -8.8917]							)
      elif model_name == "ast-finetuned-audioset-12-12-0.447":
            UpperCAmelCase          =	torch.tensor([-0.5050, -6.5833, -8.0843]							)
      elif model_name == "ast-finetuned-audioset-14-14-0.443":
            UpperCAmelCase          =	torch.tensor([-0.3826, -7.0336, -8.2413]							)
      elif model_name == "ast-finetuned-audioset-16-16-0.442":
            UpperCAmelCase          =	torch.tensor([-1.2113, -6.9101, -8.3470]							)
      elif model_name == "ast-finetuned-speech-commands-v2":
            UpperCAmelCase          =	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"""







import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
  from .tokenization_pegasus import PegasusTokenizer
else:
  _UpperCamelCase					      =  None


_UpperCamelCase					      =  logging.get_logger(__name__)


_UpperCamelCase					      =  """▁"""

_UpperCamelCase					      =  {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

_UpperCamelCase					      =  {
    """vocab_file""": {"""google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"""},
    """tokenizer_file""": {
        """google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/tokenizer.json"""
    },
}

_UpperCamelCase					      =  {
    """google/pegasus-xsum""": 512,
}







class  lowerCamelCase__ (	snake_case			):
    SCREAMING_SNAKE_CASE		=      VOCAB_FILES_NAMES
    SCREAMING_SNAKE_CASE		=      PRETRAINED_VOCAB_FILES_MAP
    SCREAMING_SNAKE_CASE		=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    SCREAMING_SNAKE_CASE		=      PegasusTokenizer
    SCREAMING_SNAKE_CASE		=      ['''input_ids''', '''attention_mask''']






    def __init__(			self     ,A=None     ,A=None     ,A="<pad>"     ,A="</s>"     ,A="<unk>"     ,A="<mask_2>"     ,A="<mask_1>"     ,A=None     ,A=103     ,**A     ,):
          UpperCAmelCase          =	offset

          if additional_special_tokens is not None:
                if not isinstance(A     ,A ):
                      raise TypeError(
                          F'''additional_special_tokens should be of type {type(A )}, but is'''
                          F''' {type(A )}''' )

                UpperCAmelCase          =	(
                    ([mask_token_sent] + additional_special_tokens)
                    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                    else additional_special_tokens
                )
                # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
                additional_special_tokens_extended += [
                    F'''<unk_{i}>''' for i in range(len(A )     ,self.offset - 1 )
                ]

                if len(set(A ) ) != len(A ):
                      raise ValueError(
                          """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
                          F''' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.''' )
                UpperCAmelCase          =	additional_special_tokens_extended
          else:
                UpperCAmelCase          =	[mask_token_sent] if mask_token_sent is not None else []
                additional_special_tokens += [F'''<unk_{i}>''' for i in range(2     ,self.offset )]

          super().__init__(
              A     ,tokenizer_file=A     ,pad_token=A     ,eos_token=A     ,unk_token=A     ,mask_token=A     ,mask_token_sent=A     ,offset=A     ,additional_special_tokens=A     ,**A     ,)
          UpperCAmelCase          =	vocab_file
          UpperCAmelCase          =	False if not self.vocab_file else True






    def     _UpperCamelCase							(			self     ,A ):
          UpperCAmelCase          =	set(self.all_special_ids )  # call it once instead of inside list comp
          all_special_ids.remove(self.unk_token_id )  # <unk> is only sometimes special

          if all_special_ids != set(range(len(self.additional_special_tokens ) + 3 ) ):
                raise ValueError(
                    """There should be 3 special tokens: mask_token, pad_token, and eos_token +"""
                    F''' {len(self.additional_special_tokens )} additional_special_tokens, but got {all_special_ids}''' )

          return [1 if x in all_special_ids else 0 for x in seq]






    def     _UpperCamelCase							(			self     ,A     ,A = None     ,A = False ):
          if already_has_special_tokens:
                return self._special_token_mask(A )
          elif token_ids_a is None:
                return self._special_token_mask(A ) + [1]
          else:
                return self._special_token_mask(token_ids_a + token_ids_a ) + [1]






    def     _UpperCamelCase							(			self     ,A     ,A=None ):
          if token_ids_a is None:
                return token_ids_a + [self.eos_token_id]
          # We don't expect to process pairs, but leave the pair logic for API consistency
          return token_ids_a + token_ids_a + [self.eos_token_id]






    def     _UpperCamelCase							(			self     ,A     ,A = None ):
          if not self.can_save_slow_tokenizer:
                raise ValueError(
                    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
                    """tokenizer.""" )

          if not os.path.isdir(A ):
                logger.error(F'''Vocabulary path ({save_directory}) should be a directory''' )
                return
          UpperCAmelCase          =	os.path.join(
              A     ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""] )

          if os.path.abspath(self.vocab_file ) != os.path.abspath(A ):
                copyfile(self.vocab_file     ,A )

          return (out_vocab_file,)



