
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     _a					(						lowerCamelCase,	lowerCamelCase=False							):
 lowerCamelCase		:	Union[str, 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		:	Optional[Any]							=      [(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     _a					(						lowerCamelCase,	lowerCamelCase,	lowerCamelCase=False							):
 for i in range(config.num_hidden_layers							):
  if base_model:
   lowerCamelCase		:	Union[str, Any]							=      """"""
  else:
   lowerCamelCase		:	Union[str, Any]							=      """vit."""
  # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
  lowerCamelCase		:	Dict							=      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		:	Optional[Any]							=      in_proj_weight[
      : config.hidden_size, :
  ]
  lowerCamelCase		:	Any							=      in_proj_bias[: config.hidden_size]
  lowerCamelCase		:	List[Any]							=      in_proj_weight[
      config.hidden_size : config.hidden_size * 2, :
  ]
  lowerCamelCase		:	Optional[int]							=      in_proj_bias[
      config.hidden_size : config.hidden_size * 2
  ]
  lowerCamelCase		:	List[str]							=      in_proj_weight[
      -config.hidden_size :, :
  ]
  lowerCamelCase		:	Tuple							=      in_proj_bias[-config.hidden_size :]





def     _a					(						lowerCamelCase							):
 lowerCamelCase		:	str							=      ["""head.weight""", """head.bias"""]
 for k in ignore_keys:
  state_dict.pop(lowerCamelCase,	lowerCamelCase							)





def     _a					(						lowerCamelCase,	lowerCamelCase,	lowerCamelCase							):
 lowerCamelCase		:	Optional[int]							=      dct.pop(lowerCamelCase							)
 lowerCamelCase		:	Union[str, Any]							=      val





def     _a					(						):
 lowerCamelCase		:	Any							=      """http://images.cocodataset.org/val2017/000000039769.jpg"""
 lowerCamelCase		:	Any							=      Image.open(requests.get(lowerCamelCase,	stream=lowerCamelCase							).raw							)
 return im





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

 # load original model from torch hub
 lowerCamelCase		:	List[str]							=      torch.hub.load("""facebookresearch/dino:main""",	lowerCamelCase							)
 original_model.eval()

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

 # load HuggingFace model
 if base_model:
  lowerCamelCase		:	Union[str, Any]							=      ViTModel(lowerCamelCase,	add_pooling_layer=lowerCamelCase							).eval()
 else:
  lowerCamelCase		:	int							=      ViTForImageClassification(lowerCamelCase							).eval()
 model.load_state_dict(lowerCamelCase							)

 # Check outputs on an image, prepared by ViTImageProcessor
 lowerCamelCase		:	List[Any]							=      ViTImageProcessor()
 lowerCamelCase		:	str							=      image_processor(images=prepare_img(),	return_tensors="""pt"""							)
 lowerCamelCase		:	Tuple							=      encoding["""pixel_values"""]
 lowerCamelCase		:	List[str]							=      model(lowerCamelCase							)

 if base_model:
  lowerCamelCase		:	List[str]							=      original_model(lowerCamelCase							)
  assert torch.allclose(lowerCamelCase,	outputs.last_hidden_state[:, 0, :],	atol=1e-1							)
 else:
  lowerCamelCase		:	List[str]							=      original_model(lowerCamelCase							)
  assert logits.shape == outputs.logits.shape
  assert torch.allclose(lowerCamelCase,	outputs.logits,	atol=1e-3							)

 Path(lowerCamelCase							).mkdir(exist_ok=lowerCamelCase							)
 print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''							)
 model.save_pretrained(lowerCamelCase							)
 print(F'''Saving image processor to {pytorch_dump_folder_path}'''							)
 image_processor.save_pretrained(lowerCamelCase							)


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)






import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


_lowerCamelCase					   =5_0_0_0_0_0

_lowerCamelCase   ,  _lowerCamelCase					   =os.path.split(__file__)
_lowerCamelCase					   =os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))





@get_duration
def     _a					(						lowerCamelCase,	**lowerCamelCase							):
 lowerCamelCase		:	Optional[Any]							=      dataset.map(**lowerCamelCase							)





@get_duration
def     _a					(						lowerCamelCase,	**lowerCamelCase							):
 lowerCamelCase		:	Optional[Any]							=      dataset.filter(**lowerCamelCase							)





def     _a					(						):
 lowerCamelCase		:	Optional[Any]							=      {"""num examples""": SPEED_TEST_N_EXAMPLES}
 with tempfile.TemporaryDirectory() as tmp_dir:
  lowerCamelCase		:	Any							=      datasets.Features({"""text""": datasets.Value("""string"""							), """numbers""": datasets.Value("""float32"""							)}							)
  lowerCamelCase		:	Tuple							=      generate_example_dataset(
      os.path.join(lowerCamelCase,	"""dataset.arrow"""							),	lowerCamelCase,	num_examples=lowerCamelCase							)

  lowerCamelCase		:	Tuple							=      transformers.AutoTokenizer.from_pretrained("""bert-base-cased""",	use_fast=lowerCamelCase							)

  def tokenize(lowerCamelCase							):
   return tokenizer(examples["""text"""]							)

  lowerCamelCase		:	List[str]							=      map(lowerCamelCase							)

  lowerCamelCase		:	int							=      map(lowerCamelCase,	batched=lowerCamelCase							)

  lowerCamelCase		:	int							=      map(lowerCamelCase,	function=lambda lowerCamelCase							: None,	batched=lowerCamelCase							)

  with dataset.formatted_as(type="""numpy"""							):
   lowerCamelCase		:	Optional[int]							=      map(lowerCamelCase,	function=lambda lowerCamelCase							: None,	batched=lowerCamelCase							)

  with dataset.formatted_as(type="""pandas"""							):
   lowerCamelCase		:	List[str]							=      map(lowerCamelCase,	function=lambda lowerCamelCase							: None,	batched=lowerCamelCase							)

  with dataset.formatted_as(type="""torch""",	columns="""numbers"""							):
   lowerCamelCase		:	List[str]							=      map(lowerCamelCase,	function=lambda lowerCamelCase							: None,	batched=lowerCamelCase							)

  with dataset.formatted_as(type="""tensorflow""",	columns="""numbers"""							):
   lowerCamelCase		:	Optional[int]							=      map(lowerCamelCase,	function=lambda lowerCamelCase							: None,	batched=lowerCamelCase							)

  lowerCamelCase		:	int							=      map(lowerCamelCase,	function=lowerCamelCase,	batched=lowerCamelCase							)

  lowerCamelCase		:	Union[str, Any]							=      filter(lowerCamelCase							)

  # Activate later when tokenizer support batched inputs
  # with dataset.formatted_as(type='numpy'):
  #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

 with open(lowerCamelCase,	"""wb"""							) as f:
  f.write(json.dumps(lowerCamelCase							).encode("""utf-8"""							)							)


if __name__ == "__main__":  # useful to run the profiler
      benchmark_map_filter()




