


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

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

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

def 				_snake_case      (   __snake_case							:      int						, __snake_case							:      str   ):




   """simple docstring"""







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

def 				_snake_case      (   __snake_case							:      Union[str, Any]   ):




   """simple docstring"""







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

def 				_snake_case      (   __snake_case							:      Optional[Any]						, __snake_case							:      Union[str, Any]						, __snake_case							:      Tuple						, __snake_case							:      Any						, __snake_case							:      Optional[Any]   ):




   """simple docstring"""







   if bloom_config_file == "":
      _lowerCamelCase	:					List[str]   =		BloomConfig()
   else:
      _lowerCamelCase	:					Optional[Any]   =		BloomConfig.from_json_file(__snake_case   )

   if shard_model:
      _lowerCamelCase	:					Any   =		os.listdir(__snake_case   )
      _lowerCamelCase	:					Any   =		sorted(filter(lambda __snake_case   : s.startswith("""layer"""   ) and "model_00" in s						, __snake_case   )   )

      _lowerCamelCase	:					Any   =		{"""weight_map""": {}, """metadata""": {}}
      _lowerCamelCase	:					Union[str, Any]   =		0

      _lowerCamelCase	:					Tuple   =		None

      _lowerCamelCase	:					Optional[Any]   =		BloomConfig()

      for j, file in enumerate(__snake_case   ):
         print("""Processing file: {}""".format(__snake_case   )   )
         _lowerCamelCase	:					str   =		None

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

            # Rename keys in the transformers names
            _lowerCamelCase	:					Optional[int]   =		list(temp.keys()   )
            for key in keys:
               _lowerCamelCase	:					str   =		temp.pop(__snake_case   )

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

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

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

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

      _lowerCamelCase	:					Tuple   =		os.listdir(__snake_case   )
      _lowerCamelCase	:					Any   =		sorted(filter(lambda __snake_case   : s.startswith("""layer"""   ) and "model_00" in s						, __snake_case   )   )

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

            # Rename keys in the transformers names
            _lowerCamelCase	:					Union[str, Any]   =		list(temp.keys()   )
            for key in keys:
               _lowerCamelCase	:					List[str]   =		temp.pop(__snake_case   )

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

            # Divide by the number of TP the weights we want to average
         for key in tensors.keys():
            if any(key.endswith(__snake_case   ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH   ):
               _lowerCamelCase	:					Union[str, Any]   =		tensors[key] / pretraining_tp

         _lowerCamelCase	:					List[Any]   =		model.load_state_dict(__snake_case						, strict=__snake_case   )
         assert not other_keys.unexpected_keys, F'The keys {other_keys.unexpected_keys} are unexpected'
         if missing_keys is None:
            _lowerCamelCase	:					Tuple   =		set(other_keys.missing_keys   )
         else:
            _lowerCamelCase	:					Tuple   =		missing_keys.intersection(set(other_keys.missing_keys   )   )

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

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


if __name__ == "__main__":
 UpperCAmelCase				=						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""",
 )
 UpperCAmelCase				=						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'''

def _SCREAMING_SNAKE_CASE (A    ,					A    ,					A			)			->			list:






       """simple docstring"""







       lowercase__      =					len(A			)
       lowercase__      =					[[0] * n for i in range(A			)]
       for i in range(A			):
              lowercase__      =					y_points[i]

       for i in range(2    ,					A			):
              for j in range(A    ,					A			):
                     lowercase__      =					(
                         (xa - x_points[j - i + 1]) * q[j][i - 1]
                         - (xa - x_points[j]) * q[j - 1][i - 1]
                     ) / (x_points[j] - x_points[j - i + 1])

       return [q[n - 1][n - 1], q]


if __name__ == "__main__":
 import doctest

 doctest.testmod()




