






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       :			Union[str, Any] ,	UpperCamelCase_       :			int ,	UpperCamelCase_       :			int=2 ,	UpperCamelCase_       :			Dict=56 ,	UpperCamelCase_       :			str=True ,	UpperCamelCase_       :			Dict=True ,	UpperCamelCase_       :			int=True ,	UpperCamelCase_       :			Dict=True ,	UpperCamelCase_       :			List[Any]=99 ,	UpperCamelCase_       :			Optional[Any]=32 ,	UpperCamelCase_       :			Tuple=2 ,	UpperCamelCase_       :			Optional[int]=2 ,	UpperCamelCase_       :			Dict=7 ,	UpperCamelCase_       :			List[Any]="gelu_new" ,	UpperCamelCase_       :			Union[str, Any]=0.1 ,	UpperCamelCase_       :			List[Any]=0.1 ,	UpperCamelCase_       :			Optional[int]=512 ,	UpperCamelCase_       :			str=16 ,	UpperCamelCase_       :			Any=2 ,	UpperCamelCase_       :			Any=0.02 ,	UpperCamelCase_       :			Any=4 ,	UpperCamelCase_       :			Union[str, Any]="block_sparse" ,	UpperCamelCase_       :			str=True ,	UpperCamelCase_       :			Dict=False ,	UpperCamelCase_       :			int=2 ,	UpperCamelCase_       :			List[str]=3 ,	):


         """simple docstring"""



         __A         =	parent
         __A         =	batch_size
         __A         =	seq_length
         __A         =	is_training
         __A         =	use_attention_mask
         __A         =	use_token_type_ids
         __A         =	use_labels
         __A         =	vocab_size
         __A         =	hidden_size
         __A         =	num_hidden_layers
         __A         =	num_attention_heads
         __A         =	intermediate_size
         __A         =	hidden_act
         __A         =	hidden_dropout_prob
         __A         =	attention_probs_dropout_prob
         __A         =	max_position_embeddings
         __A         =	type_vocab_size
         __A         =	type_sequence_label_size
         __A         =	initializer_range
         __A         =	num_choices

         __A         =	rescale_embeddings
         __A         =	attention_type
         __A         =	use_bias
         __A         =	block_size
         __A         =	num_random_blocks




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	ids_tensor([self.batch_size, self.seq_length] ,	self.vocab_size				)

         __A         =	None
         if self.use_attention_mask:
             __A         =	random_attention_mask([self.batch_size, self.seq_length]				)

         __A         =	None
         if self.use_token_type_ids:
             __A         =	ids_tensor([self.batch_size, self.seq_length] ,	self.type_vocab_size				)

         __A         =	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=UpperCamelCase_ ,	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 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	self.prepare_config_and_inputs()
         __A    ,	__A    ,	__A    ,	__A         =	config_and_inputs
         __A         =	{
             """input_ids""": input_ids,
             """token_type_ids""": token_type_ids,
             """attention_mask""": attention_mask,
         }
         return config, inputs_dict


@require_flax
class       __lowercase      (	lowercase_			,      unittest.TestCase  ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				(
         (
             FlaxBigBirdForCausalLM,
             FlaxBigBirdModel,
             FlaxBigBirdForPreTraining,
             FlaxBigBirdForMaskedLM,
             FlaxBigBirdForMultipleChoice,
             FlaxBigBirdForQuestionAnswering,
             FlaxBigBirdForSequenceClassification,
             FlaxBigBirdForTokenClassification,
         )
         if is_flax_available()
         else ()
     )

     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False




     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         __A         =	FlaxBigBirdModelTester(self				)




     @slow
     # copied from `test_modeling_flax_common` because it takes much longer than other models
     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         super().test_from_pretrained_save_pretrained()




     @slow
     # copied from `test_modeling_flax_common` because it takes much longer than other models
     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """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 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         super().test_no_automatic_init()




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


         """simple docstring"""



         super().test_hidden_states_output()




     @slow
     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         for model_class_name in self.all_model_classes:
             __A         =	model_class_name.from_pretrained("""google/bigbird-roberta-base"""				)
             self.assertIsNotNone(UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """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 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__				):
                 __A         =	self._prepare_for_class(UpperCamelCase_ ,	UpperCamelCase_				)
                 __A         =	model_class(UpperCamelCase_				)

                 @jax.jit
                 def model_jitted(UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Optional[int]=None ,	**UpperCamelCase_       :			List[str]				):
                     return model(input_ids=UpperCamelCase_ ,	attention_mask=UpperCamelCase_ ,	**UpperCamelCase_				)

                 with self.subTest("""JIT Enabled"""				):
                     __A         =	model_jitted(**UpperCamelCase_				).to_tuple()

                 with self.subTest("""JIT Disabled"""				):
                     with jax.disable_jit():
                         __A         =	model_jitted(**UpperCamelCase_				).to_tuple()

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





     def 		lowerCAmelCase_     (				self       :			Tuple ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			Optional[int]=1e-5 ,	UpperCamelCase_       :			List[str]="outputs" ,	UpperCamelCase_       :			Tuple=None				):


         """simple docstring"""



         if name.startswith("""outputs.attentions"""				):
             return
         else:
             super().check_pt_flax_outputs(UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_ ,	UpperCamelCase_				)












import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			str    )						->      str:





    """simple docstring"""

    if "cls_token" in name:
        __A         =	name.replace("""cls_token"""			,				"""vit.embeddings.cls_token"""    )
    if "mask_token" in name:
        __A         =	name.replace("""mask_token"""			,				"""decoder.mask_token"""    )
    if "decoder_pos_embed" in name:
        __A         =	name.replace("""decoder_pos_embed"""			,				"""decoder.decoder_pos_embed"""    )
    if "pos_embed" in name and "decoder" not in name:
        __A         =	name.replace("""pos_embed"""			,				"""vit.embeddings.position_embeddings"""    )
    if "patch_embed.proj" in name:
        __A         =	name.replace("""patch_embed.proj"""			,				"""vit.embeddings.patch_embeddings.projection"""    )
    if "patch_embed.norm" in name:
        __A         =	name.replace("""patch_embed.norm"""			,				"""vit.embeddings.norm"""    )
    if "decoder_blocks" in name:
        __A         =	name.replace("""decoder_blocks"""			,				"""decoder.decoder_layers"""    )
    if "blocks" in name:
        __A         =	name.replace("""blocks"""			,				"""vit.encoder.layer"""    )
    if "attn.proj" in name:
        __A         =	name.replace("""attn.proj"""			,				"""attention.output.dense"""    )
    if "attn" in name:
        __A         =	name.replace("""attn"""			,				"""attention.self"""    )
    if "norm1" in name:
        __A         =	name.replace("""norm1"""			,				"""layernorm_before"""    )
    if "norm2" in name:
        __A         =	name.replace("""norm2"""			,				"""layernorm_after"""    )
    if "mlp.fc1" in name:
        __A         =	name.replace("""mlp.fc1"""			,				"""intermediate.dense"""    )
    if "mlp.fc2" in name:
        __A         =	name.replace("""mlp.fc2"""			,				"""output.dense"""    )
    if "decoder_embed" in name:
        __A         =	name.replace("""decoder_embed"""			,				"""decoder.decoder_embed"""    )
    if "decoder_norm" in name:
        __A         =	name.replace("""decoder_norm"""			,				"""decoder.decoder_norm"""    )
    if "decoder_pred" in name:
        __A         =	name.replace("""decoder_pred"""			,				"""decoder.decoder_pred"""    )
    if "norm.weight" in name and "decoder" not in name:
        __A         =	name.replace("""norm.weight"""			,				"""vit.layernorm.weight"""    )
    if "norm.bias" in name and "decoder" not in name:
        __A         =	name.replace("""norm.bias"""			,				"""vit.layernorm.bias"""    )

    return name



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Any			,				__lowercase   :			Dict    )						->      str:





    """simple docstring"""

    for key in orig_state_dict.copy().keys():
        __A         =	orig_state_dict.pop(__lowercase    )

        if "qkv" in key:
            __A         =	key.split("""."""    )
            __A         =	int(key_split[1]    )
            if "decoder_blocks" in key:
                __A         =	config.decoder_hidden_size
                __A         =	"""decoder.decoder_layers."""
                if "weight" in key:
                    __A         =	val[:dim, :]
                    __A         =	val[dim : dim * 2, :]
                    __A         =	val[-dim:, :]
                elif "bias" in key:
                    __A         =	val[:dim]
                    __A         =	val[dim : dim * 2]
                    __A         =	val[-dim:]
            else:
                __A         =	config.hidden_size
                __A         =	"""vit.encoder.layer."""
                if "weight" in key:
                    __A         =	val[:dim, :]
                    __A         =	val[dim : dim * 2, :]
                    __A         =	val[-dim:, :]
                elif "bias" in key:
                    __A         =	val[:dim]
                    __A         =	val[dim : dim * 2]
                    __A         =	val[-dim:]

        else:
            __A         =	val

    return orig_state_dict



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Tuple			,				__lowercase   :			str    )						->      Optional[Any]:





    """simple docstring"""

    __A         =	ViTMAEConfig()
    if "large" in checkpoint_url:
        __A         =	1_0_2_4
        __A         =	4_0_9_6
        __A         =	2_4
        __A         =	1_6
    elif "huge" in checkpoint_url:
        __A         =	1_4
        __A         =	1_2_8_0
        __A         =	5_1_2_0
        __A         =	3_2
        __A         =	1_6

    __A         =	ViTMAEForPreTraining(__lowercase    )

    __A         =	torch.hub.load_state_dict_from_url(__lowercase			,				map_location="""cpu"""    )["""model"""]

    __A         =	ViTMAEImageProcessor(size=config.image_size    )

    __A         =	convert_state_dict(__lowercase			,				__lowercase    )

    model.load_state_dict(__lowercase    )
    model.eval()

    __A         =	"""https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"""

    __A         =	Image.open(requests.get(__lowercase			,				stream=__lowercase    ).raw    )
    __A         =	ViTMAEImageProcessor(size=config.image_size    )
    __A         =	image_processor(images=__lowercase			,				return_tensors="""pt"""    )

    # forward pass
    torch.manual_seed(2    )
    __A         =	model(**__lowercase    )
    __A         =	outputs.logits

    if "large" in checkpoint_url:
        __A         =	torch.tensor(
            [[-0.7_309, -0.7_128, -1.0_169], [-1.0_161, -0.9_058, -1.1_878], [-1.0_478, -0.9_411, -1.1_911]]    )
    elif "huge" in checkpoint_url:
        __A         =	torch.tensor(
            [[-1.1_599, -0.9_199, -1.2_221], [-1.1_952, -0.9_269, -1.2_307], [-1.2_143, -0.9_337, -1.2_262]]    )
    else:
        __A         =	torch.tensor(
            [[-0.9_192, -0.8_481, -1.1_259], [-1.1_349, -1.0_034, -1.2_599], [-1.1_757, -1.0_429, -1.2_726]]    )

    # verify logits
    assert torch.allclose(logits[0, :3, :3]			,				__lowercase			,				atol=1E-4    )

    print(f"Saving model to {pytorch_dump_folder_path}"    )
    model.save_pretrained(__lowercase    )

    print(f"Saving image processor to {pytorch_dump_folder_path}"    )
    image_processor.save_pretrained(__lowercase    )


if __name__ == "__main__":
  __a      :       List[str]				  =   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--checkpoint_url",
      default="https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth",
      type=str,
      help="URL of the checkpoint you'd like to convert.",
  )
  parser.add_argument(
      "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
  )

  __a      :       List[str]				  =   parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




