



'''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         =			logging.get_logger(__name__)







def    _snake_case (       __snake_case    )     ->     Dict:






      '''simple docstring'''







      UpperCAmelCase_  :	List[str]				      =		ASTConfig()

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

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

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

      return config







def    _snake_case (       __snake_case    )     ->     str:






      '''simple docstring'''







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

      return name







def    _snake_case (       __snake_case					,  __snake_case    )     ->     str:






      '''simple docstring'''







      for key in orig_state_dict.copy().keys():
            UpperCAmelCase_  :	str				      =		orig_state_dict.pop(__snake_case    )

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

      return orig_state_dict







def    _snake_case (       __snake_case    )     ->     Optional[int]:






      '''simple docstring'''







      UpperCAmelCase_  :	Optional[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(__snake_case					,  __snake_case    )







@torch.no_grad()
def    _snake_case (       __snake_case					,  __snake_case					,  __snake_case=False    )     ->     Tuple:






      '''simple docstring'''







      UpperCAmelCase_  :	List[str]				      =		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[str]				      =		model_name_to_url[model_name]
      UpperCAmelCase_  :	str				      =		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_  :	Tuple				      =		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_  :	str				      =		-4.267_7393 if "speech-commands" not in model_name else -6.84_5978
      UpperCAmelCase_  :	Optional[int]				      =		4.568_9974 if "speech-commands" not in model_name else 5.565_4526
      UpperCAmelCase_  :	int				      =		1_0_2_4 if "speech-commands" not in model_name else 1_2_8
      UpperCAmelCase_  :	str				      =		ASTFeatureExtractor(mean=__snake_case					,  std=__snake_case					,  max_length=__snake_case    )

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

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

      UpperCAmelCase_  :	Dict				      =		feature_extractor(__snake_case					,  sampling_rate=1_6_0_0_0					,  return_tensors="pt"    )

      # forward pass
      UpperCAmelCase_  :	List[Any]				      =		model(**__snake_case    )
      UpperCAmelCase_  :	Dict				      =		outputs.logits

      if model_name == "ast-finetuned-audioset-10-10-0.4593":
            UpperCAmelCase_  :	Dict				      =		torch.tensor([-0.8760, -7.0042, -8.6602]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.450":
            UpperCAmelCase_  :	int				      =		torch.tensor([-1.1986, -7.0903, -8.2718]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.448":
            UpperCAmelCase_  :	Any				      =		torch.tensor([-2.6128, -8.0080, -9.4344]    )
      elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
            UpperCAmelCase_  :	List[str]				      =		torch.tensor([-1.5080, -7.4534, -8.8917]    )
      elif model_name == "ast-finetuned-audioset-12-12-0.447":
            UpperCAmelCase_  :	int				      =		torch.tensor([-0.5050, -6.5833, -8.0843]    )
      elif model_name == "ast-finetuned-audioset-14-14-0.443":
            UpperCAmelCase_  :	List[str]				      =		torch.tensor([-0.3826, -7.0336, -8.2413]    )
      elif model_name == "ast-finetuned-audioset-16-16-0.442":
            UpperCAmelCase_  :	List[Any]				      =		torch.tensor([-1.2113, -6.9101, -8.3470]    )
      elif model_name == "ast-finetuned-speech-commands-v2":
            UpperCAmelCase_  :	List[Any]				      =		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__":
   __lowerCamelCase         =			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         =			parser.parse_args()
   convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch

import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin, SchedulerOutput


@dataclass
class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				1



       @register_to_config
       def __init__(     self     ,lowercase = 2000     ,lowercase = 0.15     ,lowercase = 0.01     ,lowercase = 1348.0     ,lowercase = 1E-5     ,lowercase = 1     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[int]				      =		sigma_max

             # setable values
             UpperCAmelCase_  :	Optional[int]				      =		None

             self.set_sigmas(lowercase     ,lowercase     ,lowercase     ,lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None):

             """simple docstring"""





             return sample



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		sampling_eps if sampling_eps is not None else self.config.sampling_eps

             UpperCAmelCase_  :	List[Any]				      =		torch.linspace(1     ,lowercase     ,lowercase     ,device=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase = None     ,lowercase = None     ,lowercase = None):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		sigma_min if sigma_min is not None else self.config.sigma_min
             UpperCAmelCase_  :	int				      =		sigma_max if sigma_max is not None else self.config.sigma_max
             UpperCAmelCase_  :	Union[str, Any]				      =		sampling_eps if sampling_eps is not None else self.config.sampling_eps
             if self.timesteps is None:
                   self.set_timesteps(lowercase     ,lowercase)

             UpperCAmelCase_  :	Union[str, Any]				      =		sigma_min * (sigma_max / sigma_min) ** (self.timesteps / sampling_eps)
             UpperCAmelCase_  :	Optional[int]				      =		torch.exp(torch.linspace(math.log(lowercase)     ,math.log(lowercase)     ,lowercase))
             UpperCAmelCase_  :	Any				      =		torch.tensor([sigma_min * (sigma_max / sigma_min) ** t for t in self.timesteps])



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





             return torch.where(
                 timesteps == 0     ,torch.zeros_like(t.to(timesteps.device))     ,self.discrete_sigmas[timesteps - 1].to(timesteps.device)     ,)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





             if self.timesteps is None:
                   raise ValueError(
                       "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler")

             UpperCAmelCase_  :	Optional[int]				      =		timestep * torch.ones(
                 sample.shape[0]     ,device=sample.device)  # torch.repeat_interleave(timestep, sample.shape[0])
             UpperCAmelCase_  :	Tuple				      =		(timestep * (len(self.timesteps) - 1)).long()

             # mps requires indices to be in the same device, so we use cpu as is the default with cuda
             UpperCAmelCase_  :	Optional[int]				      =		timesteps.to(self.discrete_sigmas.device)

             UpperCAmelCase_  :	Optional[Any]				      =		self.discrete_sigmas[timesteps].to(sample.device)
             UpperCAmelCase_  :	Optional[Any]				      =		self.get_adjacent_sigma(lowercase     ,lowercase).to(sample.device)
             UpperCAmelCase_  :	Any				      =		torch.zeros_like(lowercase)
             UpperCAmelCase_  :	Dict				      =		(sigma**2 - adjacent_sigma**2) ** 0.5

             # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
             # also equation 47 shows the analog from SDE models to ancestral sampling methods
             UpperCAmelCase_  :	Dict				      =		diffusion.flatten()
             while len(diffusion.shape) < len(sample.shape):
                   UpperCAmelCase_  :	List[str]				      =		diffusion.unsqueeze(-1)
             UpperCAmelCase_  :	List[Any]				      =		drift - diffusion**2 * model_output

             #  equation 6: sample noise for the diffusion term of
             UpperCAmelCase_  :	Union[str, Any]				      =		randn_tensor(
                 sample.shape     ,layout=sample.layout     ,generator=lowercase     ,device=sample.device     ,dtype=sample.dtype)
             UpperCAmelCase_  :	Any				      =		sample - drift  # subtract because `dt` is a small negative timestep
             # TODO is the variable diffusion the correct scaling term for the noise?
             UpperCAmelCase_  :	Tuple				      =		prev_sample_mean + diffusion * noise  # add impact of diffusion field g

             if not return_dict:
                   return (prev_sample, prev_sample_mean)

             return SdeVeOutput(prev_sample=lowercase     ,prev_sample_mean=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase = None     ,lowercase = True     ,):

             """simple docstring"""





             if self.timesteps is None:
                   raise ValueError(
                       "`self.timesteps` is not set, you need to run 'set_timesteps' after creating the scheduler")

             # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
             # sample noise for correction
             UpperCAmelCase_  :	int				      =		randn_tensor(sample.shape     ,layout=sample.layout     ,generator=lowercase).to(sample.device)

             # compute step size from the model_output, the noise, and the snr
             UpperCAmelCase_  :	Union[str, Any]				      =		torch.norm(model_output.reshape(model_output.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	Optional[Any]				      =		torch.norm(noise.reshape(noise.shape[0]     ,-1)     ,dim=-1).mean()
             UpperCAmelCase_  :	List[Any]				      =		(self.config.snr * noise_norm / grad_norm) ** 2 * 2
             UpperCAmelCase_  :	Optional[Any]				      =		step_size * torch.ones(sample.shape[0]).to(sample.device)
             # self.repeat_scalar(step_size, sample.shape[0])

             # compute corrected sample: model_output term and noise term
             UpperCAmelCase_  :	Any				      =		step_size.flatten()
             while len(step_size.shape) < len(sample.shape):
                   UpperCAmelCase_  :	Tuple				      =		step_size.unsqueeze(-1)
             UpperCAmelCase_  :	Dict				      =		sample + step_size * model_output
             UpperCAmelCase_  :	int				      =		prev_sample_mean + ((step_size * 2) ** 0.5) * noise

             if not return_dict:
                   return (prev_sample,)

             return SchedulerOutput(prev_sample=lowercase)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		timesteps.to(original_samples.device)
             UpperCAmelCase_  :	List[str]				      =		self.discrete_sigmas.to(original_samples.device)[timesteps]
             UpperCAmelCase_  :	Tuple				      =		(
                 noise * sigmas[:, None, None, None]
                 if noise is not None
                 else torch.randn_like(lowercase) * sigmas[:, None, None, None]
             )
             UpperCAmelCase_  :	Tuple				      =		noise + original_samples
             return noisy_samples



       def __len__(     self):

             """simple docstring"""





             return self.config.num_train_timesteps
