






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


def        _a  (     UpperCamelCase_					:				int   ) ->   str:
						"""simple docstring"""







						lowerCAmelCase__       =						ASTConfig()

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

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

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

						return config


def        _a  (     UpperCamelCase_					:				Tuple   ) ->   str:
						"""simple docstring"""







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

						return name


def        _a  (     UpperCamelCase_					:				Any    ,   UpperCamelCase_					:				List[Any]   ) ->   Tuple:
						"""simple docstring"""







						for key in orig_state_dict.copy().keys():
												lowerCAmelCase__       =						orig_state_dict.pop(UpperCamelCase_   )

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

						return orig_state_dict


def        _a  (     UpperCamelCase_					:				Dict   ) ->   List[str]:
						"""simple docstring"""







						lowerCAmelCase__       =						[
						    "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(UpperCamelCase_    ,   UpperCamelCase_   )


@torch.no_grad()
def        _a  (     UpperCamelCase_					:				int    ,   UpperCamelCase_					:				List[str]    ,   UpperCamelCase_					:				Tuple=False   ) ->   Union[str, Any]:
						"""simple docstring"""







						lowerCAmelCase__       =						get_audio_spectrogram_transformer_config(UpperCamelCase_   )

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

						# load 🤗 model
						lowerCAmelCase__       =						ASTForAudioClassification(UpperCamelCase_   )
						model.eval()

						model.load_state_dict(UpperCamelCase_   )

						# verify outputs on dummy input
						# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
						lowerCAmelCase__       =						-4.2_677_393 if "speech-commands" not in model_name else -6.845_978
						lowerCAmelCase__       =						4.5_689_974 if "speech-commands" not in model_name else 5.5_654_526
						lowerCAmelCase__       =						1_024 if "speech-commands" not in model_name else 128
						lowerCAmelCase__       =						ASTFeatureExtractor(mean=UpperCamelCase_    ,   std=UpperCamelCase_    ,   max_length=UpperCamelCase_   )

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

												lowerCAmelCase__      ,							lowerCAmelCase__       =						torchaudio.load(UpperCamelCase_   )
												lowerCAmelCase__       =						waveform.squeeze().numpy()

						lowerCAmelCase__       =						feature_extractor(UpperCamelCase_    ,   sampling_rate=16_000    ,   return_tensors="pt"   )

						# forward pass
						lowerCAmelCase__       =						model(**UpperCamelCase_   )
						lowerCAmelCase__       =						outputs.logits

						if model_name == "ast-finetuned-audioset-10-10-0.4593":
												lowerCAmelCase__       =						torch.tensor([-0.8_760, -7.0_042, -8.6_602]   )
						elif model_name == "ast-finetuned-audioset-10-10-0.450":
												lowerCAmelCase__       =						torch.tensor([-1.1_986, -7.0_903, -8.2_718]   )
						elif model_name == "ast-finetuned-audioset-10-10-0.448":
												lowerCAmelCase__       =						torch.tensor([-2.6_128, -8.0_080, -9.4_344]   )
						elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
												lowerCAmelCase__       =						torch.tensor([-1.5_080, -7.4_534, -8.8_917]   )
						elif model_name == "ast-finetuned-audioset-12-12-0.447":
												lowerCAmelCase__       =						torch.tensor([-0.5_050, -6.5_833, -8.0_843]   )
						elif model_name == "ast-finetuned-audioset-14-14-0.443":
												lowerCAmelCase__       =						torch.tensor([-0.3_826, -7.0_336, -8.2_413]   )
						elif model_name == "ast-finetuned-audioset-16-16-0.442":
												lowerCAmelCase__       =						torch.tensor([-1.2_113, -6.9_101, -8.3_470]   )
						elif model_name == "ast-finetuned-speech-commands-v2":
												lowerCAmelCase__       =						torch.tensor([6.1_589, -8.0_566, -8.7_984]   )
						else:
												raise ValueError("Unknown model name"   )
						if not torch.allclose(logits[0, :3]    ,   UpperCamelCase_    ,   atol=1e-4   ):
												raise ValueError("Logits don't match"   )
						print("Looks ok!"   )

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

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

	a_							=    parser.parse_args()
	convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)














import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen, xsplitext
from ..table import array_cast
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
	from .features import FeatureType


a_,						a_,						a_							=    False, False, False


@dataclass
class      lowercase__     :
	a_         =None
	a_         =True
	a_         =True
	a_         =None
	# Automatically constructed
	a_         ="dict"
	a_         =pa.struct({"""bytes""": pa.binary(), """path""": pa.string()}	)
	a_         =field(default="""Audio""",					init=_UpperCAmelCase,					repr=_UpperCAmelCase	)

	def __call__(							self       )->					Optional[int]:







							'''simple docstring'''





							return self.pa_type

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					dict:







							'''simple docstring'''





							try:
													import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
							except ImportError as err:
													raise ImportError("To support encoding audio data, please install 'soundfile'."       ) from err
							if isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
													return {"bytes": None, "path": value}
							elif isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
													return {"bytes": value, "path": None}
							elif "array" in value:
													# convert the audio array to wav bytes
													lowerCAmelCase__       =						BytesIO()
													sf.write(__UpperCAmelCase		,    value["array"]		,    value["sampling_rate"]		,    format="wav"       )
													return {"bytes": buffer.getvalue(), "path": None}
							elif value.get("path"       ) is not None and os.path.isfile(value["path"]       ):
													# we set "bytes": None to not duplicate the data if they're already available locally
													if value["path"].endswith("pcm"       ):
																			# "PCM" only has raw audio bytes
																			if value.get("sampling_rate"       ) is None:
																									# At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
																									raise KeyError("To use PCM files, please specify a 'sampling_rate' in Audio object"       )
																			if value.get("bytes"       ):
																									# If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
																									lowerCAmelCase__       =						np.frombuffer(value["bytes"]		,    dtype=np.intaa       ).astype(np.floataa       ) / 32767
																			else:
																									lowerCAmelCase__       =						np.memmap(value["path"]		,    dtype="h"		,    mode="r"       ).astype(np.floataa       ) / 32767

																			lowerCAmelCase__       =						BytesIO(bytes()       )
																			sf.write(__UpperCAmelCase		,    __UpperCAmelCase		,    value["sampling_rate"]		,    format="wav"       )
																			return {"bytes": buffer.getvalue(), "path": None}
													else:
																			return {"bytes": None, "path": value.get("path"       )}
							elif value.get("bytes"       ) is not None or value.get("path"       ) is not None:
													# store the audio bytes, and path is used to infer the audio format using the file extension
													return {"bytes": value.get("bytes"       ), "path": value.get("path"       )}
							else:
													raise ValueError(
													    F"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase = None       )->					dict:







							'''simple docstring'''





							if not self.decode:
													raise RuntimeError("Decoding is disabled for this feature. Please use Audio(decode=True) instead."       )

							lowerCAmelCase__      ,							lowerCAmelCase__       =						(value["path"], BytesIO(value["bytes"]       )) if value["bytes"] is not None else (value["path"], None)
							if path is None and file is None:
													raise ValueError(F"An audio sample should have one of 'path' or 'bytes' but both are None in {value}."       )

							try:
													import librosa
													import soundfile as sf
							except ImportError as err:
													raise ImportError("To support decoding audio files, please install 'librosa' and 'soundfile'."       ) from err

							lowerCAmelCase__       =						xsplitext(__UpperCAmelCase       )[1][1:].lower() if path is not None else None
							if not config.IS_OPUS_SUPPORTED and audio_format == "opus":
													raise RuntimeError(
													    "Decoding 'opus' files requires system library 'libsndfile'>=1.0.31, "
													    "You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. "       )
							elif not config.IS_MP3_SUPPORTED and audio_format == "mp3":
													raise RuntimeError(
													    "Decoding 'mp3' files requires system library 'libsndfile'>=1.1.0, "
													    "You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. "       )

							if file is None:
													lowerCAmelCase__       =						token_per_repo_id or {}
													lowerCAmelCase__       =						path.split("::"       )[-1]
													try:
																			lowerCAmelCase__       =						string_to_dict(__UpperCAmelCase		,    config.HUB_DATASETS_URL       )["repo_id"]
																			lowerCAmelCase__       =						token_per_repo_id[repo_id]
													except (ValueError, KeyError):
																			lowerCAmelCase__       =						None

													with xopen(__UpperCAmelCase		,    "rb"		,    use_auth_token=__UpperCAmelCase       ) as f:
																			lowerCAmelCase__      ,							lowerCAmelCase__       =						sf.read(__UpperCAmelCase       )

							else:
													lowerCAmelCase__      ,							lowerCAmelCase__       =						sf.read(__UpperCAmelCase       )

							lowerCAmelCase__       =						array.T
							if self.mono:
													lowerCAmelCase__       =						librosa.to_mono(__UpperCAmelCase       )
							if self.sampling_rate and self.sampling_rate != sampling_rate:
													lowerCAmelCase__       =						librosa.resample(__UpperCAmelCase		,    orig_sr=__UpperCAmelCase		,    target_sr=self.sampling_rate       )
													lowerCAmelCase__       =						self.sampling_rate

							return {"path": path, "array": array, "sampling_rate": sampling_rate}

	def 		UpperCAmelCase  (							self       )->					Union["FeatureType", Dict[str, "FeatureType"]]:







							'''simple docstring'''





							from .features import Value

							if self.decode:
													raise ValueError("Cannot flatten a decoded Audio feature."       )
							return {
							    "bytes": Value("binary"       ),
							    "path": Value("string"       ),
							}

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					pa.StructArray:







							'''simple docstring'''





							if pa.types.is_string(storage.type       ):
													lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.binary()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, storage]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							elif pa.types.is_binary(storage.type       ):
													lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.string()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([storage, path_array]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							elif pa.types.is_struct(storage.type       ) and storage.type.get_all_field_indices("array"       ):
													lowerCAmelCase__       =						pa.array([Audio().encode_example(__UpperCAmelCase       ) if x is not None else None for x in storage.to_pylist()]       )
							elif pa.types.is_struct(storage.type       ):
													if storage.type.get_field_index("bytes"       ) >= 0:
																			lowerCAmelCase__       =						storage.field("bytes"       )
													else:
																			lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.binary()       )
													if storage.type.get_field_index("path"       ) >= 0:
																			lowerCAmelCase__       =						storage.field("path"       )
													else:
																			lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.string()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, path_array]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							return array_cast(__UpperCAmelCase		,    self.pa_type       )


	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					pa.StructArray:







							'''simple docstring'''





							@no_op_if_value_is_null
							def path_to_bytes(__UpperCAmelCase       ):
													with xopen(__UpperCAmelCase		,    "rb"       ) as f:
																			lowerCAmelCase__       =						f.read()
													return bytes_

							lowerCAmelCase__       =						pa.array(
							    [
							        (path_to_bytes(x["path"]       ) if x["bytes"] is None else x["bytes"]) if x is not None else None
							        for x in storage.to_pylist()
							    ]		,    type=pa.binary()		,    )
							lowerCAmelCase__       =						pa.array(
							    [os.path.basename(__UpperCAmelCase       ) if path is not None else None for path in storage.field("path"       ).to_pylist()]		,    type=pa.string()		,    )
							lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, path_array]		,    ["bytes", "path"]		,    mask=bytes_array.is_null()       )
							return array_cast(__UpperCAmelCase		,    self.pa_type       )






