





import argparse
import json
from pathlib import Path

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

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


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




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]	,     __UpperCAmelCase:    List[Any]=False   )   ->			Dict:
    UpperCamelCase__							:       str      =   []
    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'''),
                ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
                ('''pre_logits.fc.bias''', '''pooler.dense.bias'''),
            ]   )

        # if just the base model, we should remove "vit" from all keys that start with "vit"
        UpperCamelCase__							:       Dict      =   [(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:    List[str]	,     __UpperCAmelCase:    List[str]	,     __UpperCAmelCase:    str=False   )   ->			Tuple:
    for i in range(config.num_hidden_layers   ):
        if base_model:
            UpperCamelCase__							:       int      =   ''''''
        else:
            UpperCamelCase__							:       str      =   '''vit.'''
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        UpperCamelCase__							:       Any      =   state_dict.pop(f"blocks.{i}.attn.qkv.weight"   )
        UpperCamelCase__							:       Optional[int]      =   state_dict.pop(f"blocks.{i}.attn.qkv.bias"   )
        # next, add query, keys and values (in that order) to the state dict
        UpperCamelCase__							:       Any      =   in_proj_weight[
            : config.hidden_size, :
        ]
        UpperCamelCase__							:       Any      =   in_proj_bias[: config.hidden_size]
        UpperCamelCase__							:       Optional[Any]      =   in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        UpperCamelCase__							:       Tuple      =   in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        UpperCamelCase__							:       int      =   in_proj_weight[
            -config.hidden_size :, :
        ]
        UpperCamelCase__							:       Tuple      =   in_proj_bias[-config.hidden_size :]




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[str]   )   ->			List[Any]:
    UpperCamelCase__							:       str      =   ['''head.weight''', '''head.bias''']
    for k in ignore_keys:
        state_dict.pop(__UpperCAmelCase	,     __UpperCAmelCase   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Union[str, Any]	,     __UpperCAmelCase:    List[str]	,     __UpperCAmelCase:    int   )   ->			str:
    UpperCamelCase__							:       Optional[Any]      =   dct.pop(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[Any]      =   val




def 	lowerCAmelCase_					(						)   ->			int:
    UpperCamelCase__							:       Any      =   '''http://images.cocodataset.org/val2017/000000039769.jpg'''
    UpperCamelCase__							:       Union[str, Any]      =   Image.open(requests.get(__UpperCAmelCase	,     stream=__UpperCAmelCase   ).raw   )
    return im




@torch.no_grad()
def 	lowerCAmelCase_					(						__UpperCAmelCase:    Tuple	,     __UpperCAmelCase:    Any   )   ->			Tuple:
    UpperCamelCase__							:       Dict      =   ViTConfig()
    UpperCamelCase__							:       Tuple      =   False
    # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
    if vit_name[-5:] == "in21k":
        UpperCamelCase__							:       Any      =   True
        UpperCamelCase__							:       str      =   int(vit_name[-12:-10]   )
        UpperCamelCase__							:       List[Any]      =   int(vit_name[-9:-6]   )
    else:
        UpperCamelCase__							:       Tuple      =   1000
        UpperCamelCase__							:       Optional[int]      =   '''huggingface/label-files'''
        UpperCamelCase__							:       str      =   '''imagenet-1k-id2label.json'''
        UpperCamelCase__							:       List[str]      =   json.load(open(hf_hub_download(__UpperCAmelCase	,     __UpperCAmelCase	,     repo_type='''dataset'''   )	,     '''r'''   )   )
        UpperCamelCase__							:       Tuple      =   {int(__UpperCAmelCase   ): v for k, v in idalabel.items()}
        UpperCamelCase__							:       List[Any]      =   idalabel
        UpperCamelCase__							:       Union[str, Any]      =   {v: k for k, v in idalabel.items()}
        UpperCamelCase__							:       Dict      =   int(vit_name[-6:-4]   )
        UpperCamelCase__							:       Optional[int]      =   int(vit_name[-3:]   )
    # size of the architecture
    if "deit" in vit_name:
        if vit_name[9:].startswith('''tiny'''   ):
            UpperCamelCase__							:       List[str]      =   192
            UpperCamelCase__							:       int      =   768
            UpperCamelCase__							:       Any      =   12
            UpperCamelCase__							:       Optional[int]      =   3
        elif vit_name[9:].startswith('''small'''   ):
            UpperCamelCase__							:       List[Any]      =   384
            UpperCamelCase__							:       Union[str, Any]      =   1536
            UpperCamelCase__							:       Optional[int]      =   12
            UpperCamelCase__							:       List[Any]      =   6
        else:
            pass
    else:
        if vit_name[4:].startswith('''small'''   ):
            UpperCamelCase__							:       str      =   768
            UpperCamelCase__							:       int      =   2304
            UpperCamelCase__							:       List[str]      =   8
            UpperCamelCase__							:       Optional[int]      =   8
        elif vit_name[4:].startswith('''base'''   ):
            pass
        elif vit_name[4:].startswith('''large'''   ):
            UpperCamelCase__							:       Optional[Any]      =   1024
            UpperCamelCase__							:       int      =   4096
            UpperCamelCase__							:       List[Any]      =   24
            UpperCamelCase__							:       Any      =   16
        elif vit_name[4:].startswith('''huge'''   ):
            UpperCamelCase__							:       Any      =   1280
            UpperCamelCase__							:       int      =   5120
            UpperCamelCase__							:       Union[str, Any]      =   32
            UpperCamelCase__							:       Optional[int]      =   16

    # load original model from timm
    UpperCamelCase__							:       Union[str, Any]      =   timm.create_model(__UpperCAmelCase	,     pretrained=__UpperCAmelCase   )
    timm_model.eval()

    # load state_dict of original model, remove and rename some keys
    UpperCamelCase__							:       Dict      =   timm_model.state_dict()
    if base_model:
        remove_classification_head_(__UpperCAmelCase   )
    UpperCamelCase__							:       Tuple      =   create_rename_keys(__UpperCAmelCase	,     __UpperCAmelCase   )
    for src, dest in rename_keys:
        rename_key(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
    read_in_q_k_v(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )

    # load HuggingFace model
    if vit_name[-5:] == "in21k":
        UpperCamelCase__							:       List[Any]      =   ViTModel(__UpperCAmelCase   ).eval()
    else:
        UpperCamelCase__							:       List[Any]      =   ViTForImageClassification(__UpperCAmelCase   ).eval()
    model.load_state_dict(__UpperCAmelCase   )

    # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
    if "deit" in vit_name:
        UpperCamelCase__							:       List[Any]      =   DeiTImageProcessor(size=config.image_size   )
    else:
        UpperCamelCase__							:       Any      =   ViTImageProcessor(size=config.image_size   )
    UpperCamelCase__							:       Optional[int]      =   image_processor(images=prepare_img()	,     return_tensors='''pt'''   )
    UpperCamelCase__							:       Any      =   encoding['''pixel_values''']
    UpperCamelCase__							:       Dict      =   model(__UpperCAmelCase   )

    if base_model:
        UpperCamelCase__							:       List[Any]      =   timm_model.forward_features(__UpperCAmelCase   )
        assert timm_pooled_output.shape == outputs.pooler_output.shape
        assert torch.allclose(__UpperCAmelCase	,     outputs.pooler_output	,     atol=1e-3   )
    else:
        UpperCamelCase__							:       Optional[int]      =   timm_model(__UpperCAmelCase   )
        assert timm_logits.shape == outputs.logits.shape
        assert torch.allclose(__UpperCAmelCase	,     outputs.logits	,     atol=1e-3   )

    Path(__UpperCAmelCase   ).mkdir(exist_ok=__UpperCAmelCase   )
    print(f"Saving model {vit_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(
      '--vit_name',
      default='vit_base_patch16_224',
      type=str,
      help='Name of the ViT timm model you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )

  UpperCAmelCase_  							=  parser.parse_args()
  convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)








import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax

  from transformers.models.big_bird.modeling_flax_big_bird import (
      FlaxBigBirdForCausalLM,
      FlaxBigBirdForMaskedLM,
      FlaxBigBirdForMultipleChoice,
      FlaxBigBirdForPreTraining,
      FlaxBigBirdForQuestionAnswering,
      FlaxBigBirdForSequenceClassification,
      FlaxBigBirdForTokenClassification,
      FlaxBigBirdModel,
  )






class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   def __init__(					self,	__magic_name__,	__magic_name__=2,	__magic_name__=56,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=True,	__magic_name__=99,	__magic_name__=32,	__magic_name__=2,	__magic_name__=2,	__magic_name__=7,	__magic_name__="gelu_new",	__magic_name__=0.1,	__magic_name__=0.1,	__magic_name__=512,	__magic_name__=16,	__magic_name__=2,	__magic_name__=0.02,	__magic_name__=4,	__magic_name__="block_sparse",	__magic_name__=True,	__magic_name__=False,	__magic_name__=2,	__magic_name__=3,	)				->   Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   parent
       UpperCamelCase__							:       List[Any]      =   batch_size
       UpperCamelCase__							:       str      =   seq_length
       UpperCamelCase__							:       List[str]      =   is_training
       UpperCamelCase__							:       List[Any]      =   use_attention_mask
       UpperCamelCase__							:       Any      =   use_token_type_ids
       UpperCamelCase__							:       Tuple      =   use_labels
       UpperCamelCase__							:       Any      =   vocab_size
       UpperCamelCase__							:       Optional[int]      =   hidden_size
       UpperCamelCase__							:       Union[str, Any]      =   num_hidden_layers
       UpperCamelCase__							:       Optional[Any]      =   num_attention_heads
       UpperCamelCase__							:       List[str]      =   intermediate_size
       UpperCamelCase__							:       List[str]      =   hidden_act
       UpperCamelCase__							:       Union[str, Any]      =   hidden_dropout_prob
       UpperCamelCase__							:       Union[str, Any]      =   attention_probs_dropout_prob
       UpperCamelCase__							:       int      =   max_position_embeddings
       UpperCamelCase__							:       List[Any]      =   type_vocab_size
       UpperCamelCase__							:       Tuple      =   type_sequence_label_size
       UpperCamelCase__							:       Optional[Any]      =   initializer_range
       UpperCamelCase__							:       Dict      =   num_choices

       UpperCamelCase__							:       Dict      =   rescale_embeddings
       UpperCamelCase__							:       Tuple      =   attention_type
       UpperCamelCase__							:       Union[str, Any]      =   use_bias
       UpperCamelCase__							:       Union[str, Any]      =   block_size
       UpperCamelCase__							:       Optional[int]      =   num_random_blocks




   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       UpperCamelCase__							:       int      =   ids_tensor([self.batch_size, self.seq_length],	self.vocab_size							)

       UpperCamelCase__							:       int      =   None
       if self.use_attention_mask:
           UpperCamelCase__							:       List[str]      =   random_attention_mask([self.batch_size, self.seq_length]							)

       UpperCamelCase__							:       Optional[int]      =   None
       if self.use_token_type_ids:
           UpperCamelCase__							:       List[str]      =   ids_tensor([self.batch_size, self.seq_length],	self.type_vocab_size							)

       UpperCamelCase__							:       str      =   BigBirdConfig(
           vocab_size=self.vocab_size,	hidden_size=self.hidden_size,	num_hidden_layers=self.num_hidden_layers,	num_attention_heads=self.num_attention_heads,	intermediate_size=self.intermediate_size,	hidden_act=self.hidden_act,	hidden_dropout_prob=self.hidden_dropout_prob,	attention_probs_dropout_prob=self.attention_probs_dropout_prob,	max_position_embeddings=self.max_position_embeddings,	type_vocab_size=self.type_vocab_size,	is_decoder=__magic_name__,	initializer_range=self.initializer_range,	attention_type=self.attention_type,	block_size=self.block_size,	num_random_blocks=self.num_random_blocks,	use_bias=self.use_bias,	rescale_embeddings=self.rescale_embeddings,	)

       return config, input_ids, token_type_ids, attention_mask




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       List[str]      =   self.prepare_config_and_inputs()
       UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__     ,UpperCamelCase__							:       Tuple      =   config_and_inputs
       UpperCamelCase__							:       Union[str, Any]      =   {
           '''input_ids''': input_ids,
           '''token_type_ids''': token_type_ids,
           '''attention_mask''': attention_mask,
       }
       return config, inputs_dict






@require_flax
class 					lowercase__ (			__lowerCamelCase    ,       unittest.TestCase       ):

   '''simple docstring'''


   a		:							Any     	=	(
       (
           FlaxBigBirdForCausalLM,
           FlaxBigBirdModel,
           FlaxBigBirdForPreTraining,
           FlaxBigBirdForMaskedLM,
           FlaxBigBirdForMultipleChoice,
           FlaxBigBirdForQuestionAnswering,
           FlaxBigBirdForSequenceClassification,
           FlaxBigBirdForTokenClassification,
       )
       if is_flax_available()
       else ()
   )

   a		:							List[str]     	=	False
   a		:							str     	=	False




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   FlaxBigBirdModelTester(self							)




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_from_pretrained_save_pretrained()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_from_pretrained_with_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Dict:






       """simple docstring"""
       super().test_no_automatic_init()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       super().test_hidden_states_output()




   @slow
   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       for model_class_name in self.all_model_classes:
           UpperCamelCase__							:       List[str]      =   model_class_name.from_pretrained('''google/bigbird-roberta-base'''							)
           self.assertIsNotNone(__magic_name__							)




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       if self.test_attn_probs:
           super().test_attention_outputs()




   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__     ,UpperCamelCase__							:       Dict      =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__							):
               UpperCamelCase__							:       Union[str, Any]      =   self._prepare_for_class(__magic_name__,	__magic_name__							)
               UpperCamelCase__							:       Tuple      =   model_class(__magic_name__							)

               @jax.jit
               def model_jitted(__magic_name__,	__magic_name__=None,	**__magic_name__							):
                   return model(input_ids=__magic_name__,	attention_mask=__magic_name__,	**__magic_name__							)

               with self.subTest('''JIT Enabled'''							):
                   UpperCamelCase__							:       str      =   model_jitted(**__magic_name__							).to_tuple()

               with self.subTest('''JIT Disabled'''							):
                   with jax.disable_jit():
                       UpperCamelCase__							:       Dict      =   model_jitted(**__magic_name__							).to_tuple()

               self.assertEqual(len(__magic_name__							),	len(__magic_name__							)							)
               for jitted_output, output in zip(__magic_name__,	__magic_name__							):
                   self.assertEqual(jitted_output.shape,	output.shape							)




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__=1E-5,	__magic_name__="outputs",	__magic_name__=None							)				->   List[Any]:






       """simple docstring"""
       # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
       # an effort was done to return `attention_probs` (yet to be verified).
       if name.startswith('''outputs.attentions'''							):
           return
       else:
           super().check_pt_flax_outputs(__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__,	__magic_name__							)

