






'''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 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase							):


   """simple docstring"""




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




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase							):


   """simple docstring"""




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




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase							):


   """simple docstring"""




   if bloom_config_file == "":
      lowercase_	:    Tuple           =     BloomConfig()
   else:
      lowercase_	:    str           =     BloomConfig.from_json_file(_UpperCamelCase							)

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

      lowercase_	:    str           =     {"weight_map": {}, "metadata": {}}
      lowercase_	:    Any           =     0

      lowercase_	:    str           =     None

      lowercase_	:    Any           =     BloomConfig()

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

         for i in range(_UpperCamelCase							):
            # load all TP files
            lowercase_	:    Union[str, Any]           =     file.replace("model_00"						,		F"""model_0{i}"""							)
            lowercase_	:    Dict           =     torch.load(os.path.join(_UpperCamelCase						,		_UpperCamelCase							)						,		map_location="cpu"							)

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

            if tensors is None:
               lowercase_	:    Union[str, Any]           =     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
                     lowercase_	:    Union[str, Any]           =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN							) else 0
                     # We concatenate these weights accross TP ranks
                     lowercase_	:    Union[str, Any]           =     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							):
               lowercase_	:    str           =     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():
            lowercase_	:    str           =     tensors[key]
            total_size += value.numel() * get_dtype_size(value.dtype							)
            if key not in index_dict["weight_map"]:
               lowercase_	:    Any           =     "pytorch_model_{}-of-{}.bin".format(
                   str(j + 1							).zfill(5							)						,		str(len(_UpperCamelCase							)							).zfill(5							)							)

      lowercase_	:    int           =     BloomConfig()
      lowercase_	:    List[Any]           =     pytorch_dump_folder_path + "/" + CONFIG_NAME
      lowercase_	:    Dict           =     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:
         lowercase_	:    Any           =     json.dumps(_UpperCamelCase						,		indent=2						,		sort_keys=_UpperCamelCase							) + "\n"
         f.write(_UpperCamelCase							)
   else:
      lowercase_	:    Tuple           =     BloomModel(_UpperCamelCase							)

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

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

            # Rename keys in the transformers names
            lowercase_	:    int           =     list(temp.keys()							)
            for key in keys:
               lowercase_	:    Optional[Any]           =     temp.pop(_UpperCamelCase							)

            if tensors is None:
               lowercase_	:    List[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(_UpperCamelCase							) for end in WEIGHTS_TO_AVERAGE_ENDSWITH							):
                     tensors[key] += temp[key]
                  else:
                     # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                     lowercase_	:    Dict           =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN							) else 0
                     # We concatenate these weights accross TP ranks
                     lowercase_	:    Any           =     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							):
               lowercase_	:    Any           =     tensors[key] / pretraining_tp

         lowercase_	:    int           =     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:
            lowercase_	:    List[str]           =     set(other_keys.missing_keys							)
         else:
            lowercase_	:    int           =     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							)
      lowercase_	:    int           =     pytorch_dump_folder_path + "/" + WEIGHTS_NAME
      lowercase_	:    str           =     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_	:    List[str]           =     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__ =			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'''
from __future__ import annotations

UpperCamelCase__ =			[-10, -5, 0, 5, 5.1, 11, 13, 21, 3, 4, -21, -10, -5, -1, 0]
UpperCamelCase__ =			[-5, 0, 5, 5.1, 11, 13, 21, -1, 4, -1, -10, -5, -1, 0, -1]




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase							):


   """simple docstring"""




   lowercase_	:    str           =     []
   lowercase_	:    List[str]           =     len(_UpperCamelCase							)

   for i in range(_UpperCamelCase							):
      lowercase_	:    float           =     -1
      for j in range(i + 1						,		_UpperCamelCase							):
         if arr[i] < arr[j]:
            lowercase_	:    Union[str, Any]           =     arr[j]
            break
      result.append(_UpperCamelCase							)
   return result




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase							):


   """simple docstring"""




   lowercase_	:    List[str]           =     []
   for i, outer in enumerate(_UpperCamelCase							):
      lowercase_	:    float           =     -1
      for inner in arr[i + 1 :]:
         if outer < inner:
            lowercase_	:    Optional[Any]           =     inner
            break
      result.append(_UpperCamelCase							)
   return result




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase							):


   """simple docstring"""




   lowercase_	:    List[str]           =     len(_UpperCamelCase							)
   lowercase_	:    list[float]           =     []
   lowercase_	:    list[float]           =     [-1] * arr_size

   for index in reversed(range(_UpperCamelCase							)							):
      if stack:
         while stack[-1] <= arr[index]:
            stack.pop()
            if not stack:
               break
      if stack:
         lowercase_	:    Optional[Any]           =     stack[-1]
      stack.append(arr[index]							)
   return result


if __name__ == "__main__":
 from doctest import testmod
 from timeit import timeit

 testmod()
 print(next_greatest_element_slow(arr))
 print(next_greatest_element_fast(arr))
 print(next_greatest_element(arr))

 UpperCamelCase__ =			(
     'from __main__ import arr, next_greatest_element_slow, '
     'next_greatest_element_fast, next_greatest_element'
 )
 print(
     'next_greatest_element_slow():',
     timeit('next_greatest_element_slow(arr)', setup=setup),
 )
 print(
     'next_greatest_element_fast():',
     timeit('next_greatest_element_fast(arr)', setup=setup),
 )
 print(
     '     next_greatest_element():',
     timeit('next_greatest_element(arr)', setup=setup),
 )




