
'''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__      :		Tuple       							=  [
    "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__      :		Dict       							=  [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[int]:




			"""simple docstring"""
			_SCREAMING_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
			_SCREAMING_SNAKE_CASE      =				int(re.match(r""".*layer_(\d*).*"""    ,      a_						)[1]						)
			layer_number -= 3
			return F"h.{layer_number}." + key
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_						)     ->     int:




			"""simple docstring"""






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




			"""simple docstring"""






			if bloom_config_file == "":
						_SCREAMING_SNAKE_CASE      =				BloomConfig()
			else:
						_SCREAMING_SNAKE_CASE      =				BloomConfig.from_json_file(a_						)

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

						_SCREAMING_SNAKE_CASE      =				{'''weight_map''': {}, '''metadata''': {}}
						_SCREAMING_SNAKE_CASE      =				0

						_SCREAMING_SNAKE_CASE      =				None

						_SCREAMING_SNAKE_CASE      =				BloomConfig()

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

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

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

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

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

						_SCREAMING_SNAKE_CASE      =				os.listdir(a_						)
						_SCREAMING_SNAKE_CASE      =				sorted(filter(lambda SCREAMING_SNAKE_CASE_						: s.startswith("""layer"""						) and "model_00" in s    ,      a_						)						)

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

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

												if tensors is None:
															_SCREAMING_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(a_						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					_SCREAMING_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
																					_SCREAMING_SNAKE_CASE      =				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						):
															_SCREAMING_SNAKE_CASE      =				tensors[key] / pretraining_tp

									_SCREAMING_SNAKE_CASE      =				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:
												_SCREAMING_SNAKE_CASE      =				set(other_keys.missing_keys						)
									else:
												_SCREAMING_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(a_    ,      exist_ok=a_						)
						_SCREAMING_SNAKE_CASE      =				pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
						_SCREAMING_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:
									_SCREAMING_SNAKE_CASE      =				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__":
					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__      :		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'''



import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
     import jax.numpy as jnp

if is_tf_available():
     import tensorflow as tf

if is_torch_available():
     import torch

class    _a				(unittest.TestCase):





       """simple docstring"""




       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				{
              """task_specific_params""": {
                  """summarization""": {"""length_penalty""": 1.0, """max_length""": 1_28, """min_length""": 12, """num_beams""": 4},
                  """summarization_cnn""": {"""length_penalty""": 2.0, """max_length""": 1_42, """min_length""": 56, """num_beams""": 4},
                  """summarization_xsum""": {"""length_penalty""": 1.0, """max_length""": 62, """min_length""": 11, """num_beams""": 6},
              }
          }
          _SCREAMING_SNAKE_CASE      =				{
              """task_specific_params.summarization.length_penalty""": 1.0,
              """task_specific_params.summarization.max_length""": 1_28,
              """task_specific_params.summarization.min_length""": 12,
              """task_specific_params.summarization.num_beams""": 4,
              """task_specific_params.summarization_cnn.length_penalty""": 2.0,
              """task_specific_params.summarization_cnn.max_length""": 1_42,
              """task_specific_params.summarization_cnn.min_length""": 56,
              """task_specific_params.summarization_cnn.num_beams""": 4,
              """task_specific_params.summarization_xsum.length_penalty""": 1.0,
              """task_specific_params.summarization_xsum.max_length""": 62,
              """task_specific_params.summarization_xsum.min_length""": 11,
              """task_specific_params.summarization_xsum.num_beams""": 6,
          }

          self.assertEqual(flatten_dict(A__      )       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          self.assertTrue(np.allclose(transpose(A__      )       ,							x.transpose()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          self.assertTrue(np.allclose(transpose(A__       ,							axes=(1, 2, 0)      )       ,							x.transpose((1, 2, 0)      )      )      )






       @require_torch
       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(transpose(A__      )       ,							transpose(A__      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(transpose(A__       ,							axes=(1, 2, 0)      )       ,							transpose(A__       ,							axes=(1, 2, 0)      ).numpy()      )      )






       @require_tf
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(transpose(A__      )       ,							transpose(A__      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(transpose(A__       ,							axes=(1, 2, 0)      )       ,							transpose(A__       ,							axes=(1, 2, 0)      ).numpy()      )      )






       @require_flax
       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(transpose(A__      )       ,							np.asarray(transpose(A__      )      )      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(transpose(A__       ,							axes=(1, 2, 0)      )       ,							np.asarray(transpose(A__       ,							axes=(1, 2, 0)      )      )      )      )






       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          self.assertTrue(np.allclose(reshape(A__       ,							(4, 3)      )       ,							np.reshape(A__       ,							(4, 3)      )      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          self.assertTrue(np.allclose(reshape(A__       ,							(12, 5)      )       ,							np.reshape(A__       ,							(12, 5)      )      )      )






       @require_torch
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(4, 3)      )       ,							reshape(A__       ,							(4, 3)      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(12, 5)      )       ,							reshape(A__       ,							(12, 5)      ).numpy()      )      )






       @require_tf
       def        UpperCamelCase				(							self      )    ->			Tuple:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(4, 3)      )       ,							reshape(A__       ,							(4, 3)      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(12, 5)      )       ,							reshape(A__       ,							(12, 5)      ).numpy()      )      )






       @require_flax
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(4, 3)      )       ,							np.asarray(reshape(A__       ,							(4, 3)      )      )      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4       ,							5      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(reshape(A__       ,							(12, 5)      )       ,							np.asarray(reshape(A__       ,							(12, 5)      )      )      )      )






       def        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							3       ,							4      )
          self.assertTrue(np.allclose(squeeze(A__      )       ,							np.squeeze(A__      )      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							4       ,							1       ,							5      )
          self.assertTrue(np.allclose(squeeze(A__       ,							axis=2      )       ,							np.squeeze(A__       ,							axis=2      )      )      )






       @require_torch
       def        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(squeeze(A__      )       ,							squeeze(A__      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							4       ,							1       ,							5      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(squeeze(A__       ,							axis=2      )       ,							squeeze(A__       ,							axis=2      ).numpy()      )      )






       @require_tf
       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(squeeze(A__      )       ,							squeeze(A__      ).numpy()      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							4       ,							1       ,							5      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(squeeze(A__       ,							axis=2      )       ,							squeeze(A__       ,							axis=2      ).numpy()      )      )






       @require_flax
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(squeeze(A__      )       ,							np.asarray(squeeze(A__      )      )      )      )

          _SCREAMING_SNAKE_CASE      =				np.random.randn(1       ,							4       ,							1       ,							5      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(squeeze(A__       ,							axis=2      )       ,							np.asarray(squeeze(A__       ,							axis=2      )      )      )      )






       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          self.assertTrue(np.allclose(expand_dims(A__       ,							axis=1      )       ,							np.expand_dims(A__       ,							axis=1      )      )      )






       @require_torch
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )
          self.assertTrue(np.allclose(expand_dims(A__       ,							axis=1      )       ,							expand_dims(A__       ,							axis=1      ).numpy()      )      )






       @require_tf
       def        UpperCamelCase				(							self      )    ->			str:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				tf.constant(A__      )
          self.assertTrue(np.allclose(expand_dims(A__       ,							axis=1      )       ,							expand_dims(A__       ,							axis=1      ).numpy()      )      )






       @require_flax
       def        UpperCamelCase				(							self      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				np.random.randn(3       ,							4      )
          _SCREAMING_SNAKE_CASE      =				jnp.array(A__      )
          self.assertTrue(np.allclose(expand_dims(A__       ,							axis=1      )       ,							np.asarray(expand_dims(A__       ,							axis=1      )      )      )      )






