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





def    _lowerCAmelCase    (     lowerCamelCase__							:							Optional[int]				)  ->    Any:
  _SCREAMING_SNAKE_CASE   :				Optional[int]							    = ASTConfig()

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

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

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

  return config





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

  return name





def    _lowerCAmelCase    (     lowerCamelCase__							:							Union[str, Any],						lowerCamelCase__							:							List[str]				)  ->    Optional[int]:
  for key in orig_state_dict.copy().keys():
    _SCREAMING_SNAKE_CASE   :				Any							    = orig_state_dict.pop(lowerCamelCase__				)

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

  return orig_state_dict





def    _lowerCAmelCase    (     lowerCamelCase__							:							Tuple				)  ->    List[str]:
  _SCREAMING_SNAKE_CASE   :				Tuple							    = [
      "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(lowerCamelCase__,						lowerCamelCase__				)





@torch.no_grad()
def    _lowerCAmelCase    (     lowerCamelCase__							:							int,						lowerCamelCase__							:							Dict,						lowerCamelCase__							:							Optional[int]=False				)  ->    Optional[Any]:
  _SCREAMING_SNAKE_CASE   :				Dict							    = get_audio_spectrogram_transformer_config(lowerCamelCase__				)

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

  # load 🤗 model
  _SCREAMING_SNAKE_CASE   :				Tuple							    = ASTForAudioClassification(lowerCamelCase__				)
  model.eval()

  model.load_state_dict(lowerCamelCase__				)

  # verify outputs on dummy input
  # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
  _SCREAMING_SNAKE_CASE   :				List[str]							    = -4.2677393 if "speech-commands" not in model_name else -6.845978
  _SCREAMING_SNAKE_CASE   :				Dict							    = 4.5689974 if "speech-commands" not in model_name else 5.5654526
  _SCREAMING_SNAKE_CASE   :				Dict							    = 1_0_2_4 if "speech-commands" not in model_name else 1_2_8
  _SCREAMING_SNAKE_CASE   :				str							    = ASTFeatureExtractor(mean=lowerCamelCase__,						std=lowerCamelCase__,						max_length=lowerCamelCase__				)

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

    _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				int							    = torchaudio.load(lowerCamelCase__				)
    _SCREAMING_SNAKE_CASE   :				int							    = waveform.squeeze().numpy()

  _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = feature_extractor(lowerCamelCase__,						sampling_rate=1_6_0_0_0,						return_tensors="pt"				)

  # forward pass
  _SCREAMING_SNAKE_CASE   :				Tuple							    = model(**lowerCamelCase__				)
  _SCREAMING_SNAKE_CASE   :				Optional[int]							    = outputs.logits

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

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

  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__":
      lowercase_       :						Tuple     					=	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.'''
      )

      lowercase_       :						Optional[Any]     					=	parser.parse_args()
      convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





"""simple docstring"""
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Model name or path of model to be trained."""}  )
     A__      =			field(
         default="""./"""      ,       metadata={"""help""": """Save dir where model repo is cloned and models updates are saved to."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot-clean-train"""      ,       metadata={"""help""": """Name or path of training dataset."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot-clean-valid"""      ,       metadata={"""help""": """Name or path of validation dataset."""}  )
     A__      =			field(default=2      ,       metadata={"""help""": """Batch size for training."""}  )
     A__      =			field(default=2      ,       metadata={"""help""": """Batch size for evaluation."""}  )
     A__      =			field(default=0.1      ,       metadata={"""help""": """Value of weight decay."""}  )
     A__      =			field(
         default=10000      ,       metadata={"""help""": """Size of buffer used to shuffle streaming dataset."""}  )
     A__      =			field(default=2e-4      ,       metadata={"""help""": """Learning rate fo training."""}  )
     A__      =			field(default="""cosine"""      ,       metadata={"""help""": """Learning rate."""}  )
     A__      =			field(
         default=750      ,       metadata={"""help""": """Number of warmup steps in the learning rate schedule."""}  )
     A__      =			field(
         default=16      ,       metadata={"""help""": """Number of gradient accumulation steps."""}  )
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Use gradient checkpointing to reduce memory footprint."""}  )
     A__      =			field(default=50000      ,       metadata={"""help""": """Maximum number of training steps."""}  )
     A__      =			field(
         default=-1      ,       metadata={"""help""": """Maximum number of evaluation steps. If -1 the full dataset is evaluated."""}  )
     A__      =			field(default=1024      ,       metadata={"""help""": """Sequence lengths used for training."""}  )
     A__      =			field(default=1      ,       metadata={"""help""": """Training seed."""}  )
     A__      =			field(
         default=1024      ,       metadata={"""help""": """Interval to save checkpoints. Measured as number of forward passes not training steps."""}      ,       )
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """States path if the training should continue from a checkpoint folder."""}  )
     A__      =			field(default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """If True the data is pretokenized."""}  )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Model name or path of model to be evaluated."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot-clean-valid"""      ,       metadata={"""help""": """Name or path of validation dataset."""}  )
     A__      =			field(default=2      ,       metadata={"""help""": """Batch size used for evaluation."""}  )
     A__      =			field(
         default=-1      ,       metadata={"""help""": """Maximum number of evaluation steps. If -1 the full dataset is evaluated."""}  )
     A__      =			field(default=1024      ,       metadata={"""help""": """Length of sequences to be evaluated."""}  )
     A__      =			field(default=1      ,       metadata={"""help""": """Random seed used for evaluation."""}  )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Model name or path of model to be evaluated."""}  )
     A__      =			field(default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Number of workers used for code evaluation."""}  )
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """The number of human-eval tasks to run. If not included all tasks are evaluated."""}      ,       )
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Sample from the language model's output distribution."""}  )
     A__      =			field(default=0.2      ,       metadata={"""help""": """Sampling temperature used for generation."""}  )
     A__      =			field(default=256      ,       metadata={"""help""": """Maximum number of newly generated tokens."""}  )
     A__      =			field(default=0      ,       metadata={"""help""": """Top-k parameter used for generation."""}  )
     A__      =			field(default=0.95      ,       metadata={"""help""": """Top-p parameter used for nucleus sampling."""}  )
     A__      =			field(default=10      ,       metadata={"""help""": """Number of generations to run in parallel."""}  )
     A__      =			field(
         default=200      ,       metadata={"""help""": """Number of completions to generate for each sample."""}  )
     A__      =			field(default=1      ,       metadata={"""help""": """Random seed used for evaluation."""}  )
     A__      =			field(
         default="""eval_results.json"""      ,       metadata={"""help""": """Random seed used for evaluation."""}  )
     A__      =			field(
         default="""0"""      ,       metadata={"""help""": """Allow `code_eval` to execute Python code on machine"""}  )
     A__      =			field(
         default=-1      ,       metadata={
             """help""": (
                 """Determine which device to run the `text-generation` Pipeline on. -1 is CPU and any zero or positive"""
                 """ number corresponds to which GPU device id to run on."""
             )
         }      ,       )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={
             """help""": """The number of CPU cores to use for parallel preprocessing. Default uses the maximum available."""
         }      ,       )
     A__      =			field(
         default="""transformersbook/codeparrot"""      ,       metadata={"""help""": """Folder or name of dataset to process."""}  )
     A__      =			field(
         default="""codeparrot-clean"""      ,       metadata={"""help""": """Folder to save processed processed dataset."""}  )
     A__      =			field(
         default=100000      ,       metadata={"""help""": """Number of files to save per JSON output file."""}  )
     A__      =			field(default="""content"""      ,       metadata={"""help""": """Column containing text data to process."""}  )
     A__      =			field(
         default=1000      ,       metadata={"""help""": """Maximum line length in file, otherwise file is filtered."""}  )
     A__      =			field(
         default=100      ,       metadata={"""help""": """Maximum mean line length in file, otherwise file is filtered."""}  )
     A__      =			field(
         default=0.25      ,       metadata={"""help""": """Maximum fraction of non-alphanumeric characters, otherwise file is filtered."""}  )
     A__      =			field(
         default=1.5      ,       metadata={"""help""": """Minimum character token ratio for the file, otherwise file is filtered."""}  )
     A__      =			field(
         default=0.7      ,       metadata={"""help""": """Probability for filtering config, test and uncommon files."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Name or path to the tokenizer."""}      ,       )
     A__      =			field(
         default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """If True, near-duplicate samples are removed."""}  )
     A__      =			field(
         default=0.85      ,       metadata={"""help""": """Jaccard threshold for near-duplicate samples."""}  )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""gpt2"""      ,       metadata={"""help""": """Base tokenizer to build new tokenizer from."""}  )
     A__      =			field(
         default="""transformersbook/codeparrot-train"""      ,       metadata={"""help""": """Dataset to train tokenizer on."""}  )
     A__      =			field(default="""content"""      ,       metadata={"""help""": """Column containing text data to process."""}  )
     A__      =			field(default=200000      ,       metadata={"""help""": """Number of examples to train tokenizer on."""}  )
     A__      =			field(
         default=32768      ,       metadata={"""help""": """Number of examples to train the tokenizer on."""}  )
     A__      =			field(default="""codeparrot"""      ,       metadata={"""help""": """Name of new tokenizer."""}  )
     A__      =			field(default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Push saved tokenizer to the hub."""}  )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Name or path to the tokenizer."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot-clean-train"""      ,       metadata={"""help""": """Name or path to the dataset to pretokenize."""}  )
     A__      =			field(
         default="""tokenized-codeparrot-train"""      ,       metadata={"""help""": """Repo name of the pretokenized data."""}  )
     A__      =			field(default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Number of workers used for code evaluation."""}  )
@dataclass
class 			UpperCamelCase      :
     A__      =			field(
         default="""gpt2-large"""      ,       metadata={"""help""": """Configuration to use for model initialization."""}  )
     A__      =			field(
         default="""codeparrot/codeparrot"""      ,       metadata={"""help""": """Tokenizer attached to model."""}  )
     A__      =			field(default="""codeparrot"""      ,       metadata={"""help""": """Name of the created model."""}  )
     A__      =			field(default=__SCREAMING_SNAKE_CASE      ,       metadata={"""help""": """Push saved tokenizer to the hub."""}  )



