






"""simple docstring"""



import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
						from transformers import WhisperFeatureExtractor

if is_torch_available():
						import torch

A_      =							random.Random()




def        _lowerCAmelCase						(							UpperCAmelCase__    :   Dict,  UpperCAmelCase__    :   Tuple=1.0,  UpperCAmelCase__    :   Optional[int]=None,  UpperCAmelCase__    :   str=None )						->Union[str, Any]:
						if rng is None:
												A__  :					Optional[int]											=					global_rng

						A__  :					Optional[Any]											=					[]
						for batch_idx in range(shape[0] ):
												values.append([] )
												for _ in range(shape[1] ):
																		values[-1].append(rng.random() * scale )

						return values






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


							def __init__(     self :  Union[str, Any]     ,      snake_case :  str     ,      snake_case :  List[str]=7     ,      snake_case :  str=400     ,      snake_case :  Optional[Any]=2000     ,      snake_case :  Union[str, Any]=10     ,      snake_case :  str=160     ,      snake_case :  List[str]=8     ,      snake_case :  List[Any]=0.0     ,      snake_case :  Optional[Any]=4000     ,      snake_case :  Any=False     ,      snake_case :  int=True     ,      ):
													'''simple docstring'''





													A__  :					Any											=					parent
													A__  :					str											=					batch_size
													A__  :					List[str]											=					min_seq_length
													A__  :					Dict											=					max_seq_length
													A__  :					str											=					(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
													A__  :					Dict											=					padding_value
													A__  :					Optional[Any]											=					sampling_rate
													A__  :					Any											=					return_attention_mask
													A__  :					Optional[int]											=					do_normalize
													A__  :					Tuple											=					feature_size
													A__  :					Optional[Any]											=					chunk_length
													A__  :					Union[str, Any]											=					hop_length


							def   _UpperCamelCase    (     self :  Union[str, Any]							):
													'''simple docstring'''





													return {
													    "feature_size": self.feature_size,
													    "hop_length": self.hop_length,
													    "chunk_length": self.chunk_length,
													    "padding_value": self.padding_value,
													    "sampling_rate": self.sampling_rate,
													    "return_attention_mask": self.return_attention_mask,
													    "do_normalize": self.do_normalize,
													}


							def   _UpperCamelCase    (     self :  Union[str, Any]     ,      snake_case :  Dict=False     ,      snake_case :  Optional[Any]=False							):
													'''simple docstring'''





													def _flatten(snake_case :  Dict							):
																			return list(itertools.chain(*snake_case							)							)

													if equal_length:
																			A__  :					Dict											=					[floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
													else:
																			# make sure that inputs increase in size
																			A__  :					Optional[int]											=					[
																			    floats_list((x, self.feature_size)							)
																			    for x in range(self.min_seq_length     ,      self.max_seq_length     ,      self.seq_length_diff							)
																			]
													if numpify:
																			A__  :					List[str]											=					[np.asarray(snake_case							) for x in speech_inputs]
													return speech_inputs






@require_torch
@require_torchaudio
class 							__SCREAMING_SNAKE_CASE  (			UpperCamelCase    ,						unittest.TestCase						):
							snake_case_       =							WhisperFeatureExtractor if is_speech_available() else None


							def   _UpperCamelCase    (     self :  Dict							):
													'''simple docstring'''





													A__  :					str											=					WhisperFeatureExtractionTester(self							)


							def   _UpperCamelCase    (     self :  int							):
													'''simple docstring'''





													A__  :					Any											=					self.feature_extraction_class(**self.feat_extract_dict							)

													with tempfile.TemporaryDirectory() as tmpdirname:
																			A__  :					List[Any]											=					feat_extract_first.save_pretrained(snake_case							)[0]
																			check_json_file_has_correct_format(snake_case							)
																			A__  :					Union[str, Any]											=					self.feature_extraction_class.from_pretrained(snake_case							)

													A__  :					str											=					feat_extract_first.to_dict()
													A__  :					Union[str, Any]											=					feat_extract_second.to_dict()
													A__  :					List[Any]											=					feat_extract_first.mel_filters
													A__  :					Optional[Any]											=					feat_extract_second.mel_filters
													self.assertTrue(np.allclose(snake_case     ,      snake_case							)							)
													self.assertEqual(snake_case     ,      snake_case							)


							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													A__  :					Any											=					self.feature_extraction_class(**self.feat_extract_dict							)

													with tempfile.TemporaryDirectory() as tmpdirname:
																			A__  :					Any											=					os.path.join(snake_case     ,      """feat_extract.json"""							)
																			feat_extract_first.to_json_file(snake_case							)
																			A__  :					int											=					self.feature_extraction_class.from_json_file(snake_case							)

													A__  :					Dict											=					feat_extract_first.to_dict()
													A__  :					str											=					feat_extract_second.to_dict()
													A__  :					str											=					feat_extract_first.mel_filters
													A__  :					Dict											=					feat_extract_second.mel_filters
													self.assertTrue(np.allclose(snake_case     ,      snake_case							)							)
													self.assertEqual(snake_case     ,      snake_case							)


							def   _UpperCamelCase    (     self :  Any							):
													'''simple docstring'''





													A__  :					Any											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
													# create three inputs of length 800, 1000, and 1200
													A__  :					str											=					[floats_list((1, x)							)[0] for x in range(800     ,      1400     ,      200							)]
													A__  :					Union[str, Any]											=					[np.asarray(snake_case							) for speech_input in speech_inputs]

													# Test feature size
													A__  :					Dict											=					feature_extractor(snake_case     ,      padding="""max_length"""     ,      return_tensors="""np"""							).input_features
													self.assertTrue(input_features.ndim == 3							)
													self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames							)
													self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size							)

													# Test not batched input
													A__  :					str											=					feature_extractor(speech_inputs[0]     ,      return_tensors="""np"""							).input_features
													A__  :					Optional[int]											=					feature_extractor(np_speech_inputs[0]     ,      return_tensors="""np"""							).input_features
													self.assertTrue(np.allclose(snake_case     ,      snake_case     ,      atol=1e-3							)							)

													# Test batched
													A__  :					List[str]											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													A__  :					List[str]											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													for enc_seq_a, enc_seq_a in zip(snake_case     ,      snake_case							):
																			self.assertTrue(np.allclose(snake_case     ,      snake_case     ,      atol=1e-3							)							)

													# Test 2-D numpy arrays are batched.
													A__  :					Tuple											=					[floats_list((1, x)							)[0] for x in (800, 800, 800)]
													A__  :					str											=					np.asarray(snake_case							)
													A__  :					List[str]											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													A__  :					Optional[int]											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													for enc_seq_a, enc_seq_a in zip(snake_case     ,      snake_case							):
																			self.assertTrue(np.allclose(snake_case     ,      snake_case     ,      atol=1e-3							)							)

													# Test truncation required
													A__  :					Optional[Any]											=					[floats_list((1, x)							)[0] for x in range(200     ,      (feature_extractor.n_samples + 500)     ,      200							)]
													A__  :					Union[str, Any]											=					[np.asarray(snake_case							) for speech_input in speech_inputs]

													A__  :					Union[str, Any]											=					[x[: feature_extractor.n_samples] for x in speech_inputs]
													A__  :					str											=					[np.asarray(snake_case							) for speech_input in speech_inputs_truncated]

													A__  :					Optional[int]											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													A__  :					str											=					feature_extractor(snake_case     ,      return_tensors="""np"""							).input_features
													for enc_seq_a, enc_seq_a in zip(snake_case     ,      snake_case							):
																			self.assertTrue(np.allclose(snake_case     ,      snake_case     ,      atol=1e-3							)							)


							def   _UpperCamelCase    (     self :  str							):
													'''simple docstring'''





													import torch

													A__  :					Tuple											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
													A__  :					List[str]											=					np.random.rand(100     ,      32							).astype(np.floataa							)
													A__  :					Tuple											=					np_speech_inputs.tolist()

													for inputs in [py_speech_inputs, np_speech_inputs]:
																			A__  :					Optional[Any]											=					feature_extractor.pad([{"""input_features""": inputs}]     ,      return_tensors="""np"""							)
																			self.assertTrue(np_processed.input_features.dtype == np.floataa							)
																			A__  :					Optional[int]											=					feature_extractor.pad([{"""input_features""": inputs}]     ,      return_tensors="""pt"""							)
																			self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)


							def   _UpperCamelCase    (     self :  Optional[Any]     ,      snake_case :  Optional[int]							):
													'''simple docstring'''





													A__  :					int											=					load_dataset("""hf-internal-testing/librispeech_asr_dummy"""     ,      """clean"""     ,      split="""validation"""							)
													# automatic decoding with librispeech
													A__  :					Union[str, Any]											=					ds.sort("""id"""							).select(range(snake_case							)							)[:num_samples]["""audio"""]

													return [x["array"] for x in speech_samples]


							def   _UpperCamelCase    (     self :  List[Any]							):
													'''simple docstring'''





													A__  :					str											=					torch.tensor(
													    [
													        0.1193, -0.0946, -0.1098, -0.0196, 0.0225, -0.0690, -0.1736, 0.0951,
													        0.0971, -0.0817, -0.0702, 0.0162, 0.0260, 0.0017, -0.0192, -0.1678,
													        0.0709, -0.1867, -0.0655, -0.0274, -0.0234, -0.1884, -0.0516, -0.0554,
													        -0.0274, -0.1425, -0.1423, 0.0837, 0.0377, -0.0854
													    ]							)
													# fmt: on

													A__  :					Optional[Any]											=					self._load_datasamples(1							)
													A__  :					Union[str, Any]											=					WhisperFeatureExtractor()
													A__  :					List[str]											=					feature_extractor(snake_case     ,      return_tensors="""pt"""							).input_features
													self.assertEqual(input_features.shape     ,      (1, 80, 3000)							)
													self.assertTrue(torch.allclose(input_features[0, 0, :30]     ,      snake_case     ,      atol=1e-4							)							)


							def   _UpperCamelCase    (     self :  Tuple							):
													'''simple docstring'''





													A__  :					Union[str, Any]											=					self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
													A__  :					Union[str, Any]											=					self._load_datasamples(1							)[0]
													A__  :					Any											=					((audio - audio.min()) / (audio.max() - audio.min())) * 6_5535  # Rescale to [0, 65535] to show issue
													A__  :					str											=					feat_extract.zero_mean_unit_var_norm([audio]     ,      attention_mask=snake_case							)[0]

													self.assertTrue(np.all(np.mean(snake_case							) < 1e-3							)							)
													self.assertTrue(np.all(np.abs(np.var(snake_case							) - 1							) < 1e-3							)							)














"""simple docstring"""



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




def        _lowerCAmelCase						(							UpperCAmelCase__    :   Union[str, Any],  UpperCAmelCase__    :   Tuple=False )						->str:
						A__  :					Optional[int]											=					[]
						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"
												A__  :					Optional[int]											=					[(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        _lowerCAmelCase						(							UpperCAmelCase__    :   Optional[int],  UpperCAmelCase__    :   Tuple,  UpperCAmelCase__    :   List[Any]=False )						->str:
						for i in range(config.num_hidden_layers ):
												if base_model:
																		A__  :					Any											=					""""""
												else:
																		A__  :					Tuple											=					"""deit."""
												# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
												A__  :					Any											=					state_dict.pop(f'blocks.{i}.attn.qkv.weight' )
												A__  :					Tuple											=					state_dict.pop(f'blocks.{i}.attn.qkv.bias' )
												# next, add query, keys and values (in that order) to the state dict
												A__  :					List[Any]											=					in_proj_weight[
												    : config.hidden_size, :
												]
												A__  :					str											=					in_proj_bias[: config.hidden_size]
												A__  :					Any											=					in_proj_weight[
												    config.hidden_size : config.hidden_size * 2, :
												]
												A__  :					Dict											=					in_proj_bias[
												    config.hidden_size : config.hidden_size * 2
												]
												A__  :					Optional[Any]											=					in_proj_weight[
												    -config.hidden_size :, :
												]
												A__  :					Any											=					in_proj_bias[-config.hidden_size :]




def        _lowerCAmelCase						(							UpperCAmelCase__    :   List[Any],  UpperCAmelCase__    :   List[Any],  UpperCAmelCase__    :   Union[str, Any] )						->Any:
						A__  :					int											=					dct.pop(UpperCAmelCase__ )
						A__  :					Tuple											=					val




def        _lowerCAmelCase						(							)						->List[Any]:
						A__  :					Optional[int]											=					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
						A__  :					int											=					Image.open(requests.get(UpperCAmelCase__,  stream=UpperCAmelCase__ ).raw )
						return im




@torch.no_grad()
def        _lowerCAmelCase						(							UpperCAmelCase__    :   Dict,  UpperCAmelCase__    :   Any )						->Tuple:
						A__  :					List[Any]											=					DeiTConfig()
						# all deit models have fine-tuned heads
						A__  :					Tuple											=					False
						# dataset (fine-tuned on ImageNet 2012), patch_size and image_size
						A__  :					str											=					1_0_0_0
						A__  :					List[str]											=					"""huggingface/label-files"""
						A__  :					Dict											=					"""imagenet-1k-id2label.json"""
						A__  :					List[str]											=					json.load(open(hf_hub_download(UpperCAmelCase__,  UpperCAmelCase__,  repo_type="""dataset""" ),  """r""" ) )
						A__  :					Dict											=					{int(UpperCAmelCase__ ): v for k, v in idalabel.items()}
						A__  :					Optional[int]											=					idalabel
						A__  :					Dict											=					{v: k for k, v in idalabel.items()}
						A__  :					List[str]											=					int(deit_name[-6:-4] )
						A__  :					str											=					int(deit_name[-3:] )
						# size of the architecture
						if deit_name[9:].startswith("""tiny""" ):
												A__  :					List[str]											=					1_9_2
												A__  :					int											=					7_6_8
												A__  :					List[Any]											=					1_2
												A__  :					Dict											=					3
						elif deit_name[9:].startswith("""small""" ):
												A__  :					List[Any]											=					3_8_4
												A__  :					List[str]											=					1_5_3_6
												A__  :					Any											=					1_2
												A__  :					Union[str, Any]											=					6
						if deit_name[9:].startswith("""base""" ):
												pass
						elif deit_name[4:].startswith("""large""" ):
												A__  :					int											=					1_0_2_4
												A__  :					str											=					4_0_9_6
												A__  :					Any											=					2_4
												A__  :					int											=					1_6

						# load original model from timm
						A__  :					Dict											=					timm.create_model(UpperCAmelCase__,  pretrained=UpperCAmelCase__ )
						timm_model.eval()

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

						# load HuggingFace model
						A__  :					str											=					DeiTForImageClassificationWithTeacher(UpperCAmelCase__ ).eval()
						model.load_state_dict(UpperCAmelCase__ )

						# Check outputs on an image, prepared by DeiTImageProcessor
						A__  :					int											=					int(
						    (2_5_6 / 2_2_4) * config.image_size )  # to maintain same ratio w.r.t. 224 images, see https://github.com/facebookresearch/deit/blob/ab5715372db8c6cad5740714b2216d55aeae052e/datasets.py#L103
						A__  :					Any											=					DeiTImageProcessor(size=UpperCAmelCase__,  crop_size=config.image_size )
						A__  :					Union[str, Any]											=					image_processor(images=prepare_img(),  return_tensors="""pt""" )
						A__  :					Optional[Any]											=					encoding["""pixel_values"""]
						A__  :					Union[str, Any]											=					model(UpperCAmelCase__ )

						A__  :					Union[str, Any]											=					timm_model(UpperCAmelCase__ )
						assert timm_logits.shape == outputs.logits.shape
						assert torch.allclose(UpperCAmelCase__,  outputs.logits,  atol=1e-3 )

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


if __name__ == "__main__":
						A_      =							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.'''
						)

						A_      =							parser.parse_args()
						convert_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)






