




from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TFDeiTForImageClassification,
           TFDeiTForImageClassificationWithTeacher,
           TFDeiTForMaskedImageModeling,
           TFDeiTModel,
       )
       from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import DeiTImageProcessor




class 		UpperCamelCase__	:

    '''simple docstring'''



    def __init__(    self     ,  UpperCamelCase__     ,  UpperCamelCase__=13     ,  UpperCamelCase__=30     ,  UpperCamelCase__=2     ,  UpperCamelCase__=3     ,  UpperCamelCase__=True     ,  UpperCamelCase__=True     ,  UpperCamelCase__=32     ,  UpperCamelCase__=2     ,  UpperCamelCase__=4     ,  UpperCamelCase__=37     ,  UpperCamelCase__="gelu"     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=10     ,  UpperCamelCase__=0.02     ,  UpperCamelCase__=3     ,  UpperCamelCase__=None     ,  UpperCamelCase__=2     ,  )     ->      Union[str, Any]:
           lowerCamelCase   :      int			 =					parent
           lowerCamelCase   :      Optional[int]			 =					batch_size
           lowerCamelCase   :      Any			 =					image_size
           lowerCamelCase   :      List[str]			 =					patch_size
           lowerCamelCase   :      Tuple			 =					num_channels
           lowerCamelCase   :      List[Any]			 =					is_training
           lowerCamelCase   :      List[str]			 =					use_labels
           lowerCamelCase   :      int			 =					hidden_size
           lowerCamelCase   :      Tuple			 =					num_hidden_layers
           lowerCamelCase   :      Tuple			 =					num_attention_heads
           lowerCamelCase   :      List[Any]			 =					intermediate_size
           lowerCamelCase   :      Any			 =					hidden_act
           lowerCamelCase   :      Dict			 =					hidden_dropout_prob
           lowerCamelCase   :      Tuple			 =					attention_probs_dropout_prob
           lowerCamelCase   :      Optional[Any]			 =					type_sequence_label_size
           lowerCamelCase   :      Union[str, Any]			 =					initializer_range
           lowerCamelCase   :      List[str]			 =					scope
           lowerCamelCase   :      List[Any]			 =					encoder_stride

           # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
           lowerCamelCase   :      List[str]			 =					(image_size // patch_size) ** 2
           lowerCamelCase   :      Dict			 =					num_patches + 2







    def     _lowercase     (    self     )     ->      Optional[Any]:
           lowerCamelCase   :      Optional[Any]			 =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

           lowerCamelCase   :      Optional[int]			 =					None
           if self.use_labels:
                  lowerCamelCase   :      Optional[int]			 =					ids_tensor([self.batch_size]     ,  self.type_sequence_label_size     )

           lowerCamelCase   :      Tuple			 =					self.get_config()

           return config, pixel_values, labels







    def     _lowercase     (    self     )     ->      Optional[int]:
           return DeiTConfig(
               image_size=self.image_size     ,  patch_size=self.patch_size     ,  num_channels=self.num_channels     ,  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=UpperCamelCase__     ,  initializer_range=self.initializer_range     ,  encoder_stride=self.encoder_stride     ,  )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      List[str]:
           lowerCamelCase   :      List[str]			 =					TFDeiTModel(config=UpperCamelCase__     )
           lowerCamelCase   :      Union[str, Any]			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size)     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      Optional[int]:
           lowerCamelCase   :      int			 =					TFDeiTForMaskedImageModeling(config=UpperCamelCase__     )
           lowerCamelCase   :      List[str]			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(
               result.reconstruction.shape     ,  (self.batch_size, self.num_channels, self.image_size, self.image_size)     )

           # test greyscale images
           lowerCamelCase   :      str			 =					1
           lowerCamelCase   :      Optional[int]			 =					TFDeiTForMaskedImageModeling(UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
           lowerCamelCase   :      Dict			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(result.reconstruction.shape     ,  (self.batch_size, 1, self.image_size, self.image_size)     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      int:
           lowerCamelCase   :      Optional[Any]			 =					self.type_sequence_label_size
           lowerCamelCase   :      Optional[int]			 =					TFDeiTForImageClassification(UpperCamelCase__     )
           lowerCamelCase   :      Tuple			 =					model(UpperCamelCase__     ,  labels=UpperCamelCase__     )
           self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.type_sequence_label_size)     )

           # test greyscale images
           lowerCamelCase   :      List[str]			 =					1
           lowerCamelCase   :      str			 =					TFDeiTForImageClassification(UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
           lowerCamelCase   :      Optional[int]			 =					model(UpperCamelCase__     ,  labels=UpperCamelCase__     )
           self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.type_sequence_label_size)     )







    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      int			 =					self.prepare_config_and_inputs()
           lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase   :      Any			 =					config_and_inputs
           lowerCamelCase   :      Dict			 =					{"pixel_values": pixel_values}
           return config, inputs_dict




@require_tf
class 		UpperCamelCase__	(lowerCAmelCase__    ,				lowerCAmelCase__    ,				unittest.TestCase     ):

    '''simple docstring'''



    lowerCamelCase_      :	List[Any]	     =      (
        (
            TFDeiTModel,
            TFDeiTForImageClassification,
            TFDeiTForImageClassificationWithTeacher,
            TFDeiTForMaskedImageModeling,
        )
        if is_tf_available()
        else ()
    )
    lowerCamelCase_      :	Optional[Any]	     =      (
        {
            """feature-extraction""": TFDeiTModel,
            """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
        }
        if is_tf_available()
        else {}
    )

    lowerCamelCase_      :	Dict	     =      False
    lowerCamelCase_      :	Tuple	     =      False
    lowerCamelCase_      :	str	     =      False
    lowerCamelCase_      :	List[Any]	     =      False







    def     _lowercase     (    self     )     ->      Union[str, Any]:
           lowerCamelCase   :      Optional[int]			 =					TFDeiTModelTester(self     )
           lowerCamelCase   :      Any			 =					ConfigTester(self     ,  config_class=UpperCamelCase__     ,  has_text_modality=UpperCamelCase__     ,  hidden_size=37     )







    def     _lowercase     (    self     )     ->      Optional[int]:
           self.config_tester.run_common_tests()







    @unittest.skip(reason="DeiT does not use inputs_embeds"     )
    def     _lowercase     (    self     )     ->      int:
           pass







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase ,      lowerCamelCase   :      List[Any]			 =					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  lowerCamelCase   :      Optional[Any]			 =					model_class(UpperCamelCase__     )
                  self.assertIsInstance(model.get_input_embeddings()     ,  (tf.keras.layers.Layer)     )
                  lowerCamelCase   :      Optional[int]			 =					model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(UpperCamelCase__     ,  tf.keras.layers.Dense     )     )







    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase ,      lowerCamelCase   :      str			 =					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  lowerCamelCase   :      Optional[Any]			 =					model_class(UpperCamelCase__     )
                  lowerCamelCase   :      int			 =					inspect.signature(model.call     )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  lowerCamelCase   :      Any			 =					[*signature.parameters.keys()]

                  lowerCamelCase   :      str			 =					["pixel_values"]
                  self.assertListEqual(arg_names[:1]     ,  UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      int			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase   :      Optional[int]			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      Optional[int]:
           lowerCamelCase   :      Union[str, Any]			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__=False     )     ->      List[str]:
           lowerCamelCase   :      Optional[Any]			 =					super()._prepare_for_class(UpperCamelCase__     ,  UpperCamelCase__     ,  return_labels=UpperCamelCase__     )

           if return_labels:
                  if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call     ).parameters:
                         del inputs_dict["labels"]

           return inputs_dict







    @slow
    def     _lowercase     (    self     )     ->      Optional[Any]:
           for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCamelCase   :      Any			 =					TFDeiTModel.from_pretrained(UpperCamelCase__     )
                  self.assertIsNotNone(UpperCamelCase__     )




def      A      (      )	-> List[str]:
       lowerCamelCase   :      str			 =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
       return image




@require_tf
@require_vision
class 		UpperCamelCase__	(unittest.TestCase     ):

    '''simple docstring'''



    @cached_property
    def     _lowercase     (    self     )     ->      List[str]:
           return (
               DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"     )
               if is_vision_available()
               else None
           )







    @slow
    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      int			 =					TFDeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"     )

           lowerCamelCase   :      str			 =					self.default_image_processor
           lowerCamelCase   :      Optional[Any]			 =					prepare_img()
           lowerCamelCase   :      Tuple			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="tf"     )

           # forward pass
           lowerCamelCase   :      Optional[int]			 =					model(**UpperCamelCase__     )

           # verify the logits
           lowerCamelCase   :      str			 =					tf.TensorShape((1, 1000)     )
           self.assertEqual(outputs.logits.shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      Optional[Any]			 =					tf.constant([-1.0266, 0.1912, -1.2861]     )

           self.assertTrue(np.allclose(outputs.logits[0, :3]     ,  UpperCamelCase__     ,  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()
SCREAMING_SNAKE_CASE__       :      str  =   logging.get_logger(__name__)

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

def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> str:
       lowerCamelCase   :      List[str]			 =					dct.pop(_SCREAMING_SNAKE_CASE    )
       lowerCamelCase   :      Any			 =					val

def      A      (      )	-> List[str]:
       lowerCamelCase   :      Union[str, Any]			 =					"http://images.cocodataset.org/val2017/000000039769.jpg"
       lowerCamelCase   :      str			 =					Image.open(requests.get(_SCREAMING_SNAKE_CASE ,stream=_SCREAMING_SNAKE_CASE    ).raw    )
       return im

@torch.no_grad()
def      A      (      _SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE    )	-> Optional[Any]:
       lowerCamelCase   :      Union[str, Any]			 =					DeiTConfig()
       # all deit models have fine-tuned heads
       lowerCamelCase   :      Optional[int]			 =					False
       # dataset (fine-tuned on ImageNet 2012), patch_size and image_size
       lowerCamelCase   :      Dict			 =					1000
       lowerCamelCase   :      Tuple			 =					"huggingface/label-files"
       lowerCamelCase   :      List[str]			 =					"imagenet-1k-id2label.json"
       lowerCamelCase   :      List[Any]			 =					json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,repo_type="dataset"    ) ,"r"    )    )
       lowerCamelCase   :      Optional[int]			 =					{int(_SCREAMING_SNAKE_CASE    ): v for k, v in idalabel.items()}
       lowerCamelCase   :      Tuple			 =					idalabel
       lowerCamelCase   :      str			 =					{v: k for k, v in idalabel.items()}
       lowerCamelCase   :      Dict			 =					int(deit_name[-6:-4]    )
       lowerCamelCase   :      Optional[Any]			 =					int(deit_name[-3:]    )
       # size of the architecture
       if deit_name[9:].startswith("tiny"    ):
              lowerCamelCase   :      Optional[Any]			 =					192
              lowerCamelCase   :      List[str]			 =					768
              lowerCamelCase   :      Tuple			 =					12
              lowerCamelCase   :      Optional[Any]			 =					3
       elif deit_name[9:].startswith("small"    ):
              lowerCamelCase   :      str			 =					384
              lowerCamelCase   :      Optional[Any]			 =					1536
              lowerCamelCase   :      Dict			 =					12
              lowerCamelCase   :      Optional[int]			 =					6
       if deit_name[9:].startswith("base"    ):
              pass
       elif deit_name[4:].startswith("large"    ):
              lowerCamelCase   :      str			 =					1024
              lowerCamelCase   :      List[str]			 =					4096
              lowerCamelCase   :      Any			 =					24
              lowerCamelCase   :      Dict			 =					16

       # load original model from timm
       lowerCamelCase   :      List[Any]			 =					timm.create_model(_SCREAMING_SNAKE_CASE ,pretrained=_SCREAMING_SNAKE_CASE    )
       timm_model.eval()

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

       # load HuggingFace model
       lowerCamelCase   :      Optional[Any]			 =					DeiTForImageClassificationWithTeacher(_SCREAMING_SNAKE_CASE    ).eval()
       model.load_state_dict(_SCREAMING_SNAKE_CASE    )

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

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

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


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

       SCREAMING_SNAKE_CASE__       :      List[str]  =   parser.parse_args()
       convert_deit_checkpoint(args.deit_name, args.pytorch_dump_folder_path)






