



import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch
			from torch import nn

			from transformers import ASTForAudioClassification, ASTModel
			from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
			    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
			)


if is_torchaudio_available():
			import torchaudio

			from transformers import ASTFeatureExtractor

class        UpperCAmelCase_  :

						'''simple docstring'''




						def __init__(   self				,							_A				,							_A=13				,							_A=2				,							_A=24				,							_A=16				,							_A=True				,							_A=True				,							_A=32				,							_A=5				,							_A=4				,							_A=37				,							_A="gelu"				,							_A=0.1				,							_A=0.1				,							_A=10				,							_A=0.0_2				,							_A=None				,							_A=2				,							_A=2				,							):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= parent
													__SCREAMING_SNAKE_CASE									= batch_size
													__SCREAMING_SNAKE_CASE									= patch_size
													__SCREAMING_SNAKE_CASE									= max_length
													__SCREAMING_SNAKE_CASE									= num_mel_bins
													__SCREAMING_SNAKE_CASE									= is_training
													__SCREAMING_SNAKE_CASE									= use_labels
													__SCREAMING_SNAKE_CASE									= hidden_size
													__SCREAMING_SNAKE_CASE									= num_hidden_layers
													__SCREAMING_SNAKE_CASE									= num_attention_heads
													__SCREAMING_SNAKE_CASE									= intermediate_size
													__SCREAMING_SNAKE_CASE									= hidden_act
													__SCREAMING_SNAKE_CASE									= hidden_dropout_prob
													__SCREAMING_SNAKE_CASE									= attention_probs_dropout_prob
													__SCREAMING_SNAKE_CASE									= type_sequence_label_size
													__SCREAMING_SNAKE_CASE									= initializer_range
													__SCREAMING_SNAKE_CASE									= scope
													__SCREAMING_SNAKE_CASE									= frequency_stride
													__SCREAMING_SNAKE_CASE									= time_stride

													# in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
													__SCREAMING_SNAKE_CASE									= (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
													__SCREAMING_SNAKE_CASE									= (self.max_length - self.patch_size) // self.time_stride + 1
													__SCREAMING_SNAKE_CASE									= frequency_out_dimension * time_out_dimension
													__SCREAMING_SNAKE_CASE									= num_patches + 2






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]		)

													__SCREAMING_SNAKE_CASE									= None
													if self.use_labels:
																				__SCREAMING_SNAKE_CASE									= ids_tensor([self.batch_size]				,							self.type_sequence_label_size		)

													__SCREAMING_SNAKE_CASE									= self.get_config()

													return config, input_values, labels






						def 	_A     (   self		):







													'''simple docstring'''

													return ASTConfig(
													    patch_size=self.patch_size				,							max_length=self.max_length				,							num_mel_bins=self.num_mel_bins				,							hidden_size=self.hidden_size				,							num_hidden_layers=self.num_hidden_layers				,							num_attention_heads=self.num_attention_heads				,							intermediate_size=self.intermediate_size				,							hidden_act=self.hidden_act				,							hidden_dropout_prob=self.hidden_dropout_prob				,							attention_probs_dropout_prob=self.attention_probs_dropout_prob				,							is_decoder=_A				,							initializer_range=self.initializer_range				,							frequency_stride=self.frequency_stride				,							time_stride=self.time_stride				,							)






						def 	_A     (   self				,							_A				,							_A				,							_A		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= ASTModel(config=_A		)
													model.to(_A		)
													model.eval()
													__SCREAMING_SNAKE_CASE									= model(_A		)
													self.parent.assertEqual(result.last_hidden_state.shape				,							(self.batch_size, self.seq_length, self.hidden_size)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.prepare_config_and_inputs()
													(
													    (
													    __SCREAMING_SNAKE_CASE	
													)							,	(
													    __SCREAMING_SNAKE_CASE	
													)							,	(
													    __SCREAMING_SNAKE_CASE	
													)							,		
													)									= config_and_inputs
													__SCREAMING_SNAKE_CASE									= {'input_values': input_values}
													return config, inputs_dict







@require_torch
class        UpperCAmelCase_  (       UpperCamelCase_   ,   UpperCamelCase_   ,   unittest.TestCase	):

						'''simple docstring'''




						UpperCamelCase__      :			List[str]      =  (
						    (
						        ASTModel,
						        ASTForAudioClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						UpperCamelCase__      :			int      =  (
						    {'''audio-classification''': ASTForAudioClassification, '''feature-extraction''': ASTModel}
						    if is_torch_available()
						    else {}
						)
						UpperCamelCase__      :			List[Any]      =  False
						UpperCamelCase__      :			Union[str, Any]      =  False
						UpperCamelCase__      :			Optional[Any]      =  False
						UpperCamelCase__      :			Tuple      =  False






						def 	_A     (   self				,							_A				,							_A				,							_A				,							_A				,							_A		):







													'''simple docstring'''

													if pipeline_test_casse_name == "AudioClassificationPipelineTests":
																				return True

													return False






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= ASTModelTester(self		)
													__SCREAMING_SNAKE_CASE									= ConfigTester(self				,							config_class=_A				,							has_text_modality=_A				,							hidden_size=37		)






						def 	_A     (   self		):







													'''simple docstring'''

													self.config_tester.run_common_tests()






						@unittest.skip(reason='AST does not use inputs_embeds'		)
						def 	_A     (   self		):







													'''simple docstring'''

													pass






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				__SCREAMING_SNAKE_CASE									= model_class(_A		)
																				self.assertIsInstance(model.get_input_embeddings()				,							(nn.Module)		)
																				__SCREAMING_SNAKE_CASE									= model.get_output_embeddings()
																				self.assertTrue(x is None or isinstance(_A				,							nn.Linear		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				__SCREAMING_SNAKE_CASE									= model_class(_A		)
																				__SCREAMING_SNAKE_CASE									= inspect.signature(model.forward		)
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				__SCREAMING_SNAKE_CASE									= [*signature.parameters.keys()]

																				__SCREAMING_SNAKE_CASE									= ['input_values']
																				self.assertListEqual(arg_names[:1]				,							_A		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*_A		)






						@slow
						def 	_A     (   self		):







													'''simple docstring'''

													for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				__SCREAMING_SNAKE_CASE									= ASTModel.from_pretrained(_A		)
																				self.assertIsNotNone(_A		)







def      __lowercase       (			)    ->			Dict:
							__SCREAMING_SNAKE_CASE									= hf_hub_download(
							    repo_id='nielsr/audio-spectogram-transformer-checkpoint'							,       filename='sample_audio.flac'							,       repo_type='dataset'       )

							__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= torchaudio.load(a__       )

							return audio, sampling_rate

@require_torch
@require_torchaudio
class        UpperCAmelCase_  (       unittest.TestCase	):

						'''simple docstring'''




						@cached_property
						def 	_A     (   self		):







													'''simple docstring'''

													return (
													    ASTFeatureExtractor.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'		)
													    if is_torchaudio_available()
													    else None
													)






						@slow
						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= self.default_feature_extractor
													__SCREAMING_SNAKE_CASE									= ASTForAudioClassification.from_pretrained('MIT/ast-finetuned-audioset-10-10-0.4593'		).to(_A		)

													__SCREAMING_SNAKE_CASE									= self.default_feature_extractor
													__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= prepare_audio()
													__SCREAMING_SNAKE_CASE									= audio.squeeze().numpy()
													__SCREAMING_SNAKE_CASE									= feature_extractor(_A				,							sampling_rate=_A				,							return_tensors='pt'		).to(_A		)

													# forward pass
													with torch.no_grad():
																				__SCREAMING_SNAKE_CASE									= model(**_A		)

													# verify the logits
													__SCREAMING_SNAKE_CASE									= torch.Size((1, 527)		)
													self.assertEqual(outputs.logits.shape				,							_A		)

													__SCREAMING_SNAKE_CASE									= torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]		).to(_A		)

													self.assertTrue(torch.allclose(outputs.logits[0, :3]				,							_A				,							atol=1e-4		)		)











import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTConfig, DeiTForImageClassificationWithTeacher, DeiTImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__		:					Optional[int]							=logging.get_logger(__name__)






def      __lowercase       (			a__							,       a__=False       )    ->			Tuple:
							__SCREAMING_SNAKE_CASE									= []
							for i in range(config.num_hidden_layers       ):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((f"""blocks.{i}.norm1.weight""", f"""deit.encoder.layer.{i}.layernorm_before.weight""")       )
														rename_keys.append((f"""blocks.{i}.norm1.bias""", f"""deit.encoder.layer.{i}.layernorm_before.bias""")       )
														rename_keys.append((f"""blocks.{i}.attn.proj.weight""", f"""deit.encoder.layer.{i}.attention.output.dense.weight""")       )
														rename_keys.append((f"""blocks.{i}.attn.proj.bias""", f"""deit.encoder.layer.{i}.attention.output.dense.bias""")       )
														rename_keys.append((f"""blocks.{i}.norm2.weight""", f"""deit.encoder.layer.{i}.layernorm_after.weight""")       )
														rename_keys.append((f"""blocks.{i}.norm2.bias""", f"""deit.encoder.layer.{i}.layernorm_after.bias""")       )
														rename_keys.append((f"""blocks.{i}.mlp.fc1.weight""", f"""deit.encoder.layer.{i}.intermediate.dense.weight""")       )
														rename_keys.append((f"""blocks.{i}.mlp.fc1.bias""", f"""deit.encoder.layer.{i}.intermediate.dense.bias""")       )
														rename_keys.append((f"""blocks.{i}.mlp.fc2.weight""", f"""deit.encoder.layer.{i}.output.dense.weight""")       )
														rename_keys.append((f"""blocks.{i}.mlp.fc2.bias""", f"""deit.encoder.layer.{i}.output.dense.bias""")       )

							# projection layer + position embeddings
							rename_keys.extend(
							    [
							        ('cls_token', 'deit.embeddings.cls_token'),
							        ('dist_token', 'deit.embeddings.distillation_token'),
							        ('patch_embed.proj.weight', 'deit.embeddings.patch_embeddings.projection.weight'),
							        ('patch_embed.proj.bias', 'deit.embeddings.patch_embeddings.projection.bias'),
							        ('pos_embed', 'deit.embeddings.position_embeddings'),
							    ]       )

							if base_model:
														# layernorm + pooler
														rename_keys.extend(
														    [
														        ('norm.weight', 'layernorm.weight'),
														        ('norm.bias', 'layernorm.bias'),
														        ('pre_logits.fc.weight', 'pooler.dense.weight'),
														        ('pre_logits.fc.bias', 'pooler.dense.bias'),
														    ]       )

														# if just the base model, we should remove "deit" from all keys that start with "deit"
														__SCREAMING_SNAKE_CASE									= [(pair[0], pair[1][4:]) if pair[1].startswith('deit'       ) else pair for pair in rename_keys]
							else:
														# layernorm + classification heads
														rename_keys.extend(
														    [
														        ('norm.weight', 'deit.layernorm.weight'),
														        ('norm.bias', 'deit.layernorm.bias'),
														        ('head.weight', 'cls_classifier.weight'),
														        ('head.bias', 'cls_classifier.bias'),
														        ('head_dist.weight', 'distillation_classifier.weight'),
														        ('head_dist.bias', 'distillation_classifier.bias'),
														    ]       )

							return rename_keys






def      __lowercase       (			a__							,       a__							,       a__=False       )    ->			Tuple:
							for i in range(config.num_hidden_layers       ):
														if base_model:
																					__SCREAMING_SNAKE_CASE									= ''
														else:
																					__SCREAMING_SNAKE_CASE									= 'deit.'
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														__SCREAMING_SNAKE_CASE									= state_dict.pop(f"""blocks.{i}.attn.qkv.weight"""       )
														__SCREAMING_SNAKE_CASE									= state_dict.pop(f"""blocks.{i}.attn.qkv.bias"""       )
														# next, add query, keys and values (in that order) to the state dict
														__SCREAMING_SNAKE_CASE									= in_proj_weight[
														    : config.hidden_size, :
														]
														__SCREAMING_SNAKE_CASE									= in_proj_bias[: config.hidden_size]
														__SCREAMING_SNAKE_CASE									= in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														__SCREAMING_SNAKE_CASE									= in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														__SCREAMING_SNAKE_CASE									= in_proj_weight[
														    -config.hidden_size :, :
														]
														__SCREAMING_SNAKE_CASE									= in_proj_bias[-config.hidden_size :]






def      __lowercase       (			a__							,       a__							,       a__       )    ->			str:
							__SCREAMING_SNAKE_CASE									= dct.pop(a__       )
							__SCREAMING_SNAKE_CASE									= val






def      __lowercase       (			)    ->			List[Any]:
							__SCREAMING_SNAKE_CASE									= 'http://images.cocodataset.org/val2017/000000039769.jpg'
							__SCREAMING_SNAKE_CASE									= Image.open(requests.get(a__							,       stream=a__       ).raw       )
							return im






@torch.no_grad()
def      __lowercase       (			a__							,       a__       )    ->			Dict:
							__SCREAMING_SNAKE_CASE									= DeiTConfig()
							# all deit models have fine-tuned heads
							__SCREAMING_SNAKE_CASE									= False
							# dataset (fine-tuned on ImageNet 2012), patch_size and image_size
							__SCREAMING_SNAKE_CASE									= 10_00
							__SCREAMING_SNAKE_CASE									= 'huggingface/label-files'
							__SCREAMING_SNAKE_CASE									= 'imagenet-1k-id2label.json'
							__SCREAMING_SNAKE_CASE									= json.load(open(hf_hub_download(a__							,       a__							,       repo_type='dataset'       )							,       'r'       )       )
							__SCREAMING_SNAKE_CASE									= {int(a__       ): v for k, v in idalabel.items()}
							__SCREAMING_SNAKE_CASE									= idalabel
							__SCREAMING_SNAKE_CASE									= {v: k for k, v in idalabel.items()}
							__SCREAMING_SNAKE_CASE									= int(deit_name[-6:-4]       )
							__SCREAMING_SNAKE_CASE									= int(deit_name[-3:]       )
							# size of the architecture
							if deit_name[9:].startswith('tiny'       ):
														__SCREAMING_SNAKE_CASE									= 1_92
														__SCREAMING_SNAKE_CASE									= 7_68
														__SCREAMING_SNAKE_CASE									= 12
														__SCREAMING_SNAKE_CASE									= 3
							elif deit_name[9:].startswith('small'       ):
														__SCREAMING_SNAKE_CASE									= 3_84
														__SCREAMING_SNAKE_CASE									= 15_36
														__SCREAMING_SNAKE_CASE									= 12
														__SCREAMING_SNAKE_CASE									= 6
							if deit_name[9:].startswith('base'       ):
														pass
							elif deit_name[4:].startswith('large'       ):
														__SCREAMING_SNAKE_CASE									= 10_24
														__SCREAMING_SNAKE_CASE									= 40_96
														__SCREAMING_SNAKE_CASE									= 24
														__SCREAMING_SNAKE_CASE									= 16

							# load original model from timm
							__SCREAMING_SNAKE_CASE									= timm.create_model(a__							,       pretrained=a__       )
							timm_model.eval()

							# load state_dict of original model, remove and rename some keys
							__SCREAMING_SNAKE_CASE									= timm_model.state_dict()
							__SCREAMING_SNAKE_CASE									= create_rename_keys(a__							,       a__       )
							for src, dest in rename_keys:
														rename_key(a__							,       a__							,       a__       )
							read_in_q_k_v(a__							,       a__							,       a__       )

							# load HuggingFace model
							__SCREAMING_SNAKE_CASE									= DeiTForImageClassificationWithTeacher(a__       ).eval()
							model.load_state_dict(a__       )

							# Check outputs on an image, prepared by DeiTImageProcessor
							__SCREAMING_SNAKE_CASE									= int(
							    (2_56 / 2_24) * config.image_size       )  # to maintain same ratio w.r.t. 224 images, see https://github.com/facebookresearch/deit/blob/ab5715372db8c6cad5740714b2216d55aeae052e/datasets.py#L103
							__SCREAMING_SNAKE_CASE									= DeiTImageProcessor(size=a__							,       crop_size=config.image_size       )
							__SCREAMING_SNAKE_CASE									= image_processor(images=prepare_img()							,       return_tensors='pt'       )
							__SCREAMING_SNAKE_CASE									= encoding['pixel_values']
							__SCREAMING_SNAKE_CASE									= model(a__       )

							__SCREAMING_SNAKE_CASE									= timm_model(a__       )
							assert timm_logits.shape == outputs.logits.shape
							assert torch.allclose(a__							,       outputs.logits							,       atol=1E-3       )

							Path(a__       ).mkdir(exist_ok=a__       )
							print(f"""Saving model {deit_name} to {pytorch_dump_folder_path}"""       )
							model.save_pretrained(a__       )
							print(f"""Saving image processor to {pytorch_dump_folder_path}"""       )
							image_processor.save_pretrained(a__       )


if __name__ == "__main__":
			lowerCAmelCase__		:					Union[str, Any]							=argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '''--deit_name''',
			    default='''vit_deit_base_distilled_patch16_224''',
			    type=str,
			    help='''Name of the DeiT timm 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.'''
			)

			lowerCAmelCase__		:					str							=parser.parse_args()
			convert_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)






