

from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


__A			:					List[str]						=			logging.get_logger(__name__)

__A			:					List[str]						=			{
    'facebook/blenderbot_small-90M': 'https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json',
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   lowerCamelCase__				=      "blenderbot-small"
   lowerCamelCase__				=      ["past_key_values"]
   lowerCamelCase__				=      {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}





   def __init__(       self					:			List[Any]     , __lowerCamelCase					:			List[str]=50265     , __lowerCamelCase					:			List[Any]=512     , __lowerCamelCase					:			Optional[Any]=8     , __lowerCamelCase					:			str=2048     , __lowerCamelCase					:			str=16     , __lowerCamelCase					:			Optional[int]=8     , __lowerCamelCase					:			Optional[int]=2048     , __lowerCamelCase					:			str=16     , __lowerCamelCase					:			int=0.0     , __lowerCamelCase					:			Any=0.0     , __lowerCamelCase					:			Tuple=True     , __lowerCamelCase					:			Optional[int]=True     , __lowerCamelCase					:			int="gelu"     , __lowerCamelCase					:			str=512     , __lowerCamelCase					:			List[Any]=0.1     , __lowerCamelCase					:			List[str]=0.0     , __lowerCamelCase					:			Optional[Any]=0.0     , __lowerCamelCase					:			Union[str, Any]=0.02     , __lowerCamelCase					:			Union[str, Any]=1     , __lowerCamelCase					:			List[Any]=False     , __lowerCamelCase					:			int=0     , __lowerCamelCase					:			Optional[Any]=1     , __lowerCamelCase					:			List[str]=2     , __lowerCamelCase					:			Optional[Any]=2     , **__lowerCamelCase					:			Optional[int]     , ):
       SCREAMING_SNAKE_CASE										=   vocab_size
       SCREAMING_SNAKE_CASE										=   max_position_embeddings
       SCREAMING_SNAKE_CASE										=   d_model
       SCREAMING_SNAKE_CASE										=   encoder_ffn_dim
       SCREAMING_SNAKE_CASE										=   encoder_layers
       SCREAMING_SNAKE_CASE										=   encoder_attention_heads
       SCREAMING_SNAKE_CASE										=   decoder_ffn_dim
       SCREAMING_SNAKE_CASE										=   decoder_layers
       SCREAMING_SNAKE_CASE										=   decoder_attention_heads
       SCREAMING_SNAKE_CASE										=   dropout
       SCREAMING_SNAKE_CASE										=   attention_dropout
       SCREAMING_SNAKE_CASE										=   activation_dropout
       SCREAMING_SNAKE_CASE										=   activation_function
       SCREAMING_SNAKE_CASE										=   init_std
       SCREAMING_SNAKE_CASE										=   encoder_layerdrop
       SCREAMING_SNAKE_CASE										=   decoder_layerdrop
       SCREAMING_SNAKE_CASE										=   use_cache
       SCREAMING_SNAKE_CASE										=   encoder_layers
       SCREAMING_SNAKE_CASE										=   scale_embedding  # scale factor will be sqrt(d_model) if True

       super().__init__(
           pad_token_id=__lowerCamelCase     , bos_token_id=__lowerCamelCase     , eos_token_id=__lowerCamelCase     , is_encoder_decoder=__lowerCamelCase     , decoder_start_token_id=__lowerCamelCase     , forced_eos_token_id=__lowerCamelCase     , **__lowerCamelCase     , )





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   @property
   def      _snake_case		(       self					:			Any			):
       if self.task in ["default", "seq2seq-lm"]:
           SCREAMING_SNAKE_CASE										=   OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
               ]			)

           if self.use_past:
               SCREAMING_SNAKE_CASE										=   {0: "batch"}
               SCREAMING_SNAKE_CASE										=   {0: "batch", 1: "past_decoder_sequence + sequence"}
           else:
               SCREAMING_SNAKE_CASE										=   {0: "batch", 1: "decoder_sequence"}
               SCREAMING_SNAKE_CASE										=   {0: "batch", 1: "decoder_sequence"}

           if self.use_past:
               self.fill_with_past_key_values_(__lowerCamelCase     , direction="inputs"			)
       elif self.task == "causal-lm":
           # TODO: figure this case out.
           SCREAMING_SNAKE_CASE										=   OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
               ]			)
           if self.use_past:
               SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_layers
               for i in range(__lowerCamelCase			):
                   SCREAMING_SNAKE_CASE										=   {0: "batch", 2: "past_sequence + sequence"}
                   SCREAMING_SNAKE_CASE										=   {0: "batch", 2: "past_sequence + sequence"}
       else:
           SCREAMING_SNAKE_CASE										=   OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                   ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                   ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
               ]			)

       return common_inputs





   @property
   def      _snake_case		(       self					:			Optional[int]			):
       if self.task in ["default", "seq2seq-lm"]:
           SCREAMING_SNAKE_CASE										=   super().outputs
       else:
           SCREAMING_SNAKE_CASE										=   super(__lowerCamelCase     , self			).outputs
           if self.use_past:
               SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_layers
               for i in range(__lowerCamelCase			):
                   SCREAMING_SNAKE_CASE										=   {0: "batch", 2: "past_sequence + sequence"}
                   SCREAMING_SNAKE_CASE										=   {0: "batch", 2: "past_sequence + sequence"}
       return common_outputs





   def      _snake_case		(       self					:			Any     , __lowerCamelCase					:			PreTrainedTokenizer     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			bool = False     , __lowerCamelCase					:			Optional[TensorType] = None     , ):
       SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)

       # Generate decoder inputs
       SCREAMING_SNAKE_CASE										=   seq_length if not self.use_past else 1
       SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
       SCREAMING_SNAKE_CASE										=   dict(**__lowerCamelCase     , **__lowerCamelCase			)

       if self.use_past:
           if not is_torch_available():
               raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."			)
           else:
               import torch
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   common_inputs["input_ids"].shape
           SCREAMING_SNAKE_CASE										=   common_inputs["decoder_input_ids"].shape[1]
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_attention_heads
           SCREAMING_SNAKE_CASE										=   (
               batch,
               num_encoder_attention_heads,
               encoder_seq_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )
           SCREAMING_SNAKE_CASE										=   decoder_seq_length + 3
           SCREAMING_SNAKE_CASE										=   (
               batch,
               num_decoder_attention_heads,
               decoder_past_length,
               self._config.hidden_size // num_decoder_attention_heads,
           )

           SCREAMING_SNAKE_CASE										=   torch.cat(
               [common_inputs["decoder_attention_mask"], torch.ones(__lowerCamelCase     , __lowerCamelCase			)]     , dim=1			)

           SCREAMING_SNAKE_CASE										=   []
           # If the number of encoder and decoder layers are present in the model configuration, both are considered
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_layers
           SCREAMING_SNAKE_CASE										=   min(__lowerCamelCase     , __lowerCamelCase			)
           SCREAMING_SNAKE_CASE										=   max(__lowerCamelCase     , __lowerCamelCase			) - min_num_layers
           SCREAMING_SNAKE_CASE										=   "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

           for _ in range(__lowerCamelCase			):
               common_inputs["past_key_values"].append(
                   (
                       torch.zeros(__lowerCamelCase			),
                       torch.zeros(__lowerCamelCase			),
                       torch.zeros(__lowerCamelCase			),
                       torch.zeros(__lowerCamelCase			),
                   )			)
           # TODO: test this.
           SCREAMING_SNAKE_CASE										=   encoder_shape if remaining_side_name == "encoder" else decoder_shape
           for _ in range(__lowerCamelCase     , __lowerCamelCase			):
               common_inputs["past_key_values"].append((torch.zeros(__lowerCamelCase			), torch.zeros(__lowerCamelCase			))			)
       return common_inputs





   def      _snake_case		(       self					:			Optional[Any]     , __lowerCamelCase					:			PreTrainedTokenizer     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			bool = False     , __lowerCamelCase					:			Optional[TensorType] = None     , ):
       SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)

       if self.use_past:
           if not is_torch_available():
               raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."			)
           else:
               import torch
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   common_inputs["input_ids"].shape
           # Not using the same length for past_key_values
           SCREAMING_SNAKE_CASE										=   seqlen + 2
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_layers
           SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE										=   self.num_attention_heads
           SCREAMING_SNAKE_CASE										=   (
               batch,
               num_encoder_attention_heads,
               past_key_values_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )

           SCREAMING_SNAKE_CASE										=   common_inputs["attention_mask"].dtype
           SCREAMING_SNAKE_CASE										=   torch.cat(
               [common_inputs["attention_mask"], torch.ones(__lowerCamelCase     , __lowerCamelCase     , dtype=__lowerCamelCase			)]     , dim=1			)
           SCREAMING_SNAKE_CASE										=   [
               (torch.zeros(__lowerCamelCase			), torch.zeros(__lowerCamelCase			)) for _ in range(__lowerCamelCase			)
           ]
       return common_inputs





   def      _snake_case		(       self					:			Tuple     , __lowerCamelCase					:			PreTrainedTokenizer     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			bool = False     , __lowerCamelCase					:			Optional[TensorType] = None     , ):
       # Copied from OnnxConfig.generate_dummy_inputs
       # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
       # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
       SCREAMING_SNAKE_CASE										=   compute_effective_axis_dimension(
           __lowerCamelCase     , fixed_dimension=OnnxConfig.default_fixed_batch     , num_token_to_add=0			)

       # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
       SCREAMING_SNAKE_CASE										=   tokenizer.num_special_tokens_to_add(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   compute_effective_axis_dimension(
           __lowerCamelCase     , fixed_dimension=OnnxConfig.default_fixed_sequence     , num_token_to_add=__lowerCamelCase			)

       # Generate dummy inputs according to compute batch and sequence
       SCREAMING_SNAKE_CASE										=   [" ".join([tokenizer.unk_token]			) * seq_length] * batch_size
       SCREAMING_SNAKE_CASE										=   dict(tokenizer(__lowerCamelCase     , return_tensors=__lowerCamelCase			)			)
       return common_inputs





   def      _snake_case		(       self					:			Optional[Any]     , __lowerCamelCase					:			PreTrainedTokenizer     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			int = -1     , __lowerCamelCase					:			bool = False     , __lowerCamelCase					:			Optional[TensorType] = None     , ):
       if self.task in ["default", "seq2seq-lm"]:
           SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
               __lowerCamelCase     , batch_size=__lowerCamelCase     , seq_length=__lowerCamelCase     , is_pair=__lowerCamelCase     , framework=__lowerCamelCase			)

       elif self.task == "causal-lm":
           SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_causal_lm(
               __lowerCamelCase     , batch_size=__lowerCamelCase     , seq_length=__lowerCamelCase     , is_pair=__lowerCamelCase     , framework=__lowerCamelCase			)
       else:
           SCREAMING_SNAKE_CASE										=   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               __lowerCamelCase     , batch_size=__lowerCamelCase     , seq_length=__lowerCamelCase     , is_pair=__lowerCamelCase     , framework=__lowerCamelCase			)

       return common_inputs





   def      _snake_case		(       self					:			Optional[int]     , __lowerCamelCase					:			Union[str, Any]     , __lowerCamelCase					:			Union[str, Any]     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			int			):
       if self.task in ["default", "seq2seq-lm"]:
           SCREAMING_SNAKE_CASE										=   super()._flatten_past_key_values_(__lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)
       else:
           SCREAMING_SNAKE_CASE										=   super(__lowerCamelCase     , self			)._flatten_past_key_values_(
               __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase     , __lowerCamelCase			)

"""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_       :    Any						      =			logging.get_logger(__name__)
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    ASTConfig()

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

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

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

					return config
def 				A_  (__a					):

					'''simple docstring'''





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

					return name
def 				A_  (__a					, __a					):

					'''simple docstring'''





					for key in orig_state_dict.copy().keys():
										A_								=    orig_state_dict.pop(__a					)

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

					return orig_state_dict
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    [
					    "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(__a					, __a					)
@torch.no_grad()
def 				A_  (__a					, __a					, __a=False					):

					'''simple docstring'''





					A_								=    get_audio_spectrogram_transformer_config(__a					)

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

					# load 🤗 model
					A_								=    ASTForAudioClassification(__a					)
					model.eval()

					model.load_state_dict(__a					)

					# verify outputs on dummy input
					# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
					A_								=    -4.2677393 if "speech-commands" not in model_name else -6.845978
					A_								=    4.5689974 if "speech-commands" not in model_name else 5.5654526
					A_								=    1024 if "speech-commands" not in model_name else 128
					A_								=    ASTFeatureExtractor(mean=__a					, std=__a					, max_length=__a					)

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

										A_			,     A_								=    torchaudio.load(__a					)
										A_								=    waveform.squeeze().numpy()

					A_								=    feature_extractor(__a					, sampling_rate=1_6000					, return_tensors="pt"					)

					# forward pass
					A_								=    model(**__a					)
					A_								=    outputs.logits

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

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

					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_       :    Dict						      =			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_       :    List[str]						      =			parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




