


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






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: List[str] ):






			"""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(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,repo_type='dataset' ),'r' ) )
			__A=	{int(_SCREAMING_SNAKE_CASE ): v for k, v in idalabel.items()}
			__A=	idalabel
			__A=	{v: k for k, v in idalabel.items()}

			return config






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[Any] ):






			"""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 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[Any],_SCREAMING_SNAKE_CASE					: Optional[int] ):






			"""simple docstring"""




			for key in orig_state_dict.copy().keys():
						__A=	orig_state_dict.pop(_SCREAMING_SNAKE_CASE )

						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 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Union[str, Any] ):






			"""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(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE )






@torch.no_grad()
def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: List[Any],_SCREAMING_SNAKE_CASE					: Tuple,_SCREAMING_SNAKE_CASE					: int=False ):






			"""simple docstring"""




			__A=	get_audio_spectrogram_transformer_config(_SCREAMING_SNAKE_CASE )

			__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(_SCREAMING_SNAKE_CASE,map_location='cpu' )
			# remove some keys
			remove_keys(_SCREAMING_SNAKE_CASE )
			# rename some keys
			__A=	convert_state_dict(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE )

			# load 🤗 model
			__A=	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
			__A=	-4.2_67_73_93 if 'speech-commands' not in model_name else -6.84_59_78
			__A=	4.5_68_99_74 if 'speech-commands' not in model_name else 5.5_65_45_26
			__A=	1024 if 'speech-commands' not in model_name else 128
			__A=	ASTFeatureExtractor(mean=_SCREAMING_SNAKE_CASE,std=_SCREAMING_SNAKE_CASE,max_length=_SCREAMING_SNAKE_CASE )

			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(_SCREAMING_SNAKE_CASE )
						__A=	waveform.squeeze().numpy()

			__A=	feature_extractor(_SCREAMING_SNAKE_CASE,sampling_rate=1_6000,return_tensors='pt' )

			# forward pass
			__A=	model(**_SCREAMING_SNAKE_CASE )
			__A=	outputs.logits

			if model_name == "ast-finetuned-audioset-10-10-0.4593":
						__A=	torch.tensor([-0.87_60, -7.00_42, -8.66_02] )
			elif model_name == "ast-finetuned-audioset-10-10-0.450":
						__A=	torch.tensor([-1.19_86, -7.09_03, -8.27_18] )
			elif model_name == "ast-finetuned-audioset-10-10-0.448":
						__A=	torch.tensor([-2.61_28, -8.00_80, -9.43_44] )
			elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
						__A=	torch.tensor([-1.50_80, -7.45_34, -8.89_17] )
			elif model_name == "ast-finetuned-audioset-12-12-0.447":
						__A=	torch.tensor([-0.50_50, -6.58_33, -8.08_43] )
			elif model_name == "ast-finetuned-audioset-14-14-0.443":
						__A=	torch.tensor([-0.38_26, -7.03_36, -8.24_13] )
			elif model_name == "ast-finetuned-audioset-16-16-0.442":
						__A=	torch.tensor([-1.21_13, -6.91_01, -8.34_70] )
			elif model_name == "ast-finetuned-speech-commands-v2":
						__A=	torch.tensor([6.15_89, -8.05_66, -8.79_84] )
			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__":
					UpperCAmelCase__  			=				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__  			=				parser.parse_args()
					convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







'''simple docstring'''





import numpy as np

from transformers import Pipeline






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Any ):






			"""simple docstring"""




			__A=	np.max(_SCREAMING_SNAKE_CASE,axis=-1,keepdims=_SCREAMING_SNAKE_CASE )
			__A=	np.exp(outputs - maxes )
			return shifted_exp / shifted_exp.sum(axis=-1,keepdims=_SCREAMING_SNAKE_CASE )






class 						a__					(					a_   ):





			'''simple docstring'''

			def 					lowerCAmelCase     (  self							:			Optional[Any]			,    **lowerCAmelCase_							:			List[str]	)      ->		List[str]:
						__A=	{}
						if "second_text" in kwargs:
									__A=	kwargs['second_text']
						return preprocess_kwargs, {}, {}
			def 					lowerCAmelCase     (  self							:			str			,    lowerCAmelCase_							:			Union[str, Any]			,    lowerCAmelCase_							:			Optional[int]=None	)      ->		int:
						return self.tokenizer(lowerCAmelCase_			,    text_pair=lowerCAmelCase_			,    return_tensors=self.framework	)
			def 					lowerCAmelCase     (  self							:			Dict			,    lowerCAmelCase_							:			List[str]	)      ->		Tuple:
						return self.model(**lowerCAmelCase_	)

			def 					lowerCAmelCase     (  self							:			str			,    lowerCAmelCase_							:			Dict	)      ->		Optional[int]:
						__A=	model_outputs.logits[0].numpy()
						__A=	softmax(lowerCAmelCase_	)

						__A=	np.argmax(lowerCAmelCase_	)
						__A=	self.model.config.idalabel[best_class]
						__A=	probabilities[best_class].item()
						__A=	logits.tolist()
						return {"label": label, "score": score, "logits": logits}



