

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

lowerCAmelCase__      		=			[
    '''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''',
]

lowerCAmelCase__      		=			[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]


def 	_A  (					A__  ,       A__   ):






 """simple docstring"""



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


def 	_A  (					A__   ):






 """simple docstring"""



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


def 	_A  (					A__  ,       A__  ,       A__  ,       A__  ,       A__   ):






 """simple docstring"""



 if bloom_config_file == "":
  __lowercase													=				BloomConfig()
 else:
  __lowercase													=				BloomConfig.from_json_file(A__   )

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

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

  __lowercase													=				None

  __lowercase													=				BloomConfig()

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

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

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

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

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

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

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

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

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

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

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

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

   __lowercase													=				model.load_state_dict(A__  ,       strict=A__   )
   assert not other_keys.unexpected_keys, F"The keys {other_keys.unexpected_keys} are unexpected"
   if missing_keys is None:
    __lowercase													=				set(other_keys.missing_keys   )
   else:
    __lowercase													=				missing_keys.intersection(set(other_keys.missing_keys   )   )

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

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


if __name__ == "__main__":
 lowerCAmelCase__      		=			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''',
 )
 lowerCAmelCase__      		=			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 copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs






class     lowercase_			(lowerCamelCase__   ):




  """simple docstring"""




  def __init__(   self     :  Optional[Any]      ,lowercase__     :  int      ,lowercase__     :  List[str]=None      ,lowercase__     :  Any=True      ,lowercase__     :  Union[str, Any]=None      ,**lowercase__     :  Dict    ):
   __lowercase													=				parent
   __lowercase													=				config_class
   __lowercase													=				has_text_modality
   __lowercase													=				kwargs
   __lowercase													=				common_properties






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				(
       ['''hidden_size''', '''num_attention_heads''', '''num_hidden_layers''']
       if self.common_properties is None
       else self.common_properties
   )

   # Add common fields for text models
   if self.has_text_modality:
    common_properties.extend(['''vocab_size''']    )

   # Test that config has the common properties as getters
   for prop in common_properties:
    self.parent.assertTrue(hasattr(lowercase__      ,lowercase__    )      ,msg=F"`{prop}` does not exist"    )

   # Test that config has the common properties as setter
   for idx, name in enumerate(lowercase__    ):
    try:
     setattr(lowercase__      ,lowercase__      ,lowercase__    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass

        # Test if config class can be called with Config(prop_name=..)
   for idx, name in enumerate(lowercase__    ):
    try:
     __lowercase													=				self.config_class(**{name: idx}    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass






  def       SCREAMING_SNAKE_CASE   (   self     :  str    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				json.loads(config.to_json_string()    )
   for key, value in self.inputs_dict.items():
    self.parent.assertEqual(obj[key]      ,lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  int    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,'''config.json'''    )
    config_first.to_json_file(lowercase__    )
    __lowercase													=				self.config_class.from_json_file(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   __lowercase													=				'''test'''
   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,lowercase__    )
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__      ,subfolder=lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict      ,num_labels=5    )
   self.parent.assertEqual(len(config.idalabel    )      ,5    )
   self.parent.assertEqual(len(config.labelaid    )      ,5    )

   __lowercase													=				3
   self.parent.assertEqual(len(config.idalabel    )      ,3    )
   self.parent.assertEqual(len(config.labelaid    )      ,3    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Tuple    ):
   if self.config_class.is_composition:
    return
   __lowercase													=				self.config_class()
   self.parent.assertIsNotNone(lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				copy.deepcopy(lowercase__    )
   __lowercase													=				self.config_class(**lowercase__    )
   __lowercase													=				[]
   for key, value in config_common_kwargs.items():
    if key == "torch_dtype":
     if not is_torch_available():
      continue
     else:
      import torch

      if config.torch_dtype != torch.floataa:
       wrong_values.append(('''torch_dtype''', config.torch_dtype, torch.floataa)    )
    elif getattr(lowercase__      ,lowercase__    ) != value:
     wrong_values.append((key, getattr(lowercase__      ,lowercase__    ), value)    )

   if len(lowercase__    ) > 0:
    __lowercase													=				'''\n'''.join([F"- {v[0]}: got {v[1]} instead of {v[2]}" for v in wrong_values]    )
    raise ValueError(F"The following keys were not properly set in the config:\n{errors}"    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[int]    ):
   self.create_and_test_config_common_properties()
   self.create_and_test_config_to_json_string()
   self.create_and_test_config_to_json_file()
   self.create_and_test_config_from_and_save_pretrained()
   self.create_and_test_config_from_and_save_pretrained_subfolder()
   self.create_and_test_config_with_num_labels()
   self.check_config_can_be_init_without_params()
   self.check_config_arguments_init()

