


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







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







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					):
			# Construct model
			if bloom_config_file == "":
						UpperCAmelCase__  :					Dict        =     BloomConfig()
			else:
						UpperCAmelCase__  :					Any        =     BloomConfig.from_json_file(UpperCamelCase__					)

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

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

						UpperCAmelCase__  :					Tuple        =     None

						UpperCAmelCase__  :					Union[str, Any]        =     BloomConfig()

						for j, file in enumerate(UpperCamelCase__					):
									print("""Processing file: {}""".format(UpperCamelCase__					)					)
									UpperCAmelCase__  :					List[str]        =     None

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

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

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

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

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

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

												# Rename keys in the transformers names
												UpperCAmelCase__  :					Any        =     list(temp.keys()					)
												for key in keys:
															UpperCAmelCase__  :					int        =     temp.pop(UpperCamelCase__					)

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

									UpperCAmelCase__  :					str        =     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:
												UpperCAmelCase__  :					Optional[int]        =     set(other_keys.missing_keys					)
									else:
												UpperCAmelCase__  :					Union[str, Any]        =     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__					)
						UpperCAmelCase__  :					int        =     pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
						UpperCAmelCase__  :					Dict        =     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() ,      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__":
		__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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A	=logging.get_logger(__name__)

__A	={
    'xlm-mlm-en-2048': 'https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json',
    'xlm-mlm-ende-1024': 'https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json',
    'xlm-mlm-enfr-1024': 'https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json',
    'xlm-mlm-enro-1024': 'https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json',
    'xlm-mlm-tlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json',
    'xlm-mlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json',
    'xlm-clm-enfr-1024': 'https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json',
    'xlm-clm-ende-1024': 'https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json',
    'xlm-mlm-17-1280': 'https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json',
    'xlm-mlm-100-1280': 'https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json',
}


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Any   =					'''xlm'''
					lowerCAmelCase      :Any   =					{
					    '''hidden_size''': '''emb_dim''',
					    '''num_attention_heads''': '''n_heads''',
					    '''num_hidden_layers''': '''n_layers''',
					    '''n_words''': '''vocab_size''',  # For backward compatibility
					}




					def __init__(      self							,							_lowerCamelCase=3_0145							,							_lowerCamelCase=2048							,							_lowerCamelCase=12							,							_lowerCamelCase=16							,							_lowerCamelCase=0.1							,							_lowerCamelCase=0.1							,							_lowerCamelCase=True							,							_lowerCamelCase=False							,							_lowerCamelCase=False							,							_lowerCamelCase=False							,							_lowerCamelCase=1							,							_lowerCamelCase=True							,							_lowerCamelCase=512							,							_lowerCamelCase=2048**-0.5							,							_lowerCamelCase=1e-1_2							,							_lowerCamelCase=0.02							,							_lowerCamelCase=0							,							_lowerCamelCase=1							,							_lowerCamelCase=2							,							_lowerCamelCase=3							,							_lowerCamelCase=5							,							_lowerCamelCase=True							,							_lowerCamelCase="first"							,							_lowerCamelCase=True							,							_lowerCamelCase=None							,							_lowerCamelCase=True							,							_lowerCamelCase=0.1							,							_lowerCamelCase=5							,							_lowerCamelCase=5							,							_lowerCamelCase=0							,							_lowerCamelCase=0							,							_lowerCamelCase=2							,							_lowerCamelCase=0							,							**_lowerCamelCase							,							):
								UpperCAmelCase__  :					Tuple        =     vocab_size
								UpperCAmelCase__  :					Tuple        =     emb_dim
								UpperCAmelCase__  :					Optional[Any]        =     n_layers
								UpperCAmelCase__  :					List[str]        =     n_heads
								UpperCAmelCase__  :					List[str]        =     dropout
								UpperCAmelCase__  :					Optional[int]        =     attention_dropout
								UpperCAmelCase__  :					Tuple        =     gelu_activation
								UpperCAmelCase__  :					Optional[Any]        =     sinusoidal_embeddings
								UpperCAmelCase__  :					int        =     causal
								UpperCAmelCase__  :					Union[str, Any]        =     asm
								UpperCAmelCase__  :					Optional[Any]        =     n_langs
								UpperCAmelCase__  :					List[Any]        =     use_lang_emb
								UpperCAmelCase__  :					Union[str, Any]        =     layer_norm_eps
								UpperCAmelCase__  :					List[str]        =     bos_index
								UpperCAmelCase__  :					List[Any]        =     eos_index
								UpperCAmelCase__  :					int        =     pad_index
								UpperCAmelCase__  :					str        =     unk_index
								UpperCAmelCase__  :					Dict        =     mask_index
								UpperCAmelCase__  :					str        =     is_encoder
								UpperCAmelCase__  :					Dict        =     max_position_embeddings
								UpperCAmelCase__  :					Any        =     embed_init_std
								UpperCAmelCase__  :					List[Any]        =     init_std
								UpperCAmelCase__  :					List[str]        =     summary_type
								UpperCAmelCase__  :					Union[str, Any]        =     summary_use_proj
								UpperCAmelCase__  :					Any        =     summary_activation
								UpperCAmelCase__  :					List[str]        =     summary_proj_to_labels
								UpperCAmelCase__  :					Union[str, Any]        =     summary_first_dropout
								UpperCAmelCase__  :					str        =     start_n_top
								UpperCAmelCase__  :					str        =     end_n_top
								UpperCAmelCase__  :					Tuple        =     mask_token_id
								UpperCAmelCase__  :					Union[str, Any]        =     lang_id

								if "n_words" in kwargs:
											UpperCAmelCase__  :					List[str]        =     kwargs["""n_words"""]

								super().__init__(pad_token_id=_lowerCamelCase							,							bos_token_id=_lowerCamelCase							,							**_lowerCamelCase)


class 		_snake_case					(     a__	):




					@property
					def 				snake_case__						(      self):
								if self.task == "multiple-choice":
											UpperCAmelCase__  :					Any        =     {0: """batch""", 1: """choice""", 2: """sequence"""}
								else:
											UpperCAmelCase__  :					List[Any]        =     {0: """batch""", 1: """sequence"""}
								return OrderedDict(
								    [
								        ("""input_ids""", dynamic_axis),
								        ("""attention_mask""", dynamic_axis),
								        ("""token_type_ids""", dynamic_axis),
								    ])