


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()
UpperCAmelCase__            =							logging.get_logger(__name__)
def     UpperCAmelCase_ (							__snake_case					,						__snake_case=False					,						__snake_case=False						)					->      str:







 """simple docstring"""


 _lowercase											='''backbone.''' if is_semantic else ''''''

 _lowercase											=[]
 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     UpperCAmelCase_ (							__snake_case					,						__snake_case					,						__snake_case=False					,						__snake_case=False						)					->      List[Any]:







 """simple docstring"""


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

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

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

  _lowercase											=gamma_a
  _lowercase											=gamma_a

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

   _lowercase											=table
   _lowercase											=index
def     UpperCAmelCase_ (							__snake_case					,						__snake_case					,						__snake_case						)					->      Tuple:







 """simple docstring"""


 _lowercase											=dct.pop(__snake_case						)
 _lowercase											=val
def     UpperCAmelCase_ (							)					->      Dict:







 """simple docstring"""


 _lowercase											='''http://images.cocodataset.org/val2017/000000039769.jpg'''
 _lowercase											=Image.open(requests.get(__snake_case					,						stream=__snake_case						).raw						)
 return im
@torch.no_grad()
def     UpperCAmelCase_ (							__snake_case					,						__snake_case						)					->      Tuple:







 """simple docstring"""


 _lowercase											=BeitConfig()
 _lowercase											=False
 _lowercase											=False
 _lowercase											='''huggingface/label-files'''
 # set config parameters based on URL
 if checkpoint_url[-9:-4] == "pt22k":
  # masked image modeling
  _lowercase											=True
  _lowercase											=True
  _lowercase											=True
 elif checkpoint_url[-9:-4] == "ft22k":
  # intermediate fine-tuning on ImageNet-22k
  _lowercase											=True
  _lowercase											=21841
  _lowercase											='''imagenet-22k-id2label.json'''
  _lowercase											=json.load(open(hf_hub_download(__snake_case					,						__snake_case					,						repo_type='''dataset'''						)					,						'''r'''						)						)
  _lowercase											={int(__snake_case						): 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[9205]
  del idalabel[15027]
  _lowercase											=idalabel
  _lowercase											={v: k for k, v in idalabel.items()}
 elif checkpoint_url[-8:-4] == "to1k":
  # fine-tuning on ImageNet-1k
  _lowercase											=True
  _lowercase											=1000
  _lowercase											='''imagenet-1k-id2label.json'''
  _lowercase											=json.load(open(hf_hub_download(__snake_case					,						__snake_case					,						repo_type='''dataset'''						)					,						'''r'''						)						)
  _lowercase											={int(__snake_case						): v for k, v in idalabel.items()}
  _lowercase											=idalabel
  _lowercase											={v: k for k, v in idalabel.items()}
  if "384" in checkpoint_url:
   _lowercase											=384
  if "512" in checkpoint_url:
   _lowercase											=512
 elif "ade20k" in checkpoint_url:
  # fine-tuning
  _lowercase											=True
  _lowercase											=150
  _lowercase											='''ade20k-id2label.json'''
  _lowercase											=json.load(open(hf_hub_download(__snake_case					,						__snake_case					,						repo_type='''dataset'''						)					,						'''r'''						)						)
  _lowercase											={int(__snake_case						): v for k, v in idalabel.items()}
  _lowercase											=idalabel
  _lowercase											={v: k for k, v in idalabel.items()}
  _lowercase											=640
  _lowercase											=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:
  _lowercase											=1024
  _lowercase											=4096
  _lowercase											=24
  _lowercase											=16
  if "ade20k" in checkpoint_url:
   _lowercase											=640
   _lowercase											=[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
 _lowercase											=torch.hub.load_state_dict_from_url(__snake_case					,						map_location='''cpu'''					,						check_hash=__snake_case						)
 _lowercase											=state_dict['''model'''] if '''ade20k''' not in checkpoint_url else state_dict['''state_dict''']

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

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

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

 _lowercase											=image_processor(images=__snake_case					,						return_tensors='''pt'''						)
 _lowercase											=encoding['''pixel_values''']

 _lowercase											=model(__snake_case						)
 _lowercase											=outputs.logits

 # verify logits
 _lowercase											=torch.Size([1, 1000]						)
 if checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k'''						):
  _lowercase											=torch.Size([1, 196, 8192]						)
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k'''						):
  _lowercase											=torch.Size([1, 196, 8192]						)
 elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft22k'''						):
  _lowercase											=torch.Size([1, 21841]						)
  _lowercase											=torch.tensor([2.22_88, 2.46_71, 0.73_95]						)
  _lowercase											=2397
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft22k'''						):
  _lowercase											=torch.Size([1, 21841]						)
  _lowercase											=torch.tensor([1.68_81, -0.27_87, 0.59_01]						)
  _lowercase											=2396
 elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft1k'''						):
  _lowercase											=torch.tensor([0.12_41, 0.07_98, -0.65_69]						)
  _lowercase											=285
 elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft22kto1k'''						):
  _lowercase											=torch.tensor([-1.23_85, -1.09_87, -1.01_08]						)
  _lowercase											=281
 elif checkpoint_url[:-4].endswith('''beit_base_patch16_384_pt22k_ft22kto1k'''						):
  _lowercase											=torch.tensor([-1.53_03, -0.94_84, -0.31_47]						)
  _lowercase											=761
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft1k'''						):
  _lowercase											=torch.tensor([0.46_10, -0.09_28, 0.20_86]						)
  _lowercase											=761
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft22kto1k'''						):
  _lowercase											=torch.tensor([-0.48_04, 0.62_57, -0.18_37]						)
  _lowercase											=761
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_384_pt22k_ft22kto1k'''						):
  _lowercase											=torch.tensor([[-0.51_22, 0.51_17, -0.21_13]]						)
  _lowercase											=761
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_512_pt22k_ft22kto1k'''						):
  _lowercase											=torch.tensor([-0.30_62, 0.72_61, 0.48_52]						)
  _lowercase											=761
 elif checkpoint_url[:-4].endswith('''beit_base_patch16_640_pt22k_ft22ktoade20k'''						):
  _lowercase											=(1, 150, 160, 160)
  _lowercase											=torch.tensor(
      [
          [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
          [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
          [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
      ]						)
 elif checkpoint_url[:-4].endswith('''beit_large_patch16_640_pt22k_ft22ktoade20k'''						):
  _lowercase											=(1, 150, 160, 160)
  _lowercase											=torch.tensor(
      [
          [[-4.33_05, -2.30_49, -3.01_61], [-2.95_91, -1.53_05, -2.22_51], [-3.41_98, -1.80_04, -2.90_62]],
          [[-5.89_22, -3.74_35, -4.39_78], [-4.20_63, -2.78_72, -3.47_55], [-4.27_91, -3.18_74, -4.16_81]],
          [[0.98_95, 4.34_67, 4.76_63], [4.24_76, 5.68_30, 6.15_18], [4.55_50, 6.24_95, 6.51_54]],
      ]						)
 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]					,						__snake_case					,						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]					,						__snake_case					,						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(__snake_case						).mkdir(exist_ok=__snake_case						)
 print(F"Saving model to {pytorch_dump_folder_path}"						)
 model.save_pretrained(__snake_case						)
 print(F"Saving image processor to {pytorch_dump_folder_path}"						)
 image_processor.save_pretrained(__snake_case						)


if __name__ == "__main__":
  UpperCAmelCase__            =							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.'''
  )
  UpperCAmelCase__            =							parser.parse_args()
  convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





def     UpperCAmelCase_ (							__snake_case						)					->      str:







 """simple docstring"""


 return " ".join(sorted(set(sentence.split()						)						)						)


if __name__ == "__main__":
  import doctest

  doctest.testmod()

