



import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402


a_								=       get_tests_dir("""fixtures/dummy-config.json""")



class 			__lowerCAmelCase  (						unittest.TestCase ):




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     0




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    self.assertIsNotNone(transformers.models.auto.__spec__					)
    self.assertIsNotNone(importlib.util.find_spec('''transformers.models.auto'''					)					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     AutoConfig.from_pretrained('''bert-base-uncased'''					)
    self.assertIsInstance(__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase					)
    self.assertIsInstance(__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase					)
    self.assertIsInstance(__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     AutoConfig.for_model('''roberta'''					)
    self.assertIsInstance(__UpperCAmelCase						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    with tempfile.TemporaryDirectory() as tmp_dir:
       # This model name contains bert and roberta, but roberta ends up being picked.
       __lowerCamelCase    				=     os.path.join(__UpperCAmelCase						,					'''fake-roberta'''					)
       os.makedirs(__UpperCAmelCase						,					exist_ok=__UpperCAmelCase					)
       with open(os.path.join(__UpperCAmelCase						,					'''config.json'''					)						,					'''w'''					) as f:
          f.write(json.dumps({}					)					)
       __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase					)
       self.assertEqual(type(__UpperCAmelCase					)						,					__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    try:
       AutoConfig.register('''custom'''						,					__UpperCAmelCase					)
       # Wrong model type will raise an error
       with self.assertRaises(__UpperCAmelCase					):
          AutoConfig.register('''model'''						,					__UpperCAmelCase					)
       # Trying to register something existing in the Transformers library will raise an error
       with self.assertRaises(__UpperCAmelCase					):
          AutoConfig.register('''bert'''						,					__UpperCAmelCase					)

       # Now that the config is registered, it can be used as any other config with the auto-API
       __lowerCamelCase    				=     CustomConfig()
       with tempfile.TemporaryDirectory() as tmp_dir:
          config.save_pretrained(__UpperCAmelCase					)
          __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase					)
          self.assertIsInstance(__UpperCAmelCase						,					__UpperCAmelCase					)

    finally:
       if "custom" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["custom"]




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    with self.assertRaisesRegex(
        __UpperCAmelCase						,					'''bert-base is not a local folder and is not a valid model identifier'''					):
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''bert-base'''					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    with self.assertRaisesRegex(
        __UpperCAmelCase						,					r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''					):
       __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase						,					revision='''aaaaaa'''					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    with self.assertRaisesRegex(
        __UpperCAmelCase						,					'''hf-internal-testing/no-config-test-repo does not appear to have a file named config.json.'''						,					):
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/no-config-test-repo'''					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    # If remote code is not set, we will time out when asking whether to load the model.
    with self.assertRaises(__UpperCAmelCase					):
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''					)
    # If remote code is disabled, we can't load this config.
    with self.assertRaises(__UpperCAmelCase					):
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''						,					trust_remote_code=__UpperCAmelCase					)

    __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''						,					trust_remote_code=__UpperCAmelCase					)
    self.assertEqual(config.__class__.__name__						,					'''NewModelConfig'''					)

    # Test config can be reloaded.
    with tempfile.TemporaryDirectory() as tmp_dir:
       config.save_pretrained(__UpperCAmelCase					)
       __lowerCamelCase    				=     AutoConfig.from_pretrained(__UpperCAmelCase						,					trust_remote_code=__UpperCAmelCase					)
    self.assertEqual(reloaded_config.__class__.__name__						,					'''NewModelConfig'''					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''



    class 			__lowerCAmelCase  (						lowerCAmelCase__ ):
     lowerCAmelCase__							       =	"""new-model"""

    try:
       AutoConfig.register('''new-model'''						,					__UpperCAmelCase					)
       # If remote code is not set, the default is to use local
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''					)
       self.assertEqual(config.__class__.__name__						,					'''NewModelConfigLocal'''					)

       # If remote code is disabled, we load the local one.
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''						,					trust_remote_code=__UpperCAmelCase					)
       self.assertEqual(config.__class__.__name__						,					'''NewModelConfigLocal'''					)

       # If remote is enabled, we load from the Hub
       __lowerCamelCase    				=     AutoConfig.from_pretrained('''hf-internal-testing/test_dynamic_model'''						,					trust_remote_code=__UpperCAmelCase					)
       self.assertEqual(config.__class__.__name__						,					'''NewModelConfig'''					)

    finally:
       if "new-model" in CONFIG_MAPPING._extra_content:
          del CONFIG_MAPPING._extra_content["new-model"]











import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor





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

   return name





def 		a__      (		_UpperCamelCase      :			Tuple			,_UpperCamelCase      :			Dict				):
   for key in orig_state_dict.copy().keys():
      __lowerCamelCase    				=     orig_state_dict.pop(_UpperCamelCase				)

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

      else:
         __lowerCamelCase    				=     val

   return orig_state_dict





def 		a__      (		_UpperCamelCase      :			Dict			,_UpperCamelCase      :			Any				):
   __lowerCamelCase    				=     ViTMAEConfig()
   if "large" in checkpoint_url:
      __lowerCamelCase    				=     10_24
      __lowerCamelCase    				=     40_96
      __lowerCamelCase    				=     24
      __lowerCamelCase    				=     16
   elif "huge" in checkpoint_url:
      __lowerCamelCase    				=     14
      __lowerCamelCase    				=     12_80
      __lowerCamelCase    				=     51_20
      __lowerCamelCase    				=     32
      __lowerCamelCase    				=     16

   __lowerCamelCase    				=     ViTMAEForPreTraining(_UpperCamelCase				)

   __lowerCamelCase    				=     torch.hub.load_state_dict_from_url(_UpperCamelCase			,map_location='''cpu'''				)['''model''']

   __lowerCamelCase    				=     ViTMAEImageProcessor(size=config.image_size				)

   __lowerCamelCase    				=     convert_state_dict(_UpperCamelCase			,_UpperCamelCase				)

   model.load_state_dict(_UpperCamelCase				)
   model.eval()

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

   __lowerCamelCase    				=     Image.open(requests.get(_UpperCamelCase			,stream=_UpperCamelCase				).raw				)
   __lowerCamelCase    				=     ViTMAEImageProcessor(size=config.image_size				)
   __lowerCamelCase    				=     image_processor(images=_UpperCamelCase			,return_tensors='''pt'''				)

   # forward pass
   torch.manual_seed(2				)
   __lowerCamelCase    				=     model(**_UpperCamelCase				)
   __lowerCamelCase    				=     outputs.logits

   if "large" in checkpoint_url:
      __lowerCamelCase    				=     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:
      __lowerCamelCase    				=     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:
      __lowerCamelCase    				=     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]			,_UpperCamelCase			,atol=1e-4				)

   print(F"""Saving model 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__":
  a_								=       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_								=       parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






