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

__a										=					[
    """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""",
]

__a										=					[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]






def    __UpperCAmelCase			(  a_:       Optional[int],					a_:       str							):
   _UpperCAmelCase							:	Optional[int]							=   {
       '''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
   _UpperCAmelCase							:	Tuple							=   int(re.match(r".*layer_(\d*).*",					__a							)[1]							)
   layer_number -= 3
   return f"""h.{layer_number}.""" + key






def    __UpperCAmelCase			(  a_:       Any							):






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






def    __UpperCAmelCase			(  a_:       int,					a_:       Dict,					a_:       Any,					a_:       List[Any],					a_:       List[Any]							):






   if bloom_config_file == "":
      _UpperCAmelCase							:	List[Any]							=   BloomConfig()
   else:
      _UpperCAmelCase							:	str							=   BloomConfig.from_json_file(__a							)

   if shard_model:
      _UpperCAmelCase							:	List[Any]							=   os.listdir(__a							)
      _UpperCAmelCase							:	List[str]							=   sorted(filter(lambda a_							: s.startswith("layer"							) and "model_00" in s,					__a							)							)

      _UpperCAmelCase							:	List[str]							=   {'''weight_map''': {}, '''metadata''': {}}
      _UpperCAmelCase							:	Any							=   0

      _UpperCAmelCase							:	Dict							=   None

      _UpperCAmelCase							:	Optional[int]							=   BloomConfig()

      for j, file in enumerate(__a							):
         print("Processing file: {}".format(__a							)							)
         _UpperCAmelCase							:	Optional[Any]							=   None

         for i in range(__a							):
            # load all TP files
            _UpperCAmelCase							:	List[str]							=   file.replace("model_00",					f"""model_0{i}"""							)
            _UpperCAmelCase							:	Tuple							=   torch.load(os.path.join(__a,					__a							),					map_location="cpu"							)

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

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

      _UpperCAmelCase							:	List[str]							=   BloomConfig()
      _UpperCAmelCase							:	List[str]							=   pytorch_dump_folder_path + '''/''' + CONFIG_NAME
      _UpperCAmelCase							:	Dict							=   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:
         _UpperCAmelCase							:	Dict							=   json.dumps(__a,					indent=2,					sort_keys=__a							) + '''\n'''
         f.write(__a							)
   else:
      _UpperCAmelCase							:	str							=   BloomModel(__a							)

      _UpperCAmelCase							:	Optional[int]							=   os.listdir(__a							)
      _UpperCAmelCase							:	int							=   sorted(filter(lambda a_							: s.startswith("layer"							) and "model_00" in s,					__a							)							)

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

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

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

         _UpperCAmelCase							:	Tuple							=   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:
            _UpperCAmelCase							:	List[Any]							=   set(other_keys.missing_keys							)
         else:
            _UpperCAmelCase							:	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(__a,					exist_ok=__a							)
      _UpperCAmelCase							:	Optional[int]							=   pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
      _UpperCAmelCase							:	Optional[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:
         _UpperCAmelCase							:	List[str]							=   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__":
     __a										=					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',
     )
     __a										=					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 typing import Dict, List, Optional, Union

import numpy as np

from transformers.utils import is_vision_available
from transformers.utils.generic import TensorType

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import logging


if is_vision_available():
  import PIL

UpperCAmelCase_   : int   			= logging.get_logger(__name__)



def 					_A			(__a							) ->       List[List[ImageInput]]:
    """simple docstring"""






    if isinstance(__a    ,				(list, tuple)							) and isinstance(videos[0]    ,				(list, tuple)							) and is_valid_image(videos[0][0]							):
        return videos

    elif isinstance(__a    ,				(list, tuple)							) and is_valid_image(videos[0]							):
        return [videos]

    elif is_valid_image(__a							):
        return [[videos]]

    raise ValueError(f'Could not make batched video from {videos}'							)





class 							lowerCAmelCase__				(			UpperCAmelCase__ ):






      '''simple docstring'''

      __UpperCamelCase	 =  ["pixel_values"]
      def __init__(						self     :  Dict						,	lowercase_     :  bool = True						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = PILImageResampling.BILINEAR						,	lowercase_     :  bool = True						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = True						,	lowercase_     :  Union[int, float] = 1 / 255						,	lowercase_     :  bool = True						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	**lowercase_     :  Dict						,	):




          '''simple docstring'''


          super().__init__(**lowercase_)
          SCREAMING_SNAKE_CASE_      :  str    						=		size if size is not None else {'''shortest_edge''': 256}
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		crop_size if crop_size is not None else {'''height''': 224, '''width''': 224}
          SCREAMING_SNAKE_CASE_      :  Dict    						=		get_size_dict(lowercase_						,	param_name='''crop_size''')

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		do_resize
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		size
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		do_center_crop
          SCREAMING_SNAKE_CASE_      :  Dict    						=		crop_size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		resample
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		do_rescale
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		rescale_factor
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		offset
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		do_normalize
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image_std if image_std is not None else IMAGENET_STANDARD_STD
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  PILImageResampling = PILImageResampling.BILINEAR						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  Any						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          if "shortest_edge" in size:
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		get_resize_output_image_size(lowercase_						,	size['''shortest_edge''']						,	default_to_square=lowercase_)
          elif "height" in size and "width" in size:
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		(size['''height'''], size['''width'''])
          else:
              raise ValueError(F'Size must have \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}')
          return resize(lowercase_						,	size=lowercase_						,	resample=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  List[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		get_size_dict(lowercase_)
          if "height" not in size or "width" not in size:
              raise ValueError(F'Size must have \'height\' and \'width\' as keys. Got {size.keys()}')
          return center_crop(lowercase_						,	size=(size['''height'''], size['''width'''])						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Union[int, float]						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  Optional[int]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		image.astype(np.floataa)
          if offset:
              SCREAMING_SNAKE_CASE_      :  Tuple    						=		image - (scale / 2)
          return rescale(lowercase_						,	scale=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Union[float, List[float]]						,	lowercase_     :  Union[float, List[float]]						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  List[str]						,	):




          '''simple docstring'''


          return normalize(lowercase_						,	mean=lowercase_						,	std=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]						,	lowercase_     :  ImageInput						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = None						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = None						,	lowercase_     :  float = None						,	lowercase_     :  bool = None						,	lowercase_     :  bool = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[ChannelDimension] = ChannelDimension.FIRST						,	):




          '''simple docstring'''


          if do_resize and size is None or resample is None:
              raise ValueError('''Size and resample must be specified if do_resize is True.''')

          if do_center_crop and crop_size is None:
              raise ValueError('''Crop size must be specified if do_center_crop is True.''')

          if do_rescale and rescale_factor is None:
              raise ValueError('''Rescale factor must be specified if do_rescale is True.''')

          if do_normalize and (image_mean is None or image_std is None):
              raise ValueError('''Image mean and std must be specified if do_normalize is True.''')

          if offset and not do_rescale:
              raise ValueError('''For offset, do_rescale must also be set to True.''')

          # All transformations expect numpy arrays.
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		to_numpy_array(lowercase_)

          if do_resize:
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.resize(image=lowercase_						,	size=lowercase_						,	resample=lowercase_)

          if do_center_crop:
              SCREAMING_SNAKE_CASE_      :  Dict    						=		self.center_crop(lowercase_						,	size=lowercase_)

          if do_rescale:
              SCREAMING_SNAKE_CASE_      :  int    						=		self.rescale(image=lowercase_						,	scale=lowercase_						,	offset=lowercase_)

          if do_normalize:
              SCREAMING_SNAKE_CASE_      :  Dict    						=		self.normalize(image=lowercase_						,	mean=lowercase_						,	std=lowercase_)

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		to_channel_dimension_format(lowercase_						,	lowercase_)
          return image



      def _SCREAMING_SNAKE_CASE (						self     :  Dict						,	lowercase_     :  ImageInput						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = None						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = None						,	lowercase_     :  float = None						,	lowercase_     :  bool = None						,	lowercase_     :  bool = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[str, TensorType]] = None						,	lowercase_     :  ChannelDimension = ChannelDimension.FIRST						,	**lowercase_     :  Optional[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		do_resize if do_resize is not None else self.do_resize
          SCREAMING_SNAKE_CASE_      :  int    						=		resample if resample is not None else self.resample
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		do_center_crop if do_center_crop is not None else self.do_center_crop
          SCREAMING_SNAKE_CASE_      :  Dict    						=		do_rescale if do_rescale is not None else self.do_rescale
          SCREAMING_SNAKE_CASE_      :  Dict    						=		rescale_factor if rescale_factor is not None else self.rescale_factor
          SCREAMING_SNAKE_CASE_      :  Dict    						=		offset if offset is not None else self.offset
          SCREAMING_SNAKE_CASE_      :  str    						=		do_normalize if do_normalize is not None else self.do_normalize
          SCREAMING_SNAKE_CASE_      :  Dict    						=		image_mean if image_mean is not None else self.image_mean
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		image_std if image_std is not None else self.image_std

          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		size if size is not None else self.size
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Any    						=		crop_size if crop_size is not None else self.crop_size
          SCREAMING_SNAKE_CASE_      :  Dict    						=		get_size_dict(lowercase_						,	param_name='''crop_size''')

          if not valid_images(lowercase_):
              raise ValueError(
                  '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                  '''torch.Tensor, tf.Tensor or jax.ndarray.''')

          SCREAMING_SNAKE_CASE_      :  Tuple    						=		make_batched(lowercase_)

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              [
                  self._preprocess_image(
                      image=lowercase_						,	do_resize=lowercase_						,	size=lowercase_						,	resample=lowercase_						,	do_center_crop=lowercase_						,	crop_size=lowercase_						,	do_rescale=lowercase_						,	rescale_factor=lowercase_						,	offset=lowercase_						,	do_normalize=lowercase_						,	image_mean=lowercase_						,	image_std=lowercase_						,	data_format=lowercase_						,	)
                  for img in video
              ]
              for video in videos
          ]

          SCREAMING_SNAKE_CASE_      :  int    						=		{'''pixel_values''': videos}
          return BatchFeature(data=lowercase_						,	tensor_type=lowercase_)


