

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

def 			a__     (		UpperCAmelCase							:			Union[str, Any]			)  ->						Optional[int]:
  UpperCAmelCase       :							str  = ASTConfig()

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

  UpperCAmelCase       :							Union[str, Any]  = '''huggingface/label-files'''
  if "speech-commands" in model_name:
    UpperCAmelCase       :							str  = 35
    UpperCAmelCase       :							Optional[Any]  = '''speech-commands-v2-id2label.json'''
  else:
    UpperCAmelCase       :							List[Any]  = 527
    UpperCAmelCase       :							int  = '''audioset-id2label.json'''

  UpperCAmelCase       :							str  = json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE__				, SCREAMING_SNAKE_CASE__				, repo_type='''dataset'''			)				, '''r'''			)			)
  UpperCAmelCase       :							List[str]  = {int(SCREAMING_SNAKE_CASE__			): v for k, v in idalabel.items()}
  UpperCAmelCase       :							Tuple  = idalabel
  UpperCAmelCase       :							List[Any]  = {v: k for k, v in idalabel.items()}

  return config

def 			a__     (		UpperCAmelCase							:			Union[str, Any]			)  ->						Union[str, Any]:
  if "module.v" in name:
    UpperCAmelCase       :							Dict  = 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       :							Optional[int]  = name.replace('''dist_token'''				, '''embeddings.distillation_token'''			)
  if "pos_embed" in name:
    UpperCAmelCase       :							int  = name.replace('''pos_embed'''				, '''embeddings.position_embeddings'''			)
  if "patch_embed.proj" in name:
    UpperCAmelCase       :							Optional[Any]  = name.replace('''patch_embed.proj'''				, '''embeddings.patch_embeddings.projection'''			)
  # transformer blocks
  if "blocks" in name:
    UpperCAmelCase       :							Any  = name.replace('''blocks'''				, '''encoder.layer'''			)
  if "attn.proj" in name:
    UpperCAmelCase       :							Tuple  = name.replace('''attn.proj'''				, '''attention.output.dense'''			)
  if "attn" in name:
    UpperCAmelCase       :							List[Any]  = name.replace('''attn'''				, '''attention.self'''			)
  if "norm1" in name:
    UpperCAmelCase       :							Tuple  = name.replace('''norm1'''				, '''layernorm_before'''			)
  if "norm2" in name:
    UpperCAmelCase       :							Optional[int]  = name.replace('''norm2'''				, '''layernorm_after'''			)
  if "mlp.fc1" in name:
    UpperCAmelCase       :							Tuple  = 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       :							Optional[int]  = name.replace('''audio_spectrogram_transformer.norm'''				, '''audio_spectrogram_transformer.layernorm'''			)
  # classifier head
  if "module.mlp_head.0" in name:
    UpperCAmelCase       :							Dict  = name.replace('''module.mlp_head.0'''				, '''classifier.layernorm'''			)
  if "module.mlp_head.1" in name:
    UpperCAmelCase       :							Dict  = name.replace('''module.mlp_head.1'''				, '''classifier.dense'''			)

  return name

def 			a__     (		UpperCAmelCase							:			Optional[int]				, UpperCAmelCase							:			Union[str, Any]			)  ->						Tuple:
  for key in orig_state_dict.copy().keys():
    UpperCAmelCase       :							List[Any]  = orig_state_dict.pop(SCREAMING_SNAKE_CASE__			)

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

  return orig_state_dict

def 			a__     (		UpperCAmelCase							:			Any			)  ->						Union[str, Any]:
  UpperCAmelCase       :							Union[str, 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(SCREAMING_SNAKE_CASE__				, SCREAMING_SNAKE_CASE__			)

@torch.no_grad()
def 			a__     (		UpperCAmelCase							:			List[Any]				, UpperCAmelCase							:			Union[str, Any]				, UpperCAmelCase							:			List[Any]=False			)  ->						Union[str, Any]:
  UpperCAmelCase       :							Tuple  = get_audio_spectrogram_transformer_config(SCREAMING_SNAKE_CASE__			)

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

  # load 🤗 model
  UpperCAmelCase       :							Union[str, Any]  = ASTForAudioClassification(SCREAMING_SNAKE_CASE__			)
  model.eval()

  model.load_state_dict(SCREAMING_SNAKE_CASE__			)

  # verify outputs on dummy input
  # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
  UpperCAmelCase       :							Union[str, Any]  = -4.2677393 if '''speech-commands''' not in model_name else -6.845978
  UpperCAmelCase       :							Union[str, Any]  = 4.5689974 if '''speech-commands''' not in model_name else 5.5654526
  UpperCAmelCase       :							Optional[int]  = 1_024 if '''speech-commands''' not in model_name else 128
  UpperCAmelCase       :							Optional[int]  = ASTFeatureExtractor(mean=SCREAMING_SNAKE_CASE__				, std=SCREAMING_SNAKE_CASE__				, max_length=SCREAMING_SNAKE_CASE__			)

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

    UpperCAmelCase ,						UpperCAmelCase       :							int  = torchaudio.load(SCREAMING_SNAKE_CASE__			)
    UpperCAmelCase       :							Dict  = waveform.squeeze().numpy()

  UpperCAmelCase       :							List[str]  = feature_extractor(SCREAMING_SNAKE_CASE__				, sampling_rate=16_000				, return_tensors='''pt'''			)

  # forward pass
  UpperCAmelCase       :							Optional[int]  = model(**SCREAMING_SNAKE_CASE__			)
  UpperCAmelCase       :							int  = outputs.logits

  if model_name == "ast-finetuned-audioset-10-10-0.4593":
    UpperCAmelCase       :							List[str]  = torch.tensor([-0.8760, -7.0042, -8.6602]			)
  elif model_name == "ast-finetuned-audioset-10-10-0.450":
    UpperCAmelCase       :							Dict  = torch.tensor([-1.1986, -7.0903, -8.2718]			)
  elif model_name == "ast-finetuned-audioset-10-10-0.448":
    UpperCAmelCase       :							Tuple  = torch.tensor([-2.6128, -8.0080, -9.4344]			)
  elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
    UpperCAmelCase       :							Optional[Any]  = torch.tensor([-1.5080, -7.4534, -8.8917]			)
  elif model_name == "ast-finetuned-audioset-12-12-0.447":
    UpperCAmelCase       :							Optional[Any]  = torch.tensor([-0.5050, -6.5833, -8.0843]			)
  elif model_name == "ast-finetuned-audioset-14-14-0.443":
    UpperCAmelCase       :							int  = torch.tensor([-0.3826, -7.0336, -8.2413]			)
  elif model_name == "ast-finetuned-audioset-16-16-0.442":
    UpperCAmelCase       :							str  = torch.tensor([-1.2113, -6.9101, -8.3470]			)
  elif model_name == "ast-finetuned-speech-commands-v2":
    UpperCAmelCase       :							str  = torch.tensor([6.1589, -8.0566, -8.7984]			)
  else:
    raise ValueError('''Unknown model name'''			)
  if not torch.allclose(logits[0, :3]				, SCREAMING_SNAKE_CASE__				, atol=1E-4			):
    raise ValueError('''Logits don\'t match'''			)
  print('''Looks ok!'''			)

  if pytorch_dump_folder_path is not None:
    Path(SCREAMING_SNAKE_CASE__			).mkdir(exist_ok=SCREAMING_SNAKE_CASE__			)
    print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''			)
    model.save_pretrained(SCREAMING_SNAKE_CASE__			)
    print(f'''Saving feature extractor to {pytorch_dump_folder_path}'''			)
    feature_extractor.save_pretrained(SCREAMING_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				:   str					   =							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)








from sklearn.metrics import mean_squared_error

import datasets


lowerCAmelCase_       	=	'''\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
'''

lowerCAmelCase_       	=	'''\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
'''


lowerCAmelCase_       	=	'''
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {"raw_values", "uniform_average"} or array-like of shape (n_outputs,), default="uniform_average"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

        "raw_values" : Returns a full set of errors in case of multioutput input.

        "uniform_average" : Errors of all outputs are averaged with uniform weight.

    squared : bool, default=True
        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.

Returns:
    mse : mean squared error.
Examples:

    >>> mse_metric = datasets.load_metric("mse")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.375}
    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)
    >>> print(rmse_result)
    {\'mse\': 0.6123724356957945}

    If you\'re using multi-dimensional lists, then set the config as follows :

    >>> mse_metric = datasets.load_metric("mse", "multilist")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.7083333333333334}
    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput=\'raw_values\')
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {\'mse\': array([0.41666667, 1. ])}
'''





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION	,       _KWARGS_DESCRIPTION					)
class 	snake_case_   (		datasets.Metric					):







			'''simple docstring'''

			def 				snake_case__(     self :				Optional[int]							)      ->List[Any]:
									return datasets.MetricInfo(
									    description=_DESCRIPTION							,   citation=_CITATION							,   inputs_description=_KWARGS_DESCRIPTION							,   features=datasets.Features(self._get_feature_types()							)							,   reference_urls=[
									        '''https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html'''
									    ]							,   )







			def 				snake_case__(     self :				List[Any]							)      ->Optional[int]:
									if self.config_name == "multilist":
															return {
															    "predictions": datasets.Sequence(datasets.Value('''float'''							)							),
															    "references": datasets.Sequence(datasets.Value('''float'''							)							),
															}
									else:
															return {
															    "predictions": datasets.Value('''float'''							),
															    "references": datasets.Value('''float'''							),
															}







			def 				snake_case__(     self :				int							,   _UpperCamelCase :				int							,   _UpperCamelCase :				Union[str, Any]							,   _UpperCamelCase :				Any=None							,   _UpperCamelCase :				Optional[int]="uniform_average"							,   _UpperCamelCase :				Tuple=True							)      ->Tuple:
									snake_case_   =	mean_squared_error(
									    _UpperCamelCase							,   _UpperCamelCase							,   sample_weight=_UpperCamelCase							,   multioutput=_UpperCamelCase							,   squared=_UpperCamelCase							)

									return {"mse": mse}