

"""simple docstring"""




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()
__UpperCAmelCase =				logging.get_logger(__name__)
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase=False      ):
	'''simple docstring'''





	UpperCAmelCase__					:  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"
		UpperCAmelCase__					:  Optional[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      ):
	'''simple docstring'''





	for i in range(config.num_hidden_layers      ):
		if base_model:
			UpperCAmelCase__					:  Tuple					   =					""""""
		else:
			UpperCAmelCase__					:  Any					   =					"""vit."""
		# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
		UpperCAmelCase__					:  Optional[Any]					   =					state_dict.pop(F"blocks.{i}.attn.qkv.weight"      )
		UpperCAmelCase__					:  Dict					   =					state_dict.pop(F"blocks.{i}.attn.qkv.bias"      )
		# next, add query, keys and values (in that order) to the state dict
		UpperCAmelCase__					:  Dict					   =					in_proj_weight[
		    : config.hidden_size, :
		]
		UpperCAmelCase__					:  str					   =					in_proj_bias[: config.hidden_size]
		UpperCAmelCase__					:  Optional[Any]					   =					in_proj_weight[
		    config.hidden_size : config.hidden_size * 2, :
		]
		UpperCAmelCase__					:  Optional[Any]					   =					in_proj_bias[
		    config.hidden_size : config.hidden_size * 2
		]
		UpperCAmelCase__					:  Union[str, Any]					   =					in_proj_weight[
		    -config.hidden_size :, :
		]
		UpperCAmelCase__					:  str					   =					in_proj_bias[-config.hidden_size :]
def   lowerCAmelCase						(					__UpperCamelCase      ):
	'''simple docstring'''





	UpperCAmelCase__					:  int					   =					["""head.weight""", """head.bias"""]
	for k in ignore_keys:
		state_dict.pop(__UpperCamelCase							,    __UpperCamelCase      )
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase      ):
	'''simple docstring'''





	UpperCAmelCase__					:  Any					   =					dct.pop(__UpperCamelCase      )
	UpperCAmelCase__					:  List[str]					   =					val
def   lowerCAmelCase						(					):
	'''simple docstring'''





	UpperCAmelCase__					:  int					   =					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
	UpperCAmelCase__					:  int					   =					Image.open(requests.get(__UpperCamelCase							,    stream=__UpperCamelCase      ).raw      )
	return im
@torch.no_grad()
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase=True      ):
	'''simple docstring'''





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

	# load original model from torch hub
	UpperCAmelCase__					:  Any					   =					torch.hub.load("""facebookresearch/dino:main"""							,    __UpperCamelCase      )
	original_model.eval()

	# load state_dict of original model, remove and rename some keys
	UpperCAmelCase__					:  Union[str, Any]					   =					original_model.state_dict()
	if base_model:
		remove_classification_head_(__UpperCamelCase      )
	UpperCAmelCase__					:  List[str]					   =					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:
		UpperCAmelCase__					:  Dict					   =					ViTModel(__UpperCamelCase							,    add_pooling_layer=__UpperCamelCase      ).eval()
	else:
		UpperCAmelCase__					:  Tuple					   =					ViTForImageClassification(__UpperCamelCase      ).eval()
	model.load_state_dict(__UpperCamelCase      )

	# Check outputs on an image, prepared by ViTImageProcessor
	UpperCAmelCase__					:  Optional[int]					   =					ViTImageProcessor()
	UpperCAmelCase__					:  int					   =					image_processor(images=prepare_img()							,    return_tensors="""pt"""      )
	UpperCAmelCase__					:  Optional[Any]					   =					encoding["""pixel_values"""]
	UpperCAmelCase__					:  str					   =					model(__UpperCamelCase      )

	if base_model:
		UpperCAmelCase__					:  List[str]					   =					original_model(__UpperCamelCase      )
		assert torch.allclose(__UpperCamelCase							,    outputs.last_hidden_state[:, 0, :]							,    atol=1e-1      )
	else:
		UpperCAmelCase__					:  List[Any]					   =					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__":
						__UpperCAmelCase =				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)
						__UpperCAmelCase =				parser.parse_args()
						convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)








"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase =				{
    'configuration_roberta': ['ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RobertaConfig', 'RobertaOnnxConfig'],
    'tokenization_roberta': ['RobertaTokenizer'],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				['RobertaTokenizerFast']

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST',
          'RobertaForCausalLM',
          'RobertaForMaskedLM',
          'RobertaForMultipleChoice',
          'RobertaForQuestionAnswering',
          'RobertaForSequenceClassification',
          'RobertaForTokenClassification',
          'RobertaModel',
          'RobertaPreTrainedModel',
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFRobertaForCausalLM',
          'TFRobertaForMaskedLM',
          'TFRobertaForMultipleChoice',
          'TFRobertaForQuestionAnswering',
          'TFRobertaForSequenceClassification',
          'TFRobertaForTokenClassification',
          'TFRobertaMainLayer',
          'TFRobertaModel',
          'TFRobertaPreTrainedModel',
      ]

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'FlaxRobertaForCausalLM',
          'FlaxRobertaForMaskedLM',
          'FlaxRobertaForMultipleChoice',
          'FlaxRobertaForQuestionAnswering',
          'FlaxRobertaForSequenceClassification',
          'FlaxRobertaForTokenClassification',
          'FlaxRobertaModel',
          'FlaxRobertaPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaOnnxConfig
      from .tokenization_roberta import RobertaTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_roberta_fast import RobertaTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_roberta import (
                ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                RobertaForCausalLM,
                RobertaForMaskedLM,
                RobertaForMultipleChoice,
                RobertaForQuestionAnswering,
                RobertaForSequenceClassification,
                RobertaForTokenClassification,
                RobertaModel,
                RobertaPreTrainedModel,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_roberta import (
                TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRobertaForCausalLM,
                TFRobertaForMaskedLM,
                TFRobertaForMultipleChoice,
                TFRobertaForQuestionAnswering,
                TFRobertaForSequenceClassification,
                TFRobertaForTokenClassification,
                TFRobertaMainLayer,
                TFRobertaModel,
                TFRobertaPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_roberta import (
                FlaxRobertaForCausalLM,
                FlaxRobertaForMaskedLM,
                FlaxRobertaForMultipleChoice,
                FlaxRobertaForQuestionAnswering,
                FlaxRobertaForSequenceClassification,
                FlaxRobertaForTokenClassification,
                FlaxRobertaModel,
                FlaxRobertaPreTrainedModel,
            )

else:
      import sys

      __UpperCAmelCase =				_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





