




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

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				Optional[Any]	   =		ASTConfig()

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

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

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

	return config

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






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

	return name

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	for key in orig_state_dict.copy().keys():
		A_  :				Optional[int]	   =		orig_state_dict.pop(_UpperCAmelCase	)

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

	return orig_state_dict

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				List[str]	   =		[
	    '''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 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase=False	):




	"""simple docstring"""






	A_  :				Tuple	   =		get_audio_spectrogram_transformer_config(_UpperCAmelCase	)

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

	# load 🤗 model
	A_  :				Optional[Any]	   =		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
	A_  :				Optional[Any]	   =		-4.267_7393 if '''speech-commands''' not in model_name else -6.84_5978
	A_  :				Dict	   =		4.568_9974 if '''speech-commands''' not in model_name else 5.565_4526
	A_  :				int	   =		1024 if '''speech-commands''' not in model_name else 128
	A_  :				List[Any]	   =		ASTFeatureExtractor(mean=_UpperCAmelCase  ,      std=_UpperCAmelCase  ,      max_length=_UpperCAmelCase	)

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

		A_   , A_  :				Optional[int]	   =		torchaudio.load(_UpperCAmelCase	)
		A_  :				int	   =		waveform.squeeze().numpy()

	A_  :				Any	   =		feature_extractor(_UpperCAmelCase  ,      sampling_rate=1_6000  ,      return_tensors='''pt'''	)

	# forward pass
	A_  :				Optional[Any]	   =		model(**_UpperCAmelCase	)
	A_  :				Optional[int]	   =		outputs.logits

	if model_name == "ast-finetuned-audioset-10-10-0.4593":
		A_  :				Optional[Any]	   =		torch.tensor([-0.8760, -7.0042, -8.6602]	)
	elif model_name == "ast-finetuned-audioset-10-10-0.450":
		A_  :				Optional[Any]	   =		torch.tensor([-1.1986, -7.0903, -8.2718]	)
	elif model_name == "ast-finetuned-audioset-10-10-0.448":
		A_  :				Tuple	   =		torch.tensor([-2.6128, -8.0080, -9.4344]	)
	elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
		A_  :				Tuple	   =		torch.tensor([-1.5080, -7.4534, -8.8917]	)
	elif model_name == "ast-finetuned-audioset-12-12-0.447":
		A_  :				Union[str, Any]	   =		torch.tensor([-0.5050, -6.5833, -8.0843]	)
	elif model_name == "ast-finetuned-audioset-14-14-0.443":
		A_  :				Tuple	   =		torch.tensor([-0.3826, -7.0336, -8.2413]	)
	elif model_name == "ast-finetuned-audioset-16-16-0.442":
		A_  :				Union[str, Any]	   =		torch.tensor([-1.2113, -6.9101, -8.3470]	)
	elif model_name == "ast-finetuned-speech-commands-v2":
		A_  :				List[str]	   =		torch.tensor([6.1589, -8.0566, -8.7984]	)
	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__":
		_lowerCamelCase  :			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.'
		)

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











"""simple docstring"""

from collections.abc import Generator
from math import sin

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	if len(_UpperCAmelCase	) != 32:
		raise ValueError('''Input must be of length 32'''	)

	A_  :				int	   =		b''''''
	for i in [3, 2, 1, 0]:
		little_endian += string_aa[8 * i : 8 * i + 8]
	return little_endian

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	if i < 0:
		raise ValueError('''Input must be non-negative'''	)

	A_  :				int	   =		format(_UpperCAmelCase  ,      '''08x'''	)[-8:]
	A_  :				Optional[Any]	   =		b''''''
	for i in [3, 2, 1, 0]:
		little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode('''utf-8'''	)
	return little_endian_hex

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				int	   =		b''''''
	for char in message:
		bit_string += format(_UpperCAmelCase  ,      '''08b'''	).encode('''utf-8'''	)
	A_  :				Optional[Any]	   =		format(len(_UpperCAmelCase	)  ,      '''064b'''	).encode('''utf-8'''	)

	# Pad bit_string to a multiple of 512 chars
	bit_string += b"1"
	while len(_UpperCAmelCase	) % 512 != 448:
		bit_string += b"0"
	bit_string += to_little_endian(start_len[32:]	) + to_little_endian(start_len[:32]	)

	return bit_string

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	if len(_UpperCAmelCase	) % 512 != 0:
		raise ValueError('''Input must have length that\'s a multiple of 512'''	)

	for pos in range(0  ,      len(_UpperCAmelCase	)  ,      512	):
		A_  :				Union[str, Any]	   =		bit_string[pos : pos + 512]
		A_  :				Optional[int]	   =		[]
		for i in range(0  ,      512  ,      32	):
			block_words.append(int(to_little_endian(block[i : i + 32]	)  ,      2	)	)
		yield block_words

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	if i < 0:
		raise ValueError('''Input must be non-negative'''	)

	A_  :				str	   =		format(_UpperCAmelCase  ,      '''032b'''	)
	A_  :				List[str]	   =		''''''
	for c in i_str:
		new_str += "1" if c == "0" else "0"
	return int(_UpperCAmelCase  ,      2	)

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	return (a + b) % 2**32

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	if i < 0:
		raise ValueError('''Input must be non-negative'''	)
	if shift < 0:
		raise ValueError('''Shift must be non-negative'''	)
	return ((i << shift) ^ (i >> (32 - shift))) % 2**32

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				Any	   =		preprocess(_UpperCAmelCase	)

	A_  :				Any	   =		[int(2**32 * abs(sin(i + 1	)	)	) for i in range(64	)]

	# Starting states
	A_  :				List[str]	   =		0x6_7_4_5_2_3_0_1
	A_  :				str	   =		0xe_f_c_d_a_b_8_9
	A_  :				Union[str, Any]	   =		0x9_8_b_a_d_c_f_e
	A_  :				Dict	   =		0x1_0_3_2_5_4_7_6

	A_  :				int	   =		[
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	]

	# Process bit string in chunks, each with 16 32-char words
	for block_words in get_block_words(_UpperCAmelCase	):
		A_  :				str	   =		aa
		A_  :				Tuple	   =		ba
		A_  :				Tuple	   =		ca
		A_  :				int	   =		da

		# Hash current chunk
		for i in range(64	):
			if i <= 15:
				# f = (b & c) | (not_32(b) & d)     # Alternate definition for f
				A_  :				str	   =		d ^ (b & (c ^ d))
				A_  :				int	   =		i
			elif i <= 31:
				# f = (d & b) | (not_32(d) & c)     # Alternate definition for f
				A_  :				List[Any]	   =		c ^ (d & (b ^ c))
				A_  :				Optional[int]	   =		(5 * i + 1) % 16
			elif i <= 47:
				A_  :				Optional[Any]	   =		b ^ c ^ d
				A_  :				List[Any]	   =		(3 * i + 5) % 16
			else:
				A_  :				Dict	   =		c ^ (b | not_aa(_UpperCAmelCase	))
				A_  :				Dict	   =		(7 * i) % 16
			A_  :				Optional[int]	   =		(f + a + added_consts[i] + block_words[g]) % 2**32
			A_  :				List[str]	   =		d
			A_  :				str	   =		c
			A_  :				Tuple	   =		b
			A_  :				Tuple	   =		sum_aa(_UpperCAmelCase  ,      left_rotate_aa(_UpperCAmelCase  ,      shift_amounts[i]	)	)

		# Add hashed chunk to running total
		A_  :				str	   =		sum_aa(_UpperCAmelCase  ,      _UpperCAmelCase	)
		A_  :				int	   =		sum_aa(_UpperCAmelCase  ,      _UpperCAmelCase	)
		A_  :				int	   =		sum_aa(_UpperCAmelCase  ,      _UpperCAmelCase	)
		A_  :				Optional[Any]	   =		sum_aa(_UpperCAmelCase  ,      _UpperCAmelCase	)

	A_  :				Dict	   =		reformat_hex(_UpperCAmelCase	) + reformat_hex(_UpperCAmelCase	) + reformat_hex(_UpperCAmelCase	) + reformat_hex(_UpperCAmelCase	)
	return digest


if __name__ == "__main__":
		import doctest

		doctest.testmod()





