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






def   lowerCAmelCase_  (						snake_case_						:     Union[str, Any]     )  ->List[Any]:
			lowerCamelCase__      :	Dict						=ASTConfig()

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

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

			lowerCamelCase__      :	List[str]						=json.load(open(hf_hub_download(__A				,       __A				,       repo_type='dataset'     )				,       'r'     )     )
			lowerCamelCase__      :	Tuple						={int(__A     ): v for k, v in idalabel.items()}
			lowerCamelCase__      :	Dict						=idalabel
			lowerCamelCase__      :	List[str]						={v: k for k, v in idalabel.items()}

			return config






def   lowerCAmelCase_  (						snake_case_						:     Dict     )  ->List[str]:




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

			return name






def   lowerCAmelCase_  (						snake_case_						:     str				,       snake_case_						:     List[str]     )  ->Tuple:




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

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

			return orig_state_dict






def   lowerCAmelCase_  (						snake_case_						:     Dict     )  ->List[str]:
			lowerCamelCase__      :	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(__A				,       __A     )






@torch.no_grad()
def   lowerCAmelCase_  (						snake_case_						:     str				,       snake_case_						:     Optional[Any]				,       snake_case_						:     Optional[Any]=False     )  ->Optional[Any]:
			lowerCamelCase__      :	Any						=get_audio_spectrogram_transformer_config(__A     )

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

			# load 🤗 model
			lowerCamelCase__      :	str						=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
			lowerCamelCase__      :	List[Any]						=-4.267_7393 if '''speech-commands''' not in model_name else -6.84_5978
			lowerCamelCase__      :	Union[str, Any]						=4.568_9974 if '''speech-commands''' not in model_name else 5.565_4526
			lowerCamelCase__      :	Optional[int]						=1_0_2_4 if '''speech-commands''' not in model_name else 1_2_8
			lowerCamelCase__      :	List[Any]						=ASTFeatureExtractor(mean=__A				,       std=__A				,       max_length=__A     )

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

						lowerCamelCase__      :	Optional[int]						=torchaudio.load(__A     )
						lowerCamelCase__      :	Any						=waveform.squeeze().numpy()

			lowerCamelCase__      :	Optional[int]						=feature_extractor(__A				,       sampling_rate=1_6_0_0_0				,       return_tensors='pt'     )

			# forward pass
			lowerCamelCase__      :	Optional[int]						=model(**__A     )
			lowerCamelCase__      :	List[Any]						=outputs.logits

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







'''simple docstring'''



from queue import PriorityQueue
from typing import Any

import numpy as np


def 				lowercase_					(     __A			:							dict       ,   __A			:							str       ,   __A			:							set       ,   __A			:							set       ,   __A			:							dict       ,   __A			:							dict       ,   __A			:							PriorityQueue       ,   __A			:							dict       ,   __A			:							float | int       ,   )							->   float | int:



 """simple docstring"""




 for nxt, d in graph[v]:
  if nxt in visited_forward:
   continue
  lowercase       : str						=cst_fwd.get(__A       ,   np.inf      )
  lowercase       : str						=cst_fwd[v] + d
  if new_cost_f < old_cost_f:
   queue.put((new_cost_f, nxt)      )
   lowercase       : int						=new_cost_f
   lowercase       : Optional[int]						=v
  if nxt in visited_backward:
   if cst_fwd[v] + d + cst_bwd[nxt] < shortest_distance:
    lowercase       : List[Any]						=cst_fwd[v] + d + cst_bwd[nxt]
 return shortest_distance


def 				lowercase_					(     __A			:							str       ,   __A			:							str       ,   __A			:							dict       ,   __A			:							dict      )							->   int:



 """simple docstring"""




 lowercase       : Optional[Any]						=-1

 lowercase       : Tuple						=set()
 lowercase       : str						=set()
 lowercase       : List[str]						={source: 0}
 lowercase       : List[str]						={destination: 0}
 lowercase       : Tuple						={source: None}
 lowercase       : List[Any]						={destination: None}
 lowercase       : PriorityQueue[Any]						=PriorityQueue()
 lowercase       : PriorityQueue[Any]						=PriorityQueue()

 lowercase       : Any						=np.inf

 queue_forward.put((0, source)      )
 queue_backward.put((0, destination)      )

 if source == destination:
  return 0

 while not queue_forward.empty() and not queue_backward.empty():
  lowercase	,		lowercase       : int						=queue_forward.get()
  visited_forward.add(__A      )

  lowercase	,		lowercase       : int						=queue_backward.get()
  visited_backward.add(__A      )

  lowercase       : Dict						=pass_and_relaxation(
      __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   )

  lowercase       : List[Any]						=pass_and_relaxation(
      __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   )

  if cst_fwd[v_fwd] + cst_bwd[v_bwd] >= shortest_distance:
   break

 if shortest_distance != np.inf:
  lowercase       : Union[str, Any]						=shortest_distance
 return shortest_path_distance


SCREAMING_SNAKE_CASE							=						{
    'B': [['C', 1]],
    'C': [['D', 1]],
    'D': [['F', 1]],
    'E': [['B', 1], ['G', 2]],
    'F': [],
    'G': [['F', 1]],
}
SCREAMING_SNAKE_CASE							=						{
    'B': [['E', 1]],
    'C': [['B', 1]],
    'D': [['C', 1]],
    'F': [['D', 1], ['G', 1]],
    'E': [[None, np.inf]],
    'G': [['E', 2]],
}

if __name__ == "__main__":
      import doctest

      doctest.testmod()



