






import warnings
from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


UpperCAmelCase__							  =						logging.get_logger(__name__)

UpperCAmelCase__							  =						{
    "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/config.json",
    # See all BART models at https://huggingface.co/models?filter=bart
}






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      __snake_case	    =    '''bart'''
      __snake_case	    =    ['''past_key_values''']
      __snake_case	    =    {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}







      def __init__( self :		List[str]					,   __UpperCAmelCase :		Any=50_265					,   __UpperCAmelCase :		Dict=1_024					,   __UpperCAmelCase :		Optional[int]=12					,   __UpperCAmelCase :		str=4_096					,   __UpperCAmelCase :		List[str]=16					,   __UpperCAmelCase :		Any=12					,   __UpperCAmelCase :		List[Any]=4_096					,   __UpperCAmelCase :		List[str]=16					,   __UpperCAmelCase :		Tuple=0.0					,   __UpperCAmelCase :		List[Any]=0.0					,   __UpperCAmelCase :		Tuple="gelu"					,   __UpperCAmelCase :		int=1_024					,   __UpperCAmelCase :		Union[str, Any]=0.1					,   __UpperCAmelCase :		Optional[Any]=0.0					,   __UpperCAmelCase :		Optional[int]=0.0					,   __UpperCAmelCase :		int=0.02					,   __UpperCAmelCase :		List[str]=0.0					,   __UpperCAmelCase :		str=False					,   __UpperCAmelCase :		Dict=True					,   __UpperCAmelCase :		int=3					,   __UpperCAmelCase :		Tuple=1					,   __UpperCAmelCase :		int=0					,   __UpperCAmelCase :		Dict=2					,   __UpperCAmelCase :		Dict=True					,   __UpperCAmelCase :		Optional[Any]=2					,   __UpperCAmelCase :		str=2					,   **__UpperCAmelCase :		int					,   )       ->List[str]:




       """simple docstring"""


       a						    = vocab_size
       a						    = max_position_embeddings
       a						    = d_model
       a						    = encoder_ffn_dim
       a						    = encoder_layers
       a						    = encoder_attention_heads
       a						    = decoder_ffn_dim
       a						    = decoder_layers
       a						    = decoder_attention_heads
       a						    = dropout
       a						    = attention_dropout
       a						    = activation_dropout
       a						    = activation_function
       a						    = init_std
       a						    = encoder_layerdrop
       a						    = decoder_layerdrop
       a						    = classifier_dropout
       a						    = use_cache
       a						    = encoder_layers
       a						    = scale_embedding  # scale factor will be sqrt(d_model) if True

       super().__init__(
           num_labels=__UpperCAmelCase					,   pad_token_id=__UpperCAmelCase					,   bos_token_id=__UpperCAmelCase					,   eos_token_id=__UpperCAmelCase					,   is_encoder_decoder=__UpperCAmelCase					,   decoder_start_token_id=__UpperCAmelCase					,   forced_eos_token_id=__UpperCAmelCase					,   **__UpperCAmelCase					,   )

       # ensure backward compatibility for BART CNN models
       if self.forced_bos_token_id is None and kwargs.get('''force_bos_token_to_be_generated'''					,   __UpperCAmelCase    ):
        a						    = self.bos_token_id
        warnings.warn(
            F"""Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. """
            '''The config can simply be saved and uploaded again to be fixed.'''    )






class 						lowercase_   (      lowercase  ):






      '''simple docstring'''




      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Mapping[str, Mapping[int, str]]:




       """simple docstring"""


       if self.task in ["default", "seq2seq-lm"]:
        a						    = OrderedDict(
            [
                ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
            ]    )

        if self.use_past:
         a						    = {0: '''batch'''}
         a						    = {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
        else:
         a						    = {0: '''batch''', 1: '''decoder_sequence'''}
         a						    = {0: '''batch''', 1: '''decoder_sequence'''}

        if self.use_past:
         self.fill_with_past_key_values_(__UpperCAmelCase					,   direction='''inputs'''    )
       elif self.task == "causal-lm":
        # TODO: figure this case out.
        a						    = OrderedDict(
            [
                ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
            ]    )
        if self.use_past:
         a    ,	a						    = self.num_layers
         for i in range(__UpperCAmelCase    ):
          a						    = {0: '''batch''', 2: '''past_sequence + sequence'''}
          a						    = {0: '''batch''', 2: '''past_sequence + sequence'''}
       else:
        a						    = OrderedDict(
            [
                ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                ('''decoder_input_ids''', {0: '''batch''', 1: '''decoder_sequence'''}),
                ('''decoder_attention_mask''', {0: '''batch''', 1: '''decoder_sequence'''}),
            ]    )

       return common_inputs







      @property
      def  __lowerCAmelCase   ( self :		Optional[int]    )       ->Mapping[str, Mapping[int, str]]:




       """simple docstring"""


       if self.task in ["default", "seq2seq-lm"]:
        a						    = super().outputs
       else:
        a						    = super(__UpperCAmelCase					,   self    ).outputs
        if self.use_past:
         a    ,	a						    = self.num_layers
         for i in range(__UpperCAmelCase    ):
          a						    = {0: '''batch''', 2: '''past_sequence + sequence'''}
          a						    = {0: '''batch''', 2: '''past_sequence + sequence'''}
       return common_outputs







      def  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		PreTrainedTokenizer					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		bool = False					,   __UpperCAmelCase :		Optional[TensorType] = None					,   )       ->Mapping[str, Any]:




       """simple docstring"""


       a						    = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )

       # Generate decoder inputs
       a						    = seq_length if not self.use_past else 1
       a						    = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )
       a						    = {F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
       a						    = dict(**__UpperCAmelCase					,   **__UpperCAmelCase    )

       if self.use_past:
        if not is_torch_available():
         raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''    )
        else:
         import torch
        a    ,	a						    = common_inputs['''input_ids'''].shape
        a						    = common_inputs['''decoder_input_ids'''].shape[1]
        a    ,	a						    = self.num_attention_heads
        a						    = (
            batch,
            num_encoder_attention_heads,
            encoder_seq_length,
            self._config.hidden_size // num_encoder_attention_heads,
        )
        a						    = decoder_seq_length + 3
        a						    = (
            batch,
            num_decoder_attention_heads,
            decoder_past_length,
            self._config.hidden_size // num_decoder_attention_heads,
        )

        a						    = torch.cat(
            [common_inputs['''decoder_attention_mask'''], torch.ones(__UpperCAmelCase					,   __UpperCAmelCase    )]					,   dim=1    )

        a						    = []
        # If the number of encoder and decoder layers are present in the model configuration, both are considered
        a    ,	a						    = self.num_layers
        a						    = min(__UpperCAmelCase					,   __UpperCAmelCase    )
        a						    = max(__UpperCAmelCase					,   __UpperCAmelCase    ) - min_num_layers
        a						    = '''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

        for _ in range(__UpperCAmelCase    ):
         common_inputs["past_key_values"].append(
             (
                 torch.zeros(__UpperCAmelCase    ),
                 torch.zeros(__UpperCAmelCase    ),
                 torch.zeros(__UpperCAmelCase    ),
                 torch.zeros(__UpperCAmelCase    ),
             )    )
        # TODO: test this.
        a						    = encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
        for _ in range(__UpperCAmelCase					,   __UpperCAmelCase    ):
         common_inputs["past_key_values"].append((torch.zeros(__UpperCAmelCase    ), torch.zeros(__UpperCAmelCase    ))    )
       return common_inputs







      def  __lowerCAmelCase   ( self :		Optional[Any]					,   __UpperCAmelCase :		PreTrainedTokenizer					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		bool = False					,   __UpperCAmelCase :		Optional[TensorType] = None					,   )       ->Mapping[str, Any]:




       """simple docstring"""


       a						    = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
           __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )

       if self.use_past:
        if not is_torch_available():
         raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''    )
        else:
         import torch
        a    ,	a						    = common_inputs['''input_ids'''].shape
        # Not using the same length for past_key_values
        a						    = seqlen + 2
        a    ,	a						    = self.num_layers
        a    ,	a						    = self.num_attention_heads
        a						    = (
            batch,
            num_encoder_attention_heads,
            past_key_values_length,
            self._config.hidden_size // num_encoder_attention_heads,
        )

        a						    = common_inputs['''attention_mask'''].dtype
        a						    = torch.cat(
            [common_inputs['''attention_mask'''], torch.ones(__UpperCAmelCase					,   __UpperCAmelCase					,   dtype=__UpperCAmelCase    )]					,   dim=1    )
        a						    = [
            (torch.zeros(__UpperCAmelCase    ), torch.zeros(__UpperCAmelCase    )) for _ in range(__UpperCAmelCase    )
        ]
       return common_inputs







      def  __lowerCAmelCase   ( self :		List[Any]					,   __UpperCAmelCase :		PreTrainedTokenizer					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		bool = False					,   __UpperCAmelCase :		Optional[TensorType] = None					,   )       ->Mapping[str, Any]:




       """simple docstring"""


       a						    = compute_effective_axis_dimension(
           __UpperCAmelCase					,   fixed_dimension=OnnxConfig.default_fixed_batch					,   num_token_to_add=0    )

       # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
       a						    = tokenizer.num_special_tokens_to_add(__UpperCAmelCase    )
       a						    = compute_effective_axis_dimension(
           __UpperCAmelCase					,   fixed_dimension=OnnxConfig.default_fixed_sequence					,   num_token_to_add=__UpperCAmelCase    )

       # Generate dummy inputs according to compute batch and sequence
       a						    = [''' '''.join([tokenizer.unk_token]    ) * seq_length] * batch_size
       a						    = dict(tokenizer(__UpperCAmelCase					,   return_tensors=__UpperCAmelCase    )    )
       return common_inputs







      def  __lowerCAmelCase   ( self :		Optional[Any]					,   __UpperCAmelCase :		PreTrainedTokenizer					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		int = -1					,   __UpperCAmelCase :		bool = False					,   __UpperCAmelCase :		Optional[TensorType] = None					,   )       ->Mapping[str, Any]:




       """simple docstring"""


       if self.task in ["default", "seq2seq-lm"]:
        a						    = self._generate_dummy_inputs_for_default_and_seqaseq_lm(
            __UpperCAmelCase					,   batch_size=__UpperCAmelCase					,   seq_length=__UpperCAmelCase					,   is_pair=__UpperCAmelCase					,   framework=__UpperCAmelCase    )

       elif self.task == "causal-lm":
        a						    = self._generate_dummy_inputs_for_causal_lm(
            __UpperCAmelCase					,   batch_size=__UpperCAmelCase					,   seq_length=__UpperCAmelCase					,   is_pair=__UpperCAmelCase					,   framework=__UpperCAmelCase    )
       else:
        a						    = self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            __UpperCAmelCase					,   batch_size=__UpperCAmelCase					,   seq_length=__UpperCAmelCase					,   is_pair=__UpperCAmelCase					,   framework=__UpperCAmelCase    )

       return common_inputs







      def  __lowerCAmelCase   ( self :		Optional[Any]					,   __UpperCAmelCase :		Optional[Any]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		int    )       ->str:




       """simple docstring"""


       if self.task in ["default", "seq2seq-lm"]:
        a						    = super()._flatten_past_key_values_(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )
       else:
        a						    = super(__UpperCAmelCase					,   self    )._flatten_past_key_values_(
            __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )










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 		_a       ( a					:List[str]      ,				a					:Tuple=False					)				->     List[Any]:
 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 		_a       ( a					:List[str]      ,				a					:int      ,				a					:Tuple=False					)				->     Any:
 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 		_a       ( a					:List[Any]					)				->     Dict:
 a						    = ['''head.weight''', '''head.bias''']
 for k in ignore_keys:
  state_dict.pop(a      ,				a					)
def 		_a       ( a					:Union[str, Any]      ,				a					:str      ,				a					:List[Any]					)				->     str:
 a						    = dct.pop(a					)
 a						    = val
def 		_a       ( )				->     Optional[Any]:
 a						    = '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 a						    = Image.open(requests.get(a      ,				stream=a					).raw					)
 return im
@torch.no_grad()
def 		_a       ( a					:Tuple      ,				a					:str      ,				a					:Optional[int]=True					)				->     Dict:
 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__":
     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)


