




'''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_						:      Optional[Any]       =						[
    '''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_						:      List[Any]       =						[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]





def      __a     (       _UpperCamelCase:   Any ,						_UpperCamelCase:   int   )						->			str:

     """simple docstring"""




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





def      __a     (       _UpperCamelCase:   Union[str, Any]   )						->			Optional[Any]:

     """simple docstring"""




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





def      __a     (       _UpperCamelCase:   Optional[Any] ,						_UpperCamelCase:   List[str] ,						_UpperCamelCase:   Dict ,						_UpperCamelCase:   str ,						_UpperCamelCase:   Tuple   )						->			int:

     """simple docstring"""




     if bloom_config_file == "":
          _snake_case							=				BloomConfig()
     else:
          _snake_case							=				BloomConfig.from_json_file(_UpperCamelCase   )

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

          _snake_case							=				{"weight_map": {}, "metadata": {}}
          _snake_case							=				0

          _snake_case							=				None

          _snake_case							=				BloomConfig()

          for j, file in enumerate(_UpperCamelCase   ):
               print("Processing file: {}".format(_UpperCamelCase   )   )
               _snake_case							=				None

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

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

                    if tensors is None:
                         _snake_case							=				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
                                   _snake_case							=				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							=				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   ):
                         _snake_case							=				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():
                    _snake_case							=				tensors[key]
                    total_size += value.numel() * get_dtype_size(value.dtype   )
                    if key not in index_dict["weight_map"]:
                         _snake_case							=				"pytorch_model_{}-of-{}.bin".format(
                             str(j + 1   ).zfill(5   ) ,						str(len(_UpperCamelCase   )   ).zfill(5   )   )

          _snake_case							=				BloomConfig()
          _snake_case							=				pytorch_dump_folder_path + "/" + CONFIG_NAME
          _snake_case							=				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:
               _snake_case							=				json.dumps(_UpperCamelCase ,						indent=2 ,						sort_keys=_UpperCamelCase   ) + "\n"
               f.write(_UpperCamelCase   )
     else:
          _snake_case							=				BloomModel(_UpperCamelCase   )

          _snake_case							=				os.listdir(_UpperCamelCase   )
          _snake_case							=				sorted(filter(lambda _UpperCamelCase   : s.startswith("layer"   ) and "model_00" in s ,						_UpperCamelCase   )   )

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

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

                    if tensors is None:
                         _snake_case							=				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
                                   _snake_case							=				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							=				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   ):
                         _snake_case							=				tensors[key] / pretraining_tp

               _snake_case							=				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:
                    _snake_case							=				set(other_keys.missing_keys   )
               else:
                    _snake_case							=				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   )
          _snake_case							=				pytorch_dump_folder_path + "/" + WEIGHTS_NAME
          _snake_case							=				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							=				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__":
   UpperCamelCase_						:      Any       =						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_						:      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,
   )








'''simple docstring'''





from __future__ import annotations

from collections.abc import Callable





def      __a     (       _UpperCamelCase:   Callable[[int | float], int | float] ,						_UpperCamelCase:   int | float ,						_UpperCamelCase:   int | float ,						_UpperCamelCase:   int = 100 ,						)						->			float:

     """simple docstring"""




     _snake_case							=				x_start
     _snake_case							=				fnc(_UpperCamelCase   )
     _snake_case							=				0.0
     for _ in range(_UpperCamelCase   ):
          # Approximates small segments of curve as linear and solve
          # for trapezoidal area
          _snake_case							=				(x_end - x_start) / steps + xa
          _snake_case							=				fnc(_UpperCamelCase   )
          area += abs(fxa + fxa   ) * (xa - xa) / 2
          # Increment step
          _snake_case							=				xa
          _snake_case							=				fxa
     return area


if __name__ == "__main__":





   def      __a     (       _UpperCamelCase:   Any   )						->			Optional[int]:

        """simple docstring"""




        return x**3 + x**2

   print('''f(x) = x^3 + x^2''')
   print('''The area between the curve, x = -5, x = 5 and the x axis is:''')
   UpperCamelCase_						:      Optional[int]       =						10
   while i <= 100000:
      print(F'with {i} steps: {trapezoidal_area(f, -5, 5, i)}')
      i *= 10


