





"""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        lowercase__     (       lowercase_		,lowercase_      )       ->		Union[str, Any]:




      """simple docstring"""

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




def        lowercase__     (       lowercase_      )       ->		Optional[Any]:




      """simple docstring"""

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




def        lowercase__     (       lowercase_		,lowercase_		,lowercase_		,lowercase_		,lowercase_      )       ->		Tuple:




      """simple docstring"""

      if bloom_config_file == "":
            _UpperCamelCase   :						Optional[Any]    =							BloomConfig()
      else:
            _UpperCamelCase   :						List[Any]    =							BloomConfig.from_json_file(lowercase_      )

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

            _UpperCamelCase   :						Optional[Any]    =							{"weight_map": {}, "metadata": {}}
            _UpperCamelCase   :						int    =							0

            _UpperCamelCase   :						Dict    =							None

            _UpperCamelCase   :						Optional[int]    =							BloomConfig()

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

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

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

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

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

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

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

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

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

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

                        if tensors is None:
                              _UpperCamelCase   :						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(lowercase_      ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH      ):
                                          tensors[key] += temp[key]
                                    else:
                                          # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                          _UpperCamelCase   :						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
                                          _UpperCamelCase   :						Optional[int]    =							torch.cat([tensors[key], temp[key]]		,dim=lowercase_      )

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

                  _UpperCamelCase   :						Dict    =							model.load_state_dict(lowercase_		,strict=lowercase_      )
                  assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
                  if missing_keys is None:
                        _UpperCamelCase   :						str    =							set(other_keys.missing_keys      )
                  else:
                        _UpperCamelCase   :						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(lowercase_		,exist_ok=lowercase_      )
            _UpperCamelCase   :						List[str]    =							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()		,lowercase_      )
            print(F'''Save configuration file to {pytorch_config_dump_path}'''      )
            with open(lowercase_		,"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 unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import GLPNImageProcessor



class   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


  def __init__(		self :							Any       , __a :							Any       , __a :							Optional[Any]=7       , __a :							Tuple=3       , __a :							int=18       , __a :							Optional[Any]=30       , __a :							Optional[Any]=400       , __a :							List[str]=True       , __a :							Any=32       , __a :							Tuple=True       , )		->					List[str]:
        _UpperCamelCase   :						List[str]    =							parent
        _UpperCamelCase   :						int    =							batch_size
        _UpperCamelCase   :						List[str]    =							num_channels
        _UpperCamelCase   :						Union[str, Any]    =							image_size
        _UpperCamelCase   :						Optional[Any]    =							min_resolution
        _UpperCamelCase   :						Union[str, Any]    =							max_resolution
        _UpperCamelCase   :						Optional[Any]    =							do_resize
        _UpperCamelCase   :						str    =							size_divisor
        _UpperCamelCase   :						Optional[Any]    =							do_rescale


  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					Union[str, Any]:
        return {
            "do_resize": self.do_resize,
            "size_divisor": self.size_divisor,
            "do_rescale": self.do_rescale,
        }





@require_torch
@require_vision
class   __SCREAMING_SNAKE_CASE     (    _UpperCamelCase		,	unittest.TestCase ):







  '''simple docstring'''


  SCREAMING_SNAKE_CASE__				:Any						  =       GLPNImageProcessor if is_vision_available() else None

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Any:
        _UpperCamelCase   :						Optional[Any]    =							GLPNImageProcessingTester(self						)

  @property
  def  __SCREAMING_SNAKE_CASE					(		self :							Optional[int]						)		->					Dict:
        return self.image_processor_tester.prepare_image_processor_dict()

  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					Union[str, Any]:
        _UpperCamelCase   :						int    =							self.image_processing_class(**self.image_processor_dict						)
        self.assertTrue(hasattr(__a       , "do_resize"						)						)
        self.assertTrue(hasattr(__a       , "size_divisor"						)						)
        self.assertTrue(hasattr(__a       , "resample"						)						)
        self.assertTrue(hasattr(__a       , "do_rescale"						)						)

  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					List[Any]:
        pass

  def  __SCREAMING_SNAKE_CASE					(		self :							int						)		->					Optional[Any]:
        # Initialize image_processing
        _UpperCamelCase   :						Optional[Any]    =							self.image_processing_class(**self.image_processor_dict						)
        # create random PIL images
        _UpperCamelCase   :						List[str]    =							prepare_image_inputs(self.image_processor_tester       , equal_resolution=__a						)
        for image in image_inputs:
              self.assertIsInstance(__a       , Image.Image						)

        # Test not batched input (GLPNImageProcessor doesn't support batching)
        _UpperCamelCase   :						Optional[int]    =							image_processing(image_inputs[0]       , return_tensors="pt"						).pixel_values
        self.assertTrue(encoded_images.shape[-1] % self.image_processor_tester.size_divisor == 0						)
        self.assertTrue(encoded_images.shape[-2] % self.image_processor_tester.size_divisor == 0						)

  def  __SCREAMING_SNAKE_CASE					(		self :							str						)		->					Dict:
        # Initialize image_processing
        _UpperCamelCase   :						Optional[int]    =							self.image_processing_class(**self.image_processor_dict						)
        # create random numpy tensors
        _UpperCamelCase   :						List[str]    =							prepare_image_inputs(self.image_processor_tester       , equal_resolution=__a       , numpify=__a						)
        for image in image_inputs:
              self.assertIsInstance(__a       , np.ndarray						)

        # Test not batched input (GLPNImageProcessor doesn't support batching)
        _UpperCamelCase   :						Dict    =							image_processing(image_inputs[0]       , return_tensors="pt"						).pixel_values
        self.assertTrue(encoded_images.shape[-1] % self.image_processor_tester.size_divisor == 0						)
        self.assertTrue(encoded_images.shape[-2] % self.image_processor_tester.size_divisor == 0						)


  def  __SCREAMING_SNAKE_CASE					(		self :							Any						)		->					Union[str, Any]:
        # Initialize image_processing
        _UpperCamelCase   :						Optional[int]    =							self.image_processing_class(**self.image_processor_dict						)
        # create random PyTorch tensors
        _UpperCamelCase   :						Dict    =							prepare_image_inputs(self.image_processor_tester       , equal_resolution=__a       , torchify=__a						)
        for image in image_inputs:
              self.assertIsInstance(__a       , torch.Tensor						)

        # Test not batched input (GLPNImageProcessor doesn't support batching)
        _UpperCamelCase   :						Optional[int]    =							image_processing(image_inputs[0]       , return_tensors="pt"						).pixel_values
        self.assertTrue(encoded_images.shape[-1] % self.image_processor_tester.size_divisor == 0						)
        self.assertTrue(encoded_images.shape[-2] % self.image_processor_tester.size_divisor == 0						)
