




import argparse
import json
from pathlib import Path

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

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase_					    = logging.get_logger(__name__)



def      __lowerCAmelCase			(					UpperCamelCase	,      UpperCamelCase=False      )     ->	List[str]:
  lowerCAmelCase__				:  Optional[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"""vit.encoder.layer.{i}.layernorm_before.weight""")      )
    rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")      )
    rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")      )
    rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")      )
    rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")      )
    rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")      )
    rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")      )
    rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")      )
    rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")      )
    rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")      )

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

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

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

  return rename_keys



def      __lowerCAmelCase			(					UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase=False      )     ->	List[str]:
  for i in range(config.num_hidden_layers      ):
    if base_model:
      lowerCAmelCase__				:  Any		       =     ''''''
    else:
      lowerCAmelCase__				:  Union[str, Any]		       =     '''vit.'''
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    lowerCAmelCase__				:  Tuple		       =     state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""      )
    lowerCAmelCase__				:  Any		       =     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__				:  str		       =     in_proj_bias[: config.hidden_size]
    lowerCAmelCase__				:  Optional[Any]		       =     in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    lowerCAmelCase__				:  int		       =     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      __lowerCAmelCase			(					UpperCamelCase      )     ->	Any:
  lowerCAmelCase__				:  Union[str, Any]		       =     ['''head.weight''', '''head.bias''']
  for k in ignore_keys:
    state_dict.pop(UpperCamelCase	,      UpperCamelCase      )



def      __lowerCAmelCase			(					UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase      )     ->	List[str]:
  lowerCAmelCase__				:  int		       =     dct.pop(UpperCamelCase      )
  lowerCAmelCase__				:  Any		       =     val



def      __lowerCAmelCase			(					)     ->	Optional[int]:
  lowerCAmelCase__				:  List[Any]		       =     '''http://images.cocodataset.org/val2017/000000039769.jpg'''
  lowerCAmelCase__				:  Dict		       =     Image.open(requests.get(UpperCamelCase	,      stream=UpperCamelCase      ).raw      )
  return im



@torch.no_grad()
def      __lowerCAmelCase			(					UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase=True      )     ->	Optional[Any]:
  lowerCAmelCase__				:  Optional[int]		       =     ViTConfig()
  # patch_size
  if model_name[-1] == "8":
    lowerCAmelCase__				:  int		       =     8
  # set labels if required
  if not base_model:
    lowerCAmelCase__				:  Tuple		       =     1000
    lowerCAmelCase__				:  Union[str, Any]		       =     '''huggingface/label-files'''
    lowerCAmelCase__				:  Optional[Any]		       =     '''imagenet-1k-id2label.json'''
    lowerCAmelCase__				:  Tuple		       =     json.load(open(hf_hub_download(UpperCamelCase	,      UpperCamelCase	,      repo_type='''dataset'''      )	,      '''r'''      )      )
    lowerCAmelCase__				:  Dict		       =     {int(UpperCamelCase      ): v for k, v in idalabel.items()}
    lowerCAmelCase__				:  int		       =     idalabel
    lowerCAmelCase__				:  List[str]		       =     {v: k for k, v in idalabel.items()}
  # size of the architecture
  if model_name in ["dino_vits8", "dino_vits16"]:
    lowerCAmelCase__				:  Any		       =     384
    lowerCAmelCase__				:  Optional[int]		       =     1536
    lowerCAmelCase__				:  int		       =     12
    lowerCAmelCase__				:  Any		       =     6

  # load original model from torch hub
  lowerCAmelCase__				:  Optional[int]		       =     torch.hub.load('''facebookresearch/dino:main'''	,      UpperCamelCase      )
  original_model.eval()

  # load state_dict of original model, remove and rename some keys
  lowerCAmelCase__				:  List[Any]		       =     original_model.state_dict()
  if base_model:
    remove_classification_head_(UpperCamelCase      )
  lowerCAmelCase__				:  Dict		       =     create_rename_keys(UpperCamelCase	,      base_model=UpperCamelCase      )
  for src, dest in rename_keys:
    rename_key(UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase      )
  read_in_q_k_v(UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase      )

  # load HuggingFace model
  if base_model:
    lowerCAmelCase__				:  Any		       =     ViTModel(UpperCamelCase	,      add_pooling_layer=UpperCamelCase      ).eval()
  else:
    lowerCAmelCase__				:  int		       =     ViTForImageClassification(UpperCamelCase      ).eval()
  model.load_state_dict(UpperCamelCase      )

  # Check outputs on an image, prepared by ViTImageProcessor
  lowerCAmelCase__				:  Union[str, Any]		       =     ViTImageProcessor()
  lowerCAmelCase__				:  Tuple		       =     image_processor(images=prepare_img()	,      return_tensors='''pt'''      )
  lowerCAmelCase__				:  List[str]		       =     encoding['''pixel_values''']
  lowerCAmelCase__				:  Optional[int]		       =     model(UpperCamelCase      )

  if base_model:
    lowerCAmelCase__				:  Tuple		       =     original_model(UpperCamelCase      )
    assert torch.allclose(UpperCamelCase	,      outputs.last_hidden_state[:, 0, :]	,      atol=1E-1      )
  else:
    lowerCAmelCase__				:  str		       =     original_model(UpperCamelCase      )
    assert logits.shape == outputs.logits.shape
    assert torch.allclose(UpperCamelCase	,      outputs.logits	,      atol=1E-3      )

  Path(UpperCamelCase      ).mkdir(exist_ok=UpperCamelCase      )
  print(F"""Saving model {model_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__":
     lowerCAmelCase_					    = argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--model_name""",
         default="""dino_vitb16""",
         type=str,
         help="""Name of the model trained with DINO 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(
         """--base_model""",
         action="""store_true""",
         help="""Whether to only convert the base model (no projection head weights).""",
     )

     parser.set_defaults(base_model=True)
     lowerCAmelCase_					    = parser.parse_args()
     convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)







from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow


if is_tf_available():
     import numpy as np
     import tensorflow as tf

     from transformers import TFCamembertModel






@require_tf
@require_sentencepiece
@require_tokenizers
class  _lowerCAmelCase     (  unittest.TestCase	):

      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Optional[Any]		       =     TFCamembertModel.from_pretrained('''jplu/tf-camembert-base'''		)

        lowerCAmelCase__				:  List[str]		       =     tf.convert_to_tensor(
            [[5, 121, 11, 660, 16, 730, 2_5543, 110, 83, 6]]      ,				dtype=tf.intaa      ,				)  # J'aime le camembert !"

        lowerCAmelCase__				:  str		       =     model(__UpperCAmelCase		)['''last_hidden_state''']
        lowerCAmelCase__				:  List[Any]		       =     tf.TensorShape((1, 10, 768)		)
        self.assertEqual(output.shape      ,				__UpperCAmelCase		)
        # compare the actual values for a slice.
        lowerCAmelCase__				:  int		       =     tf.convert_to_tensor(
            [[[-0.0254, 0.0235, 0.1027], [0.0606, -0.1811, -0.0418], [-0.1561, -0.1127, 0.2687]]]      ,				dtype=tf.floataa      ,				)
        # camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
        # camembert.eval()
        # expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

        self.assertTrue(np.allclose(output[:, :3, :3].numpy()      ,				expected_slice.numpy()      ,				atol=1e-4		)		)

