
'''simple docstring'''


import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)



def  __snake_case     (    _UpperCAmelCase					:       str,      _UpperCAmelCase					:       List[Any]=False):
   UpperCamelCase       =  []
   for i in range(config.num_hidden_layers):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((f'module.blocks.{i}.norm1.weight', f'vit.encoder.layer.{i}.layernorm_before.weight'))
      rename_keys.append((f'module.blocks.{i}.norm1.bias', f'vit.encoder.layer.{i}.layernorm_before.bias'))
      rename_keys.append(
          (f'module.blocks.{i}.attn.proj.weight', f'vit.encoder.layer.{i}.attention.output.dense.weight'))
      rename_keys.append((f'module.blocks.{i}.attn.proj.bias', f'vit.encoder.layer.{i}.attention.output.dense.bias'))
      rename_keys.append((f'module.blocks.{i}.norm2.weight', f'vit.encoder.layer.{i}.layernorm_after.weight'))
      rename_keys.append((f'module.blocks.{i}.norm2.bias', f'vit.encoder.layer.{i}.layernorm_after.bias'))
      rename_keys.append((f'module.blocks.{i}.mlp.fc1.weight', f'vit.encoder.layer.{i}.intermediate.dense.weight'))
      rename_keys.append((f'module.blocks.{i}.mlp.fc1.bias', f'vit.encoder.layer.{i}.intermediate.dense.bias'))
      rename_keys.append((f'module.blocks.{i}.mlp.fc2.weight', f'vit.encoder.layer.{i}.output.dense.weight'))
      rename_keys.append((f'module.blocks.{i}.mlp.fc2.bias', f'vit.encoder.layer.{i}.output.dense.bias'))

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ('''module.cls_token''', '''vit.embeddings.cls_token'''),
           ('''module.patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
           ('''module.patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
           ('''module.pos_embed''', '''vit.embeddings.position_embeddings'''),
       ])

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ('''module.norm.weight''', '''layernorm.weight'''),
              ('''module.norm.bias''', '''layernorm.bias'''),
          ])

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



def  __snake_case     (    _UpperCAmelCase					:       Optional[int]):
   UpperCamelCase       =  ['''head.weight''', '''head.bias''']
   for k in ignore_keys:
      state_dict.pop(_UpperCAmelCase,      _UpperCAmelCase)



def  __snake_case     (    _UpperCAmelCase					:       int):
   # projection head is used in the self-supervised pre-training in MSN,
   # for downstream task it's not needed.
   UpperCamelCase       =  [
       '''module.fc.fc1.weight''',
       '''module.fc.fc1.bias''',
       '''module.fc.bn1.weight''',
       '''module.fc.bn1.bias''',
       '''module.fc.bn1.running_mean''',
       '''module.fc.bn1.running_var''',
       '''module.fc.bn1.num_batches_tracked''',
       '''module.fc.fc2.weight''',
       '''module.fc.fc2.bias''',
       '''module.fc.bn2.weight''',
       '''module.fc.bn2.bias''',
       '''module.fc.bn2.running_mean''',
       '''module.fc.bn2.running_var''',
       '''module.fc.bn2.num_batches_tracked''',
       '''module.fc.fc3.weight''',
       '''module.fc.fc3.bias''',
   ]
   for k in ignore_keys:
      state_dict.pop(_UpperCAmelCase,      _UpperCAmelCase)



def  __snake_case     (    _UpperCAmelCase					:       List[Any],      _UpperCAmelCase					:       str,      _UpperCAmelCase					:       int):
   UpperCamelCase       =  dct.pop(_UpperCAmelCase)
   UpperCamelCase       =  val



def  __snake_case     (    _UpperCAmelCase					:       Optional[int],      _UpperCAmelCase					:       Dict):
   UpperCamelCase       =  ViTMSNConfig()
   UpperCamelCase       =  1000

   UpperCamelCase       =  '''datasets/huggingface/label-files'''
   UpperCamelCase       =  '''imagenet-1k-id2label.json'''
   UpperCamelCase       =  json.load(open(hf_hub_download(_UpperCAmelCase,      _UpperCAmelCase),      '''r'''))
   UpperCamelCase       =  {int(_UpperCAmelCase): v for k, v in idalabel.items()}
   UpperCamelCase       =  idalabel
   UpperCamelCase       =  {v: k for k, v in idalabel.items()}

   if "s16" in checkpoint_url:
      UpperCamelCase       =  384
      UpperCamelCase       =  1536
      UpperCamelCase       =  6
   elif "l16" in checkpoint_url:
      UpperCamelCase       =  1024
      UpperCamelCase       =  4096
      UpperCamelCase       =  24
      UpperCamelCase       =  16
      UpperCamelCase       =  0.1
   elif "b4" in checkpoint_url:
      UpperCamelCase       =  4
   elif "l7" in checkpoint_url:
      UpperCamelCase       =  7
      UpperCamelCase       =  1024
      UpperCamelCase       =  4096
      UpperCamelCase       =  24
      UpperCamelCase       =  16
      UpperCamelCase       =  0.1

   UpperCamelCase       =  ViTMSNModel(_UpperCAmelCase)

   UpperCamelCase       =  torch.hub.load_state_dict_from_url(_UpperCAmelCase,      map_location='''cpu''')['''target_encoder''']

   UpperCamelCase       =  ViTImageProcessor(size=config.image_size)

   remove_projection_head(_UpperCAmelCase)
   UpperCamelCase       =  create_rename_keys(_UpperCAmelCase,      base_model=_UpperCAmelCase)

   for src, dest in rename_keys:
      rename_key(_UpperCAmelCase,      _UpperCAmelCase,      _UpperCAmelCase)
   read_in_q_k_v(_UpperCAmelCase,      _UpperCAmelCase,      base_model=_UpperCAmelCase)

   model.load_state_dict(_UpperCAmelCase)
   model.eval()

   UpperCamelCase       =  '''http://images.cocodataset.org/val2017/000000039769.jpg'''

   UpperCamelCase       =  Image.open(requests.get(_UpperCAmelCase,      stream=_UpperCAmelCase).raw)
   UpperCamelCase       =  ViTImageProcessor(
       size=config.image_size,      image_mean=_UpperCAmelCase,      image_std=_UpperCAmelCase)
   UpperCamelCase       =  image_processor(images=_UpperCAmelCase,      return_tensors='''pt''')

   # forward pass
   torch.manual_seed(2)
   UpperCamelCase       =  model(**_UpperCAmelCase)
   UpperCamelCase       =  outputs.last_hidden_state

   # The following Colab Notebook was used to generate these outputs:
   # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
   if "s16" in checkpoint_url:
      UpperCamelCase       =  torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]])
   elif "b16" in checkpoint_url:
      UpperCamelCase       =  torch.tensor([[1_4.2_8_8_9, -1_8.9_0_4_5, 1_1.7_2_8_1]])
   elif "l16" in checkpoint_url:
      UpperCamelCase       =  torch.tensor([[4_1.5_0_2_8, -2_2.8_6_8_1, 4_5.6_4_7_5]])
   elif "b4" in checkpoint_url:
      UpperCamelCase       =  torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]])
   else:
      UpperCamelCase       =  torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]])

   # verify logits
   assert torch.allclose(last_hidden_state[:, 0, :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__":
   snake_case_ :  List[str]     =					argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--checkpoint_url',
       default='https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar',
       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.'
   )

   snake_case_ :  Tuple     =					parser.parse_args()
   convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







'''simple docstring'''


import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

snake_case_ :  Tuple     =					[
    'word_embeddings_layernorm.weight',
    'word_embeddings_layernorm.bias',
    'input_layernorm.weight',
    'input_layernorm.bias',
    'post_attention_layernorm.weight',
    'post_attention_layernorm.bias',
    'self_attention.dense.bias',
    'mlp.dense_4h_to_h.bias',
    'ln_f.weight',
    'ln_f.bias',
]

snake_case_ :  Union[str, Any]     =					[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]



def  __snake_case     (    _UpperCAmelCase					:       List[str],      _UpperCAmelCase					:       List[str]):
   UpperCamelCase       =  {
       '''word_embeddings.weight''': '''word_embeddings.weight''',
       '''word_embeddings.norm.weight''': '''word_embeddings_layernorm.weight''',
       '''word_embeddings.norm.bias''': '''word_embeddings_layernorm.bias''',
       '''weight''': '''ln_f.weight''',
       '''bias''': '''ln_f.bias''',
   }

   if key in layer_rename_map:
      return layer_rename_map[key]

   # Handle transformer blocks
   UpperCamelCase       =  int(re.match(R'''.*layer_(\d*).*''',      _UpperCAmelCase)[1])
   layer_number -= 3
   return f'h.{layer_number}.' + key



def  __snake_case     (    _UpperCAmelCase					:       str):
   if dtype == torch.bool:
      return 1 / 8
   UpperCamelCase       =  re.search(R'''[^\d](\d+)$''',      str(_UpperCAmelCase))
   if bit_search is None:
      raise ValueError(f'`dtype` is not a valid dtype: {dtype}.')
   UpperCamelCase       =  int(bit_search.groups()[0])
   return bit_size // 8



def  __snake_case     (    _UpperCAmelCase					:       int,      _UpperCAmelCase					:       Optional[Any],      _UpperCAmelCase					:       str,      _UpperCAmelCase					:       List[Any],      _UpperCAmelCase					:       Optional[int]):
   # Construct model
   if bloom_config_file == "":
      UpperCamelCase       =  BloomConfig()
   else:
      UpperCamelCase       =  BloomConfig.from_json_file(_UpperCAmelCase)

   if shard_model:
      UpperCamelCase       =  os.listdir(_UpperCAmelCase)
      UpperCamelCase       =  sorted(filter(lambda _UpperCAmelCase: s.startswith('''layer''') and "model_00" in s,      _UpperCAmelCase))

      UpperCamelCase       =  {'''weight_map''': {}, '''metadata''': {}}
      UpperCamelCase       =  0

      UpperCamelCase       =  None

      UpperCamelCase       =  BloomConfig()

      for j, file in enumerate(_UpperCAmelCase):
         print('''Processing file: {}'''.format(_UpperCAmelCase))
         UpperCamelCase       =  None

         for i in range(_UpperCAmelCase):
            # load all TP files
            UpperCamelCase       =  file.replace('''model_00''',      f'model_0{i}')
            UpperCamelCase       =  torch.load(os.path.join(_UpperCAmelCase,      _UpperCAmelCase),      map_location='''cpu''')

            # Rename keys in the transformers names
            UpperCamelCase       =  list(temp.keys())
            for key in keys:
               UpperCamelCase       =  temp.pop(_UpperCAmelCase)

            if tensors is None:
               UpperCamelCase       =  temp
            else:
               for key in tensors.keys():
                  if any(key.endswith(_UpperCAmelCase) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
                     # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
                     tensors[key] += temp[key]
                  else:
                     # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                     UpperCamelCase       =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
                     # We concatenate these weights accross TP ranks
                     UpperCamelCase       =  torch.cat([tensors[key], temp[key]],      dim=_UpperCAmelCase)

            # Divide by the number of TP the weights we want to average
         for key in tensors.keys():
            if any(key.endswith(_UpperCAmelCase) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
               UpperCamelCase       =  tensors[key] / pretraining_tp
         torch.save(
             _UpperCAmelCase,      os.path.join(
                 _UpperCAmelCase,      '''pytorch_model_{}-of-{}.bin'''.format(str(j + 1).zfill(5),      str(len(_UpperCAmelCase)).zfill(5)),      ),      )

         for key in tensors.keys():
            UpperCamelCase       =  tensors[key]
            total_size += value.numel() * get_dtype_size(value.dtype)
            if key not in index_dict["weight_map"]:
               UpperCamelCase       =  '''pytorch_model_{}-of-{}.bin'''.format(
                   str(j + 1).zfill(5),      str(len(_UpperCAmelCase)).zfill(5))

      UpperCamelCase       =  BloomConfig()
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + CONFIG_NAME
      UpperCamelCase       =  total_size
      with open(_UpperCAmelCase,      '''w''',      encoding='''utf-8''') as f:
         f.write(config.to_json_string())
      with open(os.path.join(_UpperCAmelCase,      WEIGHTS_NAME + '''.index.json'''),      '''w''',      encoding='''utf-8''') as f:
         UpperCamelCase       =  json.dumps(_UpperCAmelCase,      indent=2,      sort_keys=_UpperCAmelCase) + '''\n'''
         f.write(_UpperCAmelCase)
   else:
      UpperCamelCase       =  BloomModel(_UpperCAmelCase)

      UpperCamelCase       =  os.listdir(_UpperCAmelCase)
      UpperCamelCase       =  sorted(filter(lambda _UpperCAmelCase: s.startswith('''layer''') and "model_00" in s,      _UpperCAmelCase))

      UpperCamelCase       =  None
      for i, file in enumerate(_UpperCAmelCase):
         UpperCamelCase       =  None
         for i in range(_UpperCAmelCase):
            # load all TP files
            UpperCamelCase       =  file.replace('''model_00''',      f'model_0{i}')
            UpperCamelCase       =  torch.load(os.path.join(_UpperCAmelCase,      _UpperCAmelCase),      map_location='''cpu''')

            # Rename keys in the transformers names
            UpperCamelCase       =  list(temp.keys())
            for key in keys:
               UpperCamelCase       =  temp.pop(_UpperCAmelCase)

            if tensors is None:
               UpperCamelCase       =  temp
            else:
               for key in tensors.keys():
                  # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
                  if any(key.endswith(_UpperCAmelCase) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
                     tensors[key] += temp[key]
                  else:
                     # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                     UpperCamelCase       =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
                     # We concatenate these weights accross TP ranks
                     UpperCamelCase       =  torch.cat([tensors[key], temp[key]],      dim=_UpperCAmelCase)

            # Divide by the number of TP the weights we want to average
         for key in tensors.keys():
            if any(key.endswith(_UpperCAmelCase) for end in WEIGHTS_TO_AVERAGE_ENDSWITH):
               UpperCamelCase       =  tensors[key] / pretraining_tp

         UpperCamelCase       =  model.load_state_dict(_UpperCAmelCase,      strict=_UpperCAmelCase)
         assert not other_keys.unexpected_keys, f'The keys {other_keys.unexpected_keys} are unexpected'
         if missing_keys is None:
            UpperCamelCase       =  set(other_keys.missing_keys)
         else:
            UpperCamelCase       =  missing_keys.intersection(set(other_keys.missing_keys))

      assert not missing_keys, f'The keys {missing_keys} are missing'

      # Save pytorch-model
      os.makedirs(_UpperCAmelCase,      exist_ok=_UpperCAmelCase)
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + CONFIG_NAME
      print(f'Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}')
      if config.torch_dtype is not None:
         UpperCamelCase       =  model.to(config.torch_dtype)
      torch.save(model.state_dict(),      _UpperCAmelCase)
      print(f'Save configuration file to {pytorch_config_dump_path}')
      with open(_UpperCAmelCase,      '''w''',      encoding='''utf-8''') as f:
         f.write(config.to_json_string())


if __name__ == "__main__":
   snake_case_ :  Tuple     =					argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--bloom_checkpoint_path',
       default=None,
       type=str,
       required=True,
       help='Path to the Megatron-LM checkpoint path.',
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   parser.add_argument(
       '--bloom_config_file',
       default='',
       type=str,
       help=(
           'An optional config json file corresponding to the pre-trained model. \n'
           'This specifies the model architecture.'
       ),
   )
   parser.add_argument(
       '--shard_model',
       action='store_true',
       help='An optional setting to shard the output model \nThis enables sharding the converted checkpoint',
   )
   parser.add_argument(
       '--pretraining_tp',
       default=4,
       type=int,
       help='Pretraining TP rank that has been used when training the model in Megatron-LM \n',
   )
   snake_case_ :  List[str]     =					parser.parse_args()
   convert_bloom_checkpoint_to_pytorch(
       args.bloom_checkpoint_path,
       args.bloom_config_file,
       args.pytorch_dump_folder_path,
       args.shard_model,
       args.pretraining_tp,
   )





