




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			:	Optional[Any]			 =     logging.get_logger(__name__)







def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_=False		)     -> Union[str, Any]:



       """simple docstring"""

       A__   				=					[]
       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"
              A__   				=					[(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 SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_=False		)     -> List[Any]:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> Tuple:



       """simple docstring"""

       A__   				=					['''head.weight''', '''head.bias''']
       for k in ignore_keys:
              state_dict.pop(a__				,       a__		)







def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_		)     -> Dict:



       """simple docstring"""

       A__   				=					dct.pop(a__		)
       A__   				=					val







def SCREAMING_SNAKE_CASE       ( )     -> List[Any]:



       """simple docstring"""

       A__   				=					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
       A__   				=					Image.open(requests.get(a__				,       stream=a__		).raw		)
       return im







@torch.no_grad()
def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_=True		)     -> Dict:



       """simple docstring"""

       A__   				=					ViTConfig()
       # patch_size
       if model_name[-1] == "8":
              A__   				=					8
       # set labels if required
       if not base_model:
              A__   				=					1_000
              A__   				=					'''huggingface/label-files'''
              A__   				=					'''imagenet-1k-id2label.json'''
              A__   				=					json.load(open(hf_hub_download(a__				,       a__				,       repo_type='''dataset'''		)				,       '''r'''		)		)
              A__   				=					{int(a__		): v for k, v in idalabel.items()}
              A__   				=					idalabel
              A__   				=					{v: k for k, v in idalabel.items()}
       # size of the architecture
       if model_name in ["dino_vits8", "dino_vits16"]:
              A__   				=					384
              A__   				=					1_536
              A__   				=					12
              A__   				=					6

       # load original model from torch hub
       A__   				=					torch.hub.load('''facebookresearch/dino:main'''				,       a__		)
       original_model.eval()

       # load state_dict of original model, remove and rename some keys
       A__   				=					original_model.state_dict()
       if base_model:
              remove_classification_head_(a__		)
       A__   				=					create_rename_keys(a__				,       base_model=a__		)
       for src, dest in rename_keys:
              rename_key(a__				,       a__				,       a__		)
       read_in_q_k_v(a__				,       a__				,       a__		)

       # load HuggingFace model
       if base_model:
              A__   				=					ViTModel(a__				,       add_pooling_layer=a__		).eval()
       else:
              A__   				=					ViTForImageClassification(a__		).eval()
       model.load_state_dict(a__		)

       # Check outputs on an image, prepared by ViTImageProcessor
       A__   				=					ViTImageProcessor()
       A__   				=					image_processor(images=prepare_img()				,       return_tensors='''pt'''		)
       A__   				=					encoding['''pixel_values''']
       A__   				=					model(a__		)

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

       Path(a__		).mkdir(exist_ok=a__		)
       print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""		)
       model.save_pretrained(a__		)
       print(f"""Saving image processor to {pytorch_dump_folder_path}"""		)
       image_processor.save_pretrained(a__		)


if __name__ == "__main__":
  _lowerCamelCase			:	Dict			 =     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			:	str			 =     parser.parse_args()
  convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase			:	Tuple							 =     logging.get_logger(__name__)

_lowerCamelCase			:	Tuple							 =     {
    """tiiuae/falcon-40b""": """https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json""",
    """tiiuae/falcon-7b""": """https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json""",
}


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   '''falcon'''
    UpperCAmelCase__													=   ['''past_key_values''']


    def __init__(    self   :			Optional[Any]      ,				UpperCAmelCase__   :			str=65_024      ,				UpperCAmelCase__   :			List[Any]=4_544      ,				UpperCAmelCase__   :			Optional[Any]=32      ,				UpperCAmelCase__   :			Dict=71      ,				UpperCAmelCase__   :			Any=1e-5      ,				UpperCAmelCase__   :			Optional[int]=0.02      ,				UpperCAmelCase__   :			Optional[int]=True      ,				UpperCAmelCase__   :			Union[str, Any]=0.0      ,				UpperCAmelCase__   :			List[str]=0.0      ,				UpperCAmelCase__   :			str=None      ,				UpperCAmelCase__   :			Dict=False      ,				UpperCAmelCase__   :			List[str]=False      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=True      ,				UpperCAmelCase__   :			Optional[Any]=False      ,				UpperCAmelCase__   :			Dict=11      ,				UpperCAmelCase__   :			str=11      ,				**UpperCAmelCase__   :			Dict      ,				)		->Any:





           '''simple docstring'''

           A__   				=					vocab_size
           # Backward compatibility with n_embed kwarg
           A__   				=					kwargs.pop('''n_embed'''      ,				UpperCAmelCase__)
           A__   				=					hidden_size if n_embed is None else n_embed
           A__   				=					num_hidden_layers
           A__   				=					num_attention_heads
           A__   				=					layer_norm_epsilon
           A__   				=					initializer_range
           A__   				=					use_cache
           A__   				=					hidden_dropout
           A__   				=					attention_dropout

           A__   				=					bos_token_id
           A__   				=					eos_token_id
           A__   				=					num_attention_heads if num_kv_heads is None else num_kv_heads
           A__   				=					alibi
           A__   				=					new_decoder_architecture
           A__   				=					multi_query  # Ignored when new_decoder_architecture is True
           A__   				=					parallel_attn
           A__   				=					bias

           super().__init__(bos_token_id=UpperCAmelCase__      ,				eos_token_id=UpperCAmelCase__      ,				**UpperCAmelCase__)


    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->int:





           '''simple docstring'''

           return self.hidden_size // self.num_attention_heads





    @property
    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Dict:





           '''simple docstring'''

           return not self.alibi


