




import argparse
import json
from pathlib import Path

import requests
import torch
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
__SCREAMING_SNAKE_CASE     : Optional[int]      =logging.get_logger(__name__)

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[Any]  ,					lowerCamelCase:							Optional[int]=False  ,					lowerCamelCase:							List[Any]=False    )       ->      List[Any]:


      '''simple docstring'''




      lowerCamelCase__   	=					'backbone.' if is_semantic else ''

      lowerCamelCase__   	=					[]
      for i in range(config.num_hidden_layers    ):
            # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
            rename_keys.append((f"""{prefix}blocks.{i}.norm1.weight""", f"""beit.encoder.layer.{i}.layernorm_before.weight""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.norm1.bias""", f"""beit.encoder.layer.{i}.layernorm_before.bias""")    )
            rename_keys.append(
                (f"""{prefix}blocks.{i}.attn.proj.weight""", f"""beit.encoder.layer.{i}.attention.output.dense.weight""")    )
            rename_keys.append(
                (f"""{prefix}blocks.{i}.attn.proj.bias""", f"""beit.encoder.layer.{i}.attention.output.dense.bias""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.norm2.weight""", f"""beit.encoder.layer.{i}.layernorm_after.weight""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.norm2.bias""", f"""beit.encoder.layer.{i}.layernorm_after.bias""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc1.weight""", f"""beit.encoder.layer.{i}.intermediate.dense.weight""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc1.bias""", f"""beit.encoder.layer.{i}.intermediate.dense.bias""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc2.weight""", f"""beit.encoder.layer.{i}.output.dense.weight""")    )
            rename_keys.append((f"""{prefix}blocks.{i}.mlp.fc2.bias""", f"""beit.encoder.layer.{i}.output.dense.bias""")    )

      # projection layer + position embeddings
      rename_keys.extend(
          [
              (f"""{prefix}cls_token""", 'beit.embeddings.cls_token'),
              (f"""{prefix}patch_embed.proj.weight""", 'beit.embeddings.patch_embeddings.projection.weight'),
              (f"""{prefix}patch_embed.proj.bias""", 'beit.embeddings.patch_embeddings.projection.bias'),
          ]    )

      if has_lm_head:
            # mask token + shared relative position bias + layernorm
            rename_keys.extend(
                [
                    ('mask_token', 'beit.embeddings.mask_token'),
                    (
                        'rel_pos_bias.relative_position_bias_table',
                        'beit.encoder.relative_position_bias.relative_position_bias_table',
                    ),
                    (
                        'rel_pos_bias.relative_position_index',
                        'beit.encoder.relative_position_bias.relative_position_index',
                    ),
                    ('norm.weight', 'layernorm.weight'),
                    ('norm.bias', 'layernorm.bias'),
                ]    )
      elif is_semantic:
            # semantic segmentation classification heads
            rename_keys.extend(
                [
                    ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
                    ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
                    ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
                    ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
                ]    )
      else:
            # layernorm + classification head
            rename_keys.extend(
                [
                    ('fc_norm.weight', 'beit.pooler.layernorm.weight'),
                    ('fc_norm.bias', 'beit.pooler.layernorm.bias'),
                    ('head.weight', 'classifier.weight'),
                    ('head.bias', 'classifier.bias'),
                ]    )

      return rename_keys

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[int]  ,					lowerCamelCase:							int  ,					lowerCamelCase:							Dict=False  ,					lowerCamelCase:							Optional[int]=False    )       ->      List[Any]:


      '''simple docstring'''




      for i in range(config.num_hidden_layers    ):
            lowerCamelCase__   	=					'backbone.' if is_semantic else ''
            # queries, keys and values
            lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.attn.qkv.weight"""    )
            lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.attn.q_bias"""    )
            lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.attn.v_bias"""    )

            lowerCamelCase__   	=					in_proj_weight[
                : config.hidden_size, :
            ]
            lowerCamelCase__   	=					q_bias
            lowerCamelCase__   	=					in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            lowerCamelCase__   	=					in_proj_weight[
                -config.hidden_size :, :
            ]
            lowerCamelCase__   	=					v_bias

            # gamma_1 and gamma_2
            # we call them lambda because otherwise they are renamed when using .from_pretrained
            lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.gamma_1"""    )
            lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.gamma_2"""    )

            lowerCamelCase__   	=					gamma_a
            lowerCamelCase__   	=					gamma_a

            # relative_position bias table + index
            if not has_lm_head:
                  # each layer has its own relative position bias
                  lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.attn.relative_position_bias_table"""    )
                  lowerCamelCase__   	=					state_dict.pop(f"""{prefix}blocks.{i}.attn.relative_position_index"""    )

                  lowerCamelCase__   	=					table
                  lowerCamelCase__   	=					index

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[Any]  ,					lowerCamelCase:							Any  ,					lowerCamelCase:							int    )       ->      Any:


      '''simple docstring'''




      lowerCamelCase__   	=					dct.pop(lowerCamelCase    )
      lowerCamelCase__   	=					val

def       lowerCAmelCase_      (	)       ->      List[str]:


      '''simple docstring'''




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

@torch.no_grad()
def       lowerCAmelCase_      (	lowerCamelCase:							List[str]  ,					lowerCamelCase:							Optional[Any]    )       ->      Optional[Any]:


      '''simple docstring'''




      lowerCamelCase__   	=					BeitConfig()
      lowerCamelCase__   	=					False
      lowerCamelCase__   	=					False
      lowerCamelCase__   	=					'huggingface/label-files'
      # set config parameters based on URL
      if checkpoint_url[-9:-4] == "pt22k":
            # masked image modeling
            lowerCamelCase__   	=					True
            lowerCamelCase__   	=					True
            lowerCamelCase__   	=					True
      elif checkpoint_url[-9:-4] == "ft22k":
            # intermediate fine-tuning on ImageNet-22k
            lowerCamelCase__   	=					True
            lowerCamelCase__   	=					2_18_41
            lowerCamelCase__   	=					'imagenet-22k-id2label.json'
            lowerCamelCase__   	=					json.load(open(hf_hub_download(lowerCamelCase  ,					lowerCamelCase  ,					repo_type='dataset'    )  ,					'r'    )    )
            lowerCamelCase__   	=					{int(lowerCamelCase    ): v for k, v in idalabel.items()}
            # this dataset contains 21843 labels but the model only has 21841
            # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
            del idalabel[92_05]
            del idalabel[1_50_27]
            lowerCamelCase__   	=					idalabel
            lowerCamelCase__   	=					{v: k for k, v in idalabel.items()}
      elif checkpoint_url[-8:-4] == "to1k":
            # fine-tuning on ImageNet-1k
            lowerCamelCase__   	=					True
            lowerCamelCase__   	=					10_00
            lowerCamelCase__   	=					'imagenet-1k-id2label.json'
            lowerCamelCase__   	=					json.load(open(hf_hub_download(lowerCamelCase  ,					lowerCamelCase  ,					repo_type='dataset'    )  ,					'r'    )    )
            lowerCamelCase__   	=					{int(lowerCamelCase    ): v for k, v in idalabel.items()}
            lowerCamelCase__   	=					idalabel
            lowerCamelCase__   	=					{v: k for k, v in idalabel.items()}
            if "384" in checkpoint_url:
                  lowerCamelCase__   	=					3_84
            if "512" in checkpoint_url:
                  lowerCamelCase__   	=					5_12
      elif "ade20k" in checkpoint_url:
            # fine-tuning
            lowerCamelCase__   	=					True
            lowerCamelCase__   	=					1_50
            lowerCamelCase__   	=					'ade20k-id2label.json'
            lowerCamelCase__   	=					json.load(open(hf_hub_download(lowerCamelCase  ,					lowerCamelCase  ,					repo_type='dataset'    )  ,					'r'    )    )
            lowerCamelCase__   	=					{int(lowerCamelCase    ): v for k, v in idalabel.items()}
            lowerCamelCase__   	=					idalabel
            lowerCamelCase__   	=					{v: k for k, v in idalabel.items()}
            lowerCamelCase__   	=					6_40
            lowerCamelCase__   	=					True
      else:
            raise ValueError('Checkpoint not supported, URL should either end with \'pt22k\', \'ft22k\', \'to1k\' or \'ade20k\''    )

      # size of the architecture
      if "base" in checkpoint_url:
            pass
      elif "large" in checkpoint_url:
            lowerCamelCase__   	=					10_24
            lowerCamelCase__   	=					40_96
            lowerCamelCase__   	=					24
            lowerCamelCase__   	=					16
            if "ade20k" in checkpoint_url:
                  lowerCamelCase__   	=					6_40
                  lowerCamelCase__   	=					[7, 11, 15, 23]
      else:
            raise ValueError('Should either find \'base\' or \'large\' in checkpoint URL'    )

      # load state_dict of original model, remove and rename some keys
      lowerCamelCase__   	=					torch.hub.load_state_dict_from_url(lowerCamelCase  ,					map_location='cpu'  ,					check_hash=lowerCamelCase    )
      lowerCamelCase__   	=					state_dict['model'] if 'ade20k' not in checkpoint_url else state_dict['state_dict']

      lowerCamelCase__   	=					create_rename_keys(lowerCamelCase  ,					has_lm_head=lowerCamelCase  ,					is_semantic=lowerCamelCase    )
      for src, dest in rename_keys:
            rename_key(lowerCamelCase  ,					lowerCamelCase  ,					lowerCamelCase    )
      read_in_q_k_v(lowerCamelCase  ,					lowerCamelCase  ,					has_lm_head=lowerCamelCase  ,					is_semantic=lowerCamelCase    )
      if is_semantic:
            # add prefix to decoder keys
            for key, val in state_dict.copy().items():
                  lowerCamelCase__   	=					state_dict.pop(lowerCamelCase    )
                  if key.startswith('backbone.fpn'    ):
                        lowerCamelCase__   	=					key.replace('backbone.fpn'  ,					'fpn'    )
                  lowerCamelCase__   	=					val

    # load HuggingFace model
      if checkpoint_url[-9:-4] == "pt22k":
            lowerCamelCase__   	=					BeitForMaskedImageModeling(lowerCamelCase    )
      elif "ade20k" in checkpoint_url:
            lowerCamelCase__   	=					BeitForSemanticSegmentation(lowerCamelCase    )
      else:
            lowerCamelCase__   	=					BeitForImageClassification(lowerCamelCase    )
      model.eval()
      model.load_state_dict(lowerCamelCase    )

      # Check outputs on an image
      if is_semantic:
            lowerCamelCase__   	=					BeitImageProcessor(size=config.image_size  ,					do_center_crop=lowerCamelCase    )
            lowerCamelCase__   	=					load_dataset('hf-internal-testing/fixtures_ade20k'  ,					split='test'    )
            lowerCamelCase__   	=					Image.open(ds[0]['file']    )
      else:
            lowerCamelCase__   	=					BeitImageProcessor(
                size=config.image_size  ,					resample=PILImageResampling.BILINEAR  ,					do_center_crop=lowerCamelCase    )
            lowerCamelCase__   	=					prepare_img()

      lowerCamelCase__   	=					image_processor(images=lowerCamelCase  ,					return_tensors='pt'    )
      lowerCamelCase__   	=					encoding['pixel_values']

      lowerCamelCase__   	=					model(lowerCamelCase    )
      lowerCamelCase__   	=					outputs.logits

      # verify logits
      lowerCamelCase__   	=					torch.Size([1, 10_00]    )
      if checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k'    ):
            lowerCamelCase__   	=					torch.Size([1, 1_96, 81_92]    )
      elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k'    ):
            lowerCamelCase__   	=					torch.Size([1, 1_96, 81_92]    )
      elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft22k'    ):
            lowerCamelCase__   	=					torch.Size([1, 2_18_41]    )
            lowerCamelCase__   	=					torch.tensor([2.2288, 2.4671, 0.7395]    )
            lowerCamelCase__   	=					23_97
      elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft22k'    ):
            lowerCamelCase__   	=					torch.Size([1, 2_18_41]    )
            lowerCamelCase__   	=					torch.tensor([1.6881, -0.2787, 0.5901]    )
            lowerCamelCase__   	=					23_96
      elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft1k'    ):
            lowerCamelCase__   	=					torch.tensor([0.1241, 0.0798, -0.6569]    )
            lowerCamelCase__   	=					2_85
      elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft22kto1k'    ):
            lowerCamelCase__   	=					torch.tensor([-1.2385, -1.0987, -1.0108]    )
            lowerCamelCase__   	=					2_81
      elif checkpoint_url[:-4].endswith('beit_base_patch16_384_pt22k_ft22kto1k'    ):
            lowerCamelCase__   	=					torch.tensor([-1.5303, -0.9484, -0.3147]    )
            lowerCamelCase__   	=					7_61
      elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft1k'    ):
            lowerCamelCase__   	=					torch.tensor([0.4610, -0.0928, 0.2086]    )
            lowerCamelCase__   	=					7_61
      elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft22kto1k'    ):
            lowerCamelCase__   	=					torch.tensor([-0.4804, 0.6257, -0.1837]    )
            lowerCamelCase__   	=					7_61
      elif checkpoint_url[:-4].endswith('beit_large_patch16_384_pt22k_ft22kto1k'    ):
            lowerCamelCase__   	=					torch.tensor([[-0.5122, 0.5117, -0.2113]]    )
            lowerCamelCase__   	=					7_61
      elif checkpoint_url[:-4].endswith('beit_large_patch16_512_pt22k_ft22kto1k'    ):
            lowerCamelCase__   	=					torch.tensor([-0.3062, 0.7261, 0.4852]    )
            lowerCamelCase__   	=					7_61
      elif checkpoint_url[:-4].endswith('beit_base_patch16_640_pt22k_ft22ktoade20k'    ):
            lowerCamelCase__   	=					(1, 1_50, 1_60, 1_60)
            lowerCamelCase__   	=					torch.tensor(
                [
                    [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
                    [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
                    [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
                ]    )
      elif checkpoint_url[:-4].endswith('beit_large_patch16_640_pt22k_ft22ktoade20k'    ):
            lowerCamelCase__   	=					(1, 1_50, 1_60, 1_60)
            lowerCamelCase__   	=					torch.tensor(
                [
                    [[-4.3305, -2.3049, -3.0161], [-2.9591, -1.5305, -2.2251], [-3.4198, -1.8004, -2.9062]],
                    [[-5.8922, -3.7435, -4.3978], [-4.2063, -2.7872, -3.4755], [-4.2791, -3.1874, -4.1681]],
                    [[0.9895, 4.3467, 4.7663], [4.2476, 5.6830, 6.1518], [4.5550, 6.2495, 6.5154]],
                ]    )
      else:
            raise ValueError('Can\'t verify logits as model is not supported'    )

      assert logits.shape == expected_shape, "Shape of logits not as expected"
      if not has_lm_head:
            if is_semantic:
                  assert torch.allclose(
                      logits[0, :3, :3, :3]  ,					lowerCamelCase  ,					atol=1e-3    ), "First elements of logits not as expected"
            else:
                  print('Predicted class idx:'  ,					logits.argmax(-1    ).item()    )
                  assert torch.allclose(
                      logits[0, :3]  ,					lowerCamelCase  ,					atol=1e-3    ), "First elements of logits not as expected"
                  assert logits.argmax(-1    ).item() == expected_class_idx, "Predicted class index not as expected"

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


if __name__ == "__main__":
      __SCREAMING_SNAKE_CASE     : Union[str, Any]      =argparse.ArgumentParser()

      parser.add_argument(
          """--checkpoint_url""",
          default="""https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth""",
          type=str,
          help="""URL to the original PyTorch checkpoint (.pth file).""",
      )
      parser.add_argument(
          """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
      )
      __SCREAMING_SNAKE_CASE     : Optional[int]      =parser.parse_args()
      convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





def       lowerCAmelCase_      (	lowerCamelCase:							str    )       ->      int:


      '''simple docstring'''




      lowerCamelCase__   	=					str(lowerCamelCase    ).strip()
      if not bin_string:
            raise ValueError('Empty string was passed to the function'    )
      lowerCamelCase__   	=					bin_string[0] == '-'
      if is_negative:
            lowerCamelCase__   	=					bin_string[1:]
      if not all(char in '01' for char in bin_string    ):
            raise ValueError('Non-binary value was passed to the function'    )
      lowerCamelCase__   	=					0
      for char in bin_string:
            lowerCamelCase__   	=					2 * decimal_number + int(lowerCamelCase    )
      return -decimal_number if is_negative else decimal_number


if __name__ == "__main__":
      from doctest import testmod

      testmod()