"""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__  :  Dict						     =	[
    '''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__  :  str						     =	[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]



def   a_						(					lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							{
				    '''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__             =							int(re.match(r'.*layer_(\d*).*'					,			_a       )[1]       )
				layer_number -= 3
				return f'''h.{layer_number}.''' + key



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



def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				# Construct model
				if bloom_config_file == "":
								UpperCAmelCase__             =							BloomConfig()
				else:
								UpperCAmelCase__             =							BloomConfig.from_json_file(_a       )

				if shard_model:
								UpperCAmelCase__             =							os.listdir(_a       )
								UpperCAmelCase__             =							sorted(filter(lambda lowerCamelCase       : s.startswith('layer'       ) and "model_00" in s					,			_a       )       )

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

								UpperCAmelCase__             =							None

								UpperCAmelCase__             =							BloomConfig()

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

												for i in range(_a       ):
																# load all TP files
																UpperCAmelCase__             =							file.replace('model_00'					,			f'''model_0{i}'''       )
																UpperCAmelCase__             =							torch.load(os.path.join(_a					,			_a       )					,			map_location='cpu'       )

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

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

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

								UpperCAmelCase__             =							os.listdir(_a       )
								UpperCAmelCase__             =							sorted(filter(lambda lowerCamelCase       : s.startswith('layer'       ) and "model_00" in s					,			_a       )       )

								UpperCAmelCase__             =							None
								for i, file in enumerate(_a       ):
												UpperCAmelCase__             =							None
												for i in range(_a       ):
																# load all TP files
																UpperCAmelCase__             =							file.replace('model_00'					,			f'''model_0{i}'''       )
																UpperCAmelCase__             =							torch.load(os.path.join(_a					,			_a       )					,			map_location='cpu'       )

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

																if tensors is None:
																				UpperCAmelCase__             =							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__             =							1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN       ) else 0
																												# We concatenate these weights accross TP ranks
																												UpperCAmelCase__             =							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__             =							tensors[key] / pretraining_tp

												UpperCAmelCase__             =							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__             =							set(other_keys.missing_keys       )
												else:
																UpperCAmelCase__             =							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__             =							pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
								UpperCAmelCase__             =							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__             =							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__":
						lowerCAmelCase__  :  str						     =	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__  :  List[Any]						     =	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

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
         OpenAIGPTConfig,
         OpenAIGPTDoubleHeadsModel,
         OpenAIGPTForSequenceClassification,
         OpenAIGPTLMHeadModel,
         OpenAIGPTModel,
     )




class        _UpperCAmelCase       :

 def __init__(     self		:       Union[str, Any]  ,     lowercase_		:       List[Any]  ,     lowercase_		:       int=13  ,     lowercase_		:       Optional[int]=7  ,     lowercase_		:       Any=True  ,     lowercase_		:       Dict=True  ,     lowercase_		:       Dict=True  ,     lowercase_		:       Optional[Any]=99  ,     lowercase_		:       Union[str, Any]=32  ,     lowercase_		:       str=5  ,     lowercase_		:       Union[str, Any]=4  ,     lowercase_		:       Any=37  ,     lowercase_		:       Tuple="gelu"  ,     lowercase_		:       Dict=0.1  ,     lowercase_		:       Tuple=0.1  ,     lowercase_		:       Optional[int]=512  ,     lowercase_		:       Optional[Any]=16  ,     lowercase_		:       Optional[Any]=2  ,     lowercase_		:       Optional[Any]=0.02  ,     lowercase_		:       List[Any]=3  ,     lowercase_		:       Union[str, Any]=4  ,     lowercase_		:       List[Any]=None  ,     ):
     snake_case_ :     Any				=	parent
     snake_case_ :     List[str]				=	batch_size
     snake_case_ :     List[Any]				=	seq_length
     snake_case_ :     Optional[int]				=	is_training
     snake_case_ :     Union[str, Any]				=	use_token_type_ids
     snake_case_ :     Optional[Any]				=	use_labels
     snake_case_ :     Union[str, Any]				=	vocab_size
     snake_case_ :     Any				=	hidden_size
     snake_case_ :     List[Any]				=	num_hidden_layers
     snake_case_ :     Any				=	num_attention_heads
     snake_case_ :     Dict				=	intermediate_size
     snake_case_ :     Union[str, Any]				=	hidden_act
     snake_case_ :     Optional[int]				=	hidden_dropout_prob
     snake_case_ :     Optional[Any]				=	attention_probs_dropout_prob
     snake_case_ :     Tuple				=	max_position_embeddings
     snake_case_ :     int				=	type_vocab_size
     snake_case_ :     Tuple				=	type_sequence_label_size
     snake_case_ :     str				=	initializer_range
     snake_case_ :     Tuple				=	num_labels
     snake_case_ :     str				=	num_choices
     snake_case_ :     Any				=	scope
     snake_case_ :     Dict				=	self.vocab_size - 1

 def 						_snake_case					(     self		:       int				):
     snake_case_ :     Optional[Any]				=	ids_tensor([self.batch_size, self.seq_length]  ,     self.vocab_size				)

     snake_case_ :     Optional[Any]				=	None
     if self.use_token_type_ids:
         snake_case_ :     Tuple				=	ids_tensor([self.batch_size, self.seq_length]  ,     self.type_vocab_size				)

     snake_case_ :     str				=	None
     snake_case_ :     Dict				=	None
     snake_case_ :     str				=	None
     if self.use_labels:
         snake_case_ :     Optional[int]				=	ids_tensor([self.batch_size]  ,     self.type_sequence_label_size				)
         snake_case_ :     List[Any]				=	ids_tensor([self.batch_size, self.seq_length]  ,     self.num_labels				)
         snake_case_ :     Tuple				=	ids_tensor([self.batch_size]  ,     self.num_choices				)

     snake_case_ :     int				=	OpenAIGPTConfig(
         vocab_size=self.vocab_size  ,     n_embd=self.hidden_size  ,     n_layer=self.num_hidden_layers  ,     n_head=self.num_attention_heads  ,     n_positions=self.max_position_embeddings  ,     pad_token_id=self.pad_token_id  ,     )

     snake_case_ :     Any				=	ids_tensor([self.num_hidden_layers, self.num_attention_heads]  ,     2				)

     return (
         config,
         input_ids,
         head_mask,
         token_type_ids,
         sequence_labels,
         token_labels,
         choice_labels,
     )

 def 						_snake_case					(     self		:       Tuple  ,     lowercase_		:       Any  ,     lowercase_		:       Union[str, Any]  ,     lowercase_		:       str  ,     lowercase_		:       Dict  ,     *lowercase_		:       Dict				):
     snake_case_ :     List[Any]				=	OpenAIGPTModel(config=lowercase_				)
     model.to(lowercase_				)
     model.eval()

     snake_case_ :     Any				=	model(lowercase_  ,     token_type_ids=lowercase_  ,     head_mask=lowercase_				)
     snake_case_ :     Optional[Any]				=	model(lowercase_  ,     token_type_ids=lowercase_				)
     snake_case_ :     Optional[Any]				=	model(lowercase_				)

     self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length, self.hidden_size)				)

 def 						_snake_case					(     self		:       Tuple  ,     lowercase_		:       Dict  ,     lowercase_		:       str  ,     lowercase_		:       Optional[Any]  ,     lowercase_		:       List[Any]  ,     *lowercase_		:       Optional[Any]				):
     snake_case_ :     Union[str, Any]				=	OpenAIGPTLMHeadModel(lowercase_				)
     model.to(lowercase_				)
     model.eval()

     snake_case_ :     Union[str, Any]				=	model(lowercase_  ,     token_type_ids=lowercase_  ,     labels=lowercase_				)
     self.parent.assertEqual(result.loss.shape  ,     ()				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)

 def 						_snake_case					(     self		:       List[str]  ,     lowercase_		:       Dict  ,     lowercase_		:       List[str]  ,     lowercase_		:       Any  ,     lowercase_		:       Dict  ,     *lowercase_		:       Union[str, Any]				):
     snake_case_ :     Tuple				=	OpenAIGPTDoubleHeadsModel(lowercase_				)
     model.to(lowercase_				)
     model.eval()

     snake_case_ :     Dict				=	model(lowercase_  ,     token_type_ids=lowercase_  ,     labels=lowercase_				)
     self.parent.assertEqual(result.loss.shape  ,     ()				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)

 def 						_snake_case					(     self		:       Any  ,     lowercase_		:       str  ,     lowercase_		:       List[str]  ,     lowercase_		:       Optional[Any]  ,     lowercase_		:       Optional[Any]  ,     *lowercase_		:       Any				):
     snake_case_ :     int				=	self.num_labels
     snake_case_ :     Any				=	OpenAIGPTForSequenceClassification(lowercase_				)
     model.to(lowercase_				)
     model.eval()

     snake_case_ :     Optional[int]				=	ids_tensor([self.batch_size]  ,     self.type_sequence_label_size				)
     snake_case_ :     Optional[Any]				=	model(lowercase_  ,     token_type_ids=lowercase_  ,     labels=lowercase_				)
     self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_labels)				)






 def 						_snake_case					(     self		:       int				):
     snake_case_ :     Dict				=	self.prepare_config_and_inputs()
     (
         (
         snake_case_
     ),	(
         snake_case_
     ),	(
         snake_case_
     ),	(
         snake_case_
     ),	(
         snake_case_
     ),	(
         snake_case_
     ),	(
         snake_case_
     ),	
     ) :     str				=	config_and_inputs
     snake_case_ :     str				=	{
         '''input_ids''': input_ids,
         '''token_type_ids''': token_type_ids,
         '''head_mask''': head_mask,
     }

     return config, inputs_dict





@require_torch
class        _UpperCAmelCase       (			lowerCAmelCase__    ,     lowerCAmelCase__    ,     lowerCAmelCase__    ,     unittest.TestCase):
 _lowerCAmelCase   :			Dict							=		(
     (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
     if is_torch_available()
     else ()
 )
 _lowerCAmelCase   :			int							=		(
     (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
 )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
 _lowerCAmelCase   :			Union[str, Any]							=		(
     {
         """feature-extraction""": OpenAIGPTModel,
         """text-classification""": OpenAIGPTForSequenceClassification,
         """text-generation""": OpenAIGPTLMHeadModel,
         """zero-shot""": OpenAIGPTForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 def 						_snake_case					(     self		:       Tuple  ,     lowercase_		:       Optional[int]  ,     lowercase_		:       int  ,     lowercase_		:       List[Any]  ,     lowercase_		:       List[Any]  ,     lowercase_		:       Union[str, Any]				):
     if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
         # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
         # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
         # tiny config could not be created.
         return True

     return False

 def 						_snake_case					(     self		:       Optional[int]  ,     lowercase_		:       List[Any]  ,     lowercase_		:       Optional[int]  ,     lowercase_		:       List[str]=False				):
     snake_case_ :     Dict				=	super()._prepare_for_class(lowercase_  ,     lowercase_  ,     return_labels=lowercase_				)

     if return_labels:
         if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
             snake_case_ :     List[str]				=	torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)  ,     dtype=torch.long  ,     device=lowercase_  ,     )
             snake_case_ :     int				=	inputs_dict['''labels''']
             snake_case_ :     Optional[Any]				=	inputs_dict['''labels''']
             snake_case_ :     int				=	torch.zeros(
                 (self.model_tester.batch_size, self.model_tester.num_choices)  ,     dtype=torch.long  ,     device=lowercase_  ,     )
             snake_case_ :     Tuple				=	torch.zeros(
                 self.model_tester.batch_size  ,     dtype=torch.long  ,     device=lowercase_				)
     return inputs_dict

 def 						_snake_case					(     self		:       Any				):
     snake_case_ :     List[str]				=	OpenAIGPTModelTester(self				)
     snake_case_ :     Dict				=	ConfigTester(self  ,     config_class=lowercase_  ,     n_embd=37				)

 def 						_snake_case					(     self		:       List[str]				):
     self.config_tester.run_common_tests()

 def 						_snake_case					(     self		:       Optional[Any]				):
     snake_case_ :     Tuple				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_openai_gpt_model(*lowercase_				)

 def 						_snake_case					(     self		:       List[str]				):
     snake_case_ :     List[str]				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_lm_head_model(*lowercase_				)

 def 						_snake_case					(     self		:       int				):
     snake_case_ :     Any				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_double_lm_head_model(*lowercase_				)

 def 						_snake_case					(     self		:       List[str]				):
     snake_case_ :     Optional[int]				=	self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*lowercase_				)






 @slow
 def 						_snake_case					(     self		:       Dict				):
     for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case_ :     Optional[Any]				=	OpenAIGPTModel.from_pretrained(lowercase_				)
         self.assertIsNotNone(lowercase_				)





@require_torch
class        _UpperCAmelCase       (			unittest.TestCase):

 @slow
 def 						_snake_case					(     self		:       Optional[int]				):
     snake_case_ :     Optional[Any]				=	OpenAIGPTLMHeadModel.from_pretrained('''openai-gpt'''				)
     model.to(lowercase_				)
     snake_case_ :     List[str]				=	torch.tensor([[481, 4735, 544]]  ,     dtype=torch.long  ,     device=lowercase_				)  # the president is
     snake_case_ :     List[Any]				=	[
         481,
         4735,
         544,
         246,
         963,
         870,
         762,
         239,
         244,
         40477,
         244,
         249,
         719,
         881,
         487,
         544,
         240,
         244,
         603,
         481,
     ]  # the president is a very good man. " \n " i\'m sure he is, " said the

     snake_case_ :     Optional[Any]				=	model.generate(lowercase_  ,     do_sample=lowercase_				)
     self.assertListEqual(output_ids[0].tolist()  ,     lowercase_				)





