




'''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()

lowercase      :   Optional[int]      =       [
    """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""",
]

lowercase      :   int      =       [
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]



def        lowerCamelCase__				(      __lowercase     ,   __lowercase				):
     snake_case			:      Dict        =			{
         """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
     snake_case			:      int        =			int(re.match(r""".*layer_(\d*).*"""     ,   __lowercase				)[1]				)
     layer_number -= 3
     return F'''h.{layer_number}.''' + key



def        lowerCamelCase__				(      __lowercase				):
     if dtype == torch.bool:
          return 1 / 8
     snake_case			:      Dict        =			re.search(r"""[^\d](\d+)$"""     ,   str(__lowercase				)				)
     if bit_search is None:
          raise ValueError(F'''`dtype` is not a valid dtype: {dtype}.'''				)
     snake_case			:      str        =			int(bit_search.groups()[0]				)
     return bit_size // 8



def        lowerCamelCase__				(      __lowercase     ,   __lowercase     ,   __lowercase     ,   __lowercase     ,   __lowercase				):
     # Construct model
     if bloom_config_file == "":
          snake_case			:      Union[str, Any]        =			BloomConfig()
     else:
          snake_case			:      int        =			BloomConfig.from_json_file(__lowercase				)

     if shard_model:
          snake_case			:      Union[str, Any]        =			os.listdir(__lowercase				)
          snake_case			:      Dict        =			sorted(filter(lambda __lowercase				: s.startswith("""layer"""				) and "model_00" in s     ,   __lowercase				)				)

          snake_case			:      Optional[int]        =			{"""weight_map""": {}, """metadata""": {}}
          snake_case			:      Any        =			0

          snake_case			:      int        =			None

          snake_case			:      Any        =			BloomConfig()

          for j, file in enumerate(__lowercase				):
               print("""Processing file: {}""".format(__lowercase				)				)
               snake_case			:      List[Any]        =			None

               for i in range(__lowercase				):
                    # load all TP files
                    snake_case			:      List[Any]        =			file.replace("""model_00"""     ,   F'''model_0{i}'''				)
                    snake_case			:      Optional[int]        =			torch.load(os.path.join(__lowercase     ,   __lowercase				)     ,   map_location="""cpu"""				)

                    # Rename keys in the transformers names
                    snake_case			:      Dict        =			list(temp.keys()				)
                    for key in keys:
                         snake_case			:      Any        =			temp.pop(__lowercase				)

                    if tensors is None:
                         snake_case			:      str        =			temp
                    else:
                         for key in tensors.keys():
                              if any(key.endswith(__lowercase				) 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
                                   snake_case			:      Tuple        =			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN				) else 0
                                   # We concatenate these weights accross TP ranks
                                   snake_case			:      Optional[Any]        =			torch.cat([tensors[key], temp[key]]     ,   dim=__lowercase				)

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

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

          snake_case			:      List[Any]        =			BloomConfig()
          snake_case			:      Optional[int]        =			pytorch_dump_folder_path + """/""" + CONFIG_NAME
          snake_case			:      List[str]        =			total_size
          with open(__lowercase     ,   """w"""     ,   encoding="""utf-8"""				) as f:
               f.write(config.to_json_string()				)
          with open(os.path.join(__lowercase     ,   WEIGHTS_NAME + """.index.json"""				)     ,   """w"""     ,   encoding="""utf-8"""				) as f:
               snake_case			:      Dict        =			json.dumps(__lowercase     ,   indent=2     ,   sort_keys=__lowercase				) + """\n"""
               f.write(__lowercase				)
     else:
          snake_case			:      List[Any]        =			BloomModel(__lowercase				)

          snake_case			:      Tuple        =			os.listdir(__lowercase				)
          snake_case			:      int        =			sorted(filter(lambda __lowercase				: s.startswith("""layer"""				) and "model_00" in s     ,   __lowercase				)				)

          snake_case			:      Tuple        =			None
          for i, file in enumerate(__lowercase				):
               snake_case			:      List[Any]        =			None
               for i in range(__lowercase				):
                    # load all TP files
                    snake_case			:      Union[str, Any]        =			file.replace("""model_00"""     ,   F'''model_0{i}'''				)
                    snake_case			:      List[str]        =			torch.load(os.path.join(__lowercase     ,   __lowercase				)     ,   map_location="""cpu"""				)

                    # Rename keys in the transformers names
                    snake_case			:      Dict        =			list(temp.keys()				)
                    for key in keys:
                         snake_case			:      Optional[int]        =			temp.pop(__lowercase				)

                    if tensors is None:
                         snake_case			:      Any        =			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(__lowercase				) for end in WEIGHTS_TO_AVERAGE_ENDSWITH				):
                                   tensors[key] += temp[key]
                              else:
                                   # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                   snake_case			:      Tuple        =			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN				) else 0
                                   # We concatenate these weights accross TP ranks
                                   snake_case			:      int        =			torch.cat([tensors[key], temp[key]]     ,   dim=__lowercase				)

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

               snake_case			:      int        =			model.load_state_dict(__lowercase     ,   strict=__lowercase				)
               assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
               if missing_keys is None:
                    snake_case			:      Union[str, Any]        =			set(other_keys.missing_keys				)
               else:
                    snake_case			:      Dict        =			missing_keys.intersection(set(other_keys.missing_keys				)				)

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

          # Save pytorch-model
          os.makedirs(__lowercase     ,   exist_ok=__lowercase				)
          snake_case			:      Dict        =			pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
          snake_case			:      Union[str, Any]        =			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:
               snake_case			:      int        =			model.to(config.torch_dtype				)
          torch.save(model.state_dict()     ,   __lowercase				)
          print(F'''Save configuration file to {pytorch_config_dump_path}'''				)
          with open(__lowercase     ,   """w"""     ,   encoding="""utf-8"""				) as f:
               f.write(config.to_json_string()				)


if __name__ == "__main__":
      lowercase      :   Dict      =       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""",
      )
      lowercase      :   Tuple      =       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,
      )









'''simple docstring'''







import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_lxmert import LxmertTokenizer


lowercase      :   Tuple      =       {"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

lowercase      :   Dict      =       {
    """vocab_file""": {
        """unc-nlp/lxmert-base-uncased""": """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/vocab.txt""",
    },
    """tokenizer_file""": {
        """unc-nlp/lxmert-base-uncased""": (
            """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/tokenizer.json"""
        ),
    },
}

lowercase      :   List[Any]      =       {
    """unc-nlp/lxmert-base-uncased""": 512,
}

lowercase      :   Tuple      =       {
    """unc-nlp/lxmert-base-uncased""": {"""do_lower_case""": True},
}







class 					_a (a__	):





     '''simple docstring'''



     lowerCAmelCase_  :     Optional[int]          =					VOCAB_FILES_NAMES
     lowerCAmelCase_  :     Optional[int]          =					PRETRAINED_VOCAB_FILES_MAP
     lowerCAmelCase_  :     Any          =					PRETRAINED_INIT_CONFIGURATION
     lowerCAmelCase_  :     Optional[int]          =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     lowerCAmelCase_  :     Optional[int]          =					LxmertTokenizer







     def __init__(				self							,__a=None							,__a=None							,__a=True							,__a="[UNK]"							,__a="[SEP]"							,__a="[PAD]"							,__a="[CLS]"							,__a="[MASK]"							,__a=True							,__a=None							,**__a							,)      ->				str:
          super().__init__(
              __a							,tokenizer_file=__a							,do_lower_case=__a							,unk_token=__a							,sep_token=__a							,pad_token=__a							,cls_token=__a							,mask_token=__a							,tokenize_chinese_chars=__a							,strip_accents=__a							,**__a							,)

          snake_case			:      List[Any]        =			json.loads(self.backend_tokenizer.normalizer.__getstate__()							)
          if (
              normalizer_state.get("""lowercase"""							,__a							) != do_lower_case
              or normalizer_state.get("""strip_accents"""							,__a							) != strip_accents
              or normalizer_state.get("""handle_chinese_chars"""							,__a							) != tokenize_chinese_chars
          ):
               snake_case			:      Union[str, Any]        =			getattr(__a							,normalizer_state.pop("""type"""							)							)
               snake_case			:      Optional[int]        =			do_lower_case
               snake_case			:      Optional[int]        =			strip_accents
               snake_case			:      List[Any]        =			tokenize_chinese_chars
               snake_case			:      Union[str, Any]        =			normalizer_class(**__a							)

          snake_case			:      str        =			do_lower_case







     def 			snake_case_      (				self							,__a							,__a=None							)      ->				Optional[int]:
          snake_case			:      Tuple        =			[self.cls_token_id] + token_ids_a + [self.sep_token_id]

          if token_ids_a:
               output += token_ids_a + [self.sep_token_id]

          return output







     def 			snake_case_      (				self							,__a							,__a = None							)      ->				List[int]:
          snake_case			:      int        =			[self.sep_token_id]
          snake_case			:      Any        =			[self.cls_token_id]
          if token_ids_a is None:
               return len(cls + token_ids_a + sep							) * [0]
          return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]







     def 			snake_case_      (				self							,__a							,__a = None							)      ->				Tuple[str]:
          snake_case			:      Optional[int]        =			self._tokenizer.model.save(__a							,name=__a							)
          return tuple(__a							)



