


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






def       lowerCamelCase					(a_						:Any)     ->							Union[str, Any]:
      if dtype == torch.bool:
            return 1 / 8
      lowercase		:Tuple       							=				re.search(R'''[^\d](\d+)$'''				,				str(UpperCAmelCase__))
      if bit_search is None:
            raise ValueError(F"""`dtype` is not a valid dtype: {dtype}.""")
      lowercase		:Optional[int]       							=				int(bit_search.groups()[0])
      return bit_size // 8






def       lowerCamelCase					(a_						:Dict				,				a_						:Optional[Any]				,				a_						:Optional[int]				,				a_						:Union[str, Any]				,				a_						:Optional[int])     ->							Optional[Any]:
      # Construct model
      if bloom_config_file == "":
            lowercase		:Union[str, Any]       							=				BloomConfig()
      else:
            lowercase		:int       							=				BloomConfig.from_json_file(UpperCAmelCase__)

      if shard_model:
            lowercase		:Optional[Any]       							=				os.listdir(UpperCAmelCase__)
            lowercase		:Union[str, Any]       							=				sorted(filter(lambda a_: s.startswith('''layer''') and "model_00" in s				,				UpperCAmelCase__))

            lowercase		:Union[str, Any]       							=				{"""weight_map""": {}, """metadata""": {}}
            lowercase		:Optional[int]       							=				0

            lowercase		:int       							=				None

            lowercase		:str       							=				BloomConfig()

            for j, file in enumerate(UpperCAmelCase__):
                  print('''Processing file: {}'''.format(UpperCAmelCase__))
                  lowercase		:Any       							=				None

                  for i in range(UpperCAmelCase__):
                        # load all TP files
                        lowercase		:Any       							=				file.replace('''model_00'''				,				F"""model_0{i}""")
                        lowercase		: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		:Any       							=				temp.pop(UpperCAmelCase__)

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

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

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

            lowercase		:Union[str, Any]       							=				None
            for i, file in enumerate(UpperCAmelCase__):
                  lowercase		:Tuple       							=				None
                  for i in range(UpperCAmelCase__):
                        # load all TP files
                        lowercase		:List[str]       							=				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		:int       							=				temp.pop(UpperCAmelCase__)

                        if tensors is None:
                              lowercase		:Dict       							=				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		:Optional[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		:int       							=				tensors[key] / pretraining_tp

                  lowercase		:Optional[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		:Union[str, Any]       							=				set(other_keys.missing_keys)
                  else:
                        lowercase		:str       							=				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		:Union[str, Any]       							=				pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
            lowercase		: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:
                  lowercase		:Union[str, Any]       							=				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"""



import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()
class 					__magic_name__		(			__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      Tuple        =   PriorTransformer
       __A			:      Any        =   "hidden_states"






       @property
       def  __snake_case						(   self  :    str				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				4
              lowercase		:List[Any]       							=				8
              lowercase		:Any       							=				7

              lowercase		:Optional[Any]       							=				floats_tensor((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:List[Any]       							=				floats_tensor((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Union[str, Any]       							=				floats_tensor((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       def  __snake_case						(   self  :    Dict  ,     snake_case__  :    Optional[Any]=0				):

              '''simple docstring'''






              torch.manual_seed(snake_case__				)
              lowercase		:Union[str, Any]       							=				4
              lowercase		:Optional[int]       							=				8
              lowercase		:str       							=				7

              lowercase		:Dict       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:Union[str, Any]       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Tuple       							=				torch.randn((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       @property
       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              return (4, 8)






       @property
       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              return (4, 8)






       def  __snake_case						(   self  :    Optional[Any]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				{
                  '''num_attention_heads''': 2,
                  '''attention_head_dim''': 4,
                  '''num_layers''': 2,
                  '''embedding_dim''': 8,
                  '''num_embeddings''': 7,
                  '''additional_embeddings''': 4,
              }
              lowercase		:List[str]       							=				self.dummy_input
              return init_dict, inputs_dict






       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase							,  lowercase		:Optional[int]       							=				PriorTransformer.from_pretrained(
                  '''hf-internal-testing/prior-dummy'''  ,     output_loading_info=snake_case__				)
              self.assertIsNotNone(snake_case__				)
              self.assertEqual(len(loading_info['''missing_keys''']				)  ,     0				)

              model.to(snake_case__				)
              lowercase		:List[Any]       							=				model(**self.dummy_input				)[0]

              assert hidden_states is not None, "Make sure output is not None"






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

              '''simple docstring'''






              lowercase							,  lowercase		:List[Any]       							=				self.prepare_init_args_and_inputs_for_common()

              lowercase		:str       							=				self.model_class(**snake_case__				)
              lowercase		:Optional[Any]       							=				inspect.signature(model.forward				)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              lowercase		:List[Any]       							=				[*signature.parameters.keys()]

              lowercase		:Optional[Any]       							=				['''hidden_states''', '''timestep''']
              self.assertListEqual(arg_names[:2]  ,     snake_case__				)







       def  __snake_case						(   self  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				PriorTransformer.from_pretrained('''hf-internal-testing/prior-dummy'''				)
              lowercase		:List[Any]       							=				model.to(snake_case__				)

              if hasattr(snake_case__  ,     '''set_default_attn_processor'''				):
                     model.set_default_attn_processor()

              lowercase		:List[str]       							=				self.get_dummy_seed_input()

              with torch.no_grad():
                     lowercase		:List[Any]       							=				model(**snake_case__				)[0]

              lowercase		:Dict       							=				output[0, :5].flatten().cpu()
              print(snake_case__				)

              # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
              # the expected output slices are not the same for CPU and GPU.
              lowercase		:str       							=				torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39]				)
              self.assertTrue(torch_all_close(snake_case__  ,     snake_case__  ,     rtol=1e-2				)				)







@slow
class 					__magic_name__		(			unittest.TestCase     ):






       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    int=1  ,     snake_case__  :    int=7_6_8  ,     snake_case__  :    Optional[Any]=7_7  ,     snake_case__  :    Union[str, Any]=0				):

              '''simple docstring'''






              torch.manual_seed(snake_case__				)
              lowercase		:Any       							=				batch_size
              lowercase		:Dict       							=				embedding_dim
              lowercase		:int       							=				num_embeddings

              lowercase		:int       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:List[str]       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Optional[Any]       							=				torch.randn((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()







       @parameterized.expand(
           [
               # fmt: off
               [1_3, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
               [3_7, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
               # fmt: on
           ]				)
       def  __snake_case						(   self  :    Dict  ,     snake_case__  :    Dict  ,     snake_case__  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:str       							=				PriorTransformer.from_pretrained('''kandinsky-community/kandinsky-2-1-prior'''  ,     subfolder='''prior'''				)
              model.to(snake_case__				)
              lowercase		:Union[str, Any]       							=				self.get_dummy_seed_input(seed=snake_case__				)

              with torch.no_grad():
                     lowercase		:int       							=				model(**snake_case__				)[0]

              assert list(sample.shape				) == [1, 7_6_8]

              lowercase		:List[Any]       							=				sample[0, :8].flatten().cpu()
              print(snake_case__				)
              lowercase		:Union[str, Any]       							=				torch.tensor(snake_case__				)

              assert torch_all_close(snake_case__  ,     snake_case__  ,     atol=1e-3				)

