



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    :		List[Any]		      =    logging.get_logger(__name__)

def        __a			(   __lowerCAmelCase   )     ->     Optional[int]:
					SCREAMING_SNAKE_CASE							:     Any						=				ASTConfig()

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

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

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

					return config

def        __a			(   __lowerCAmelCase   )     ->     Any:





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

					return name

def        __a			(   __lowerCAmelCase      ,	__lowerCAmelCase   )     ->     List[Any]:





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

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

					return orig_state_dict

def        __a			(   __lowerCAmelCase   )     ->     Union[str, Any]:
					SCREAMING_SNAKE_CASE							:     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        __a			(   __lowerCAmelCase      ,	__lowerCAmelCase      ,	__lowerCAmelCase=False   )     ->     List[Any]:
					SCREAMING_SNAKE_CASE							:     List[str]						=				get_audio_spectrogram_transformer_config(A_   )

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

					# load 🤗 model
					SCREAMING_SNAKE_CASE							:     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
					SCREAMING_SNAKE_CASE							:     Dict						=				-4.2_677_393 if 'speech-commands' not in model_name else -6.845_978
					SCREAMING_SNAKE_CASE							:     Optional[Any]						=				4.5_689_974 if 'speech-commands' not in model_name else 5.5_654_526
					SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				1024 if 'speech-commands' not in model_name else 128
					SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				ASTFeatureExtractor(mean=A_      ,	std=A_      ,	max_length=A_   )

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

										SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     int						=				torchaudio.load(A_   )
										SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				waveform.squeeze().numpy()

					SCREAMING_SNAKE_CASE							:     Dict						=				feature_extractor(A_      ,	sampling_rate=1_6000      ,	return_tensors='pt'   )

					# forward pass
					SCREAMING_SNAKE_CASE							:     Dict						=				model(**A_   )
					SCREAMING_SNAKE_CASE							:     Any						=				outputs.logits

					if model_name == "ast-finetuned-audioset-10-10-0.4593":
										SCREAMING_SNAKE_CASE							:     Tuple						=				torch.tensor([-0.8_760, -7.0_042, -8.6_602]   )
					elif model_name == "ast-finetuned-audioset-10-10-0.450":
										SCREAMING_SNAKE_CASE							:     Optional[int]						=				torch.tensor([-1.1_986, -7.0_903, -8.2_718]   )
					elif model_name == "ast-finetuned-audioset-10-10-0.448":
										SCREAMING_SNAKE_CASE							:     Optional[Any]						=				torch.tensor([-2.6_128, -8.0_080, -9.4_344]   )
					elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
										SCREAMING_SNAKE_CASE							:     List[str]						=				torch.tensor([-1.5_080, -7.4_534, -8.8_917]   )
					elif model_name == "ast-finetuned-audioset-12-12-0.447":
										SCREAMING_SNAKE_CASE							:     str						=				torch.tensor([-0.5_050, -6.5_833, -8.0_843]   )
					elif model_name == "ast-finetuned-audioset-14-14-0.443":
										SCREAMING_SNAKE_CASE							:     List[Any]						=				torch.tensor([-0.3_826, -7.0_336, -8.2_413]   )
					elif model_name == "ast-finetuned-audioset-16-16-0.442":
										SCREAMING_SNAKE_CASE							:     List[str]						=				torch.tensor([-1.2_113, -6.9_101, -8.3_470]   )
					elif model_name == "ast-finetuned-speech-commands-v2":
										SCREAMING_SNAKE_CASE							:     str						=				torch.tensor([6.1_589, -8.0_566, -8.7_984]   )
					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    :		Optional[int]		      =    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    :		Dict		      =    parser.parse_args()
			convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCAmelCase     :	Dict       			=  logging.get_logger(__name__)

__lowerCAmelCase     :	int       			=  {
    'microsoft/trocr-base-handwritten': (
        'https://huggingface.co/microsoft/trocr-base-handwritten/resolve/main/config.json'
    ),
    # See all TrOCR models at https://huggingface.co/models?filter=trocr
}




class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  a__       			=						"""trocr"""
  a__       			=						["""past_key_values"""]
  a__       			=						{
      """num_attention_heads""": """decoder_attention_heads""",
      """hidden_size""": """d_model""",
      """num_hidden_layers""": """decoder_layers""",
  }
  def __init__(					self     :					Dict     ,				UpperCamelCase__     :					Optional[Any]=5_0265     ,				UpperCamelCase__     :					int=1024     ,				UpperCamelCase__     :					int=12     ,				UpperCamelCase__     :					List[Any]=16     ,				UpperCamelCase__     :					Union[str, Any]=4096     ,				UpperCamelCase__     :					List[Any]="gelu"     ,				UpperCamelCase__     :					Optional[int]=512     ,				UpperCamelCase__     :					Optional[Any]=0.1     ,				UpperCamelCase__     :					int=0.0     ,				UpperCamelCase__     :					str=0.0     ,				UpperCamelCase__     :					Optional[int]=2     ,				UpperCamelCase__     :					List[str]=0.02     ,				UpperCamelCase__     :					Dict=0.0     ,				UpperCamelCase__     :					List[str]=True     ,				UpperCamelCase__     :					str=False     ,				UpperCamelCase__     :					List[str]=True     ,				UpperCamelCase__     :					Tuple=True     ,				UpperCamelCase__     :					str=1     ,				UpperCamelCase__     :					Union[str, Any]=0     ,				UpperCamelCase__     :					List[str]=2     ,				**UpperCamelCase__     :					Tuple     ,				)     ->      Dict:







    """simple docstring"""






    __magic_name__								=			vocab_size
    __magic_name__								=			d_model
    __magic_name__								=			decoder_layers
    __magic_name__								=			decoder_attention_heads
    __magic_name__								=			decoder_ffn_dim
    __magic_name__								=			activation_function
    __magic_name__								=			max_position_embeddings
    __magic_name__								=			dropout
    __magic_name__								=			attention_dropout
    __magic_name__								=			activation_dropout
    __magic_name__								=			init_std
    __magic_name__								=			decoder_layerdrop
    __magic_name__								=			use_cache
    __magic_name__								=			scale_embedding
    __magic_name__								=			use_learned_position_embeddings
    __magic_name__								=			layernorm_embedding

    super().__init__(
        pad_token_id=UpperCamelCase__     ,				bos_token_id=UpperCamelCase__     ,				eos_token_id=UpperCamelCase__     ,				decoder_start_token_id=UpperCamelCase__     ,				**UpperCamelCase__     ,				)

