





import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig






def 				_lowerCAmelCase			(	_a :       str       , _a :       Any       , _a :       Union[str, Any]=0      )						->      Any:
     # Format the message.
     if name is None:
          lowerCAmelCase_       :							List[Any]			 =		None
     else:
          lowerCAmelCase_       :							Union[str, Any]			 =		""".""" * max(0       , spaces - 2      ) + """# {:""" + str(50 - spaces      ) + """s}"""
          lowerCAmelCase_       :							Tuple			 =		fmt.format(_a      )

     # Print and recurse (if needed).
     if isinstance(_a       , _a      ):
          if msg is not None:
               print(_a      )
          for k in val.keys():
               recursive_print(_a       , val[k]       , spaces + 2      )
     elif isinstance(_a       , torch.Tensor      ):
          print(_a       , """:"""       , val.size()      )
     else:
          print(_a       , """:"""       , _a      )






def 				_lowerCAmelCase			(	_a :       str       , _a :       Optional[int]       , _a :       str       , _a :       Dict      )						->      Tuple:
     # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
     # for compatibility with later versions of NVIDIA Megatron-LM.
     # The inverse operation is performed inside Megatron-LM to read checkpoints:
     # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
     # If param is the weight tensor of the self-attention block, the returned tensor
     # will have to be transposed one more time to be read by HuggingFace GPT2.
     lowerCAmelCase_       :							Dict			 =		param.size()
     # other versions store [num_heads * num_splits * hidden_size, :]
     lowerCAmelCase_       :							List[Any]			 =		(num_heads, num_splits, hidden_size) + input_shape[1:]
     lowerCAmelCase_       :							Union[str, Any]			 =		param.view(*_a      )
     lowerCAmelCase_       :							Any			 =		param.transpose(0       , 1      ).contiguous()
     lowerCAmelCase_       :							str			 =		param.view(*_a      )
     return param






def 				_lowerCAmelCase			(	_a :       Tuple       , _a :       Union[str, Any]      )						->      Any:
     lowerCAmelCase_       :							int			 =		config.n_positions
     lowerCAmelCase_       :							Union[str, Any]			 =		config.n_layer
     lowerCAmelCase_       :							List[str]			 =		config.vocab_size
     lowerCAmelCase_       :							Optional[int]			 =		config.n_head
     lowerCAmelCase_       :							Dict			 =		config.n_embd // config.n_head

     lowerCAmelCase_       :							Tuple			 =		sd_megatron["""model.language_model.embedding.word_embeddings.weight"""][:vocab_size, :]
     lowerCAmelCase_       :							Union[str, Any]			 =		{
         """transformer.wte.weight""": word_embeddings,
         """transformer.wpe.weight""": sd_megatron["""model.language_model.embedding.position_embeddings.weight"""],
         """transformer.ln_f.weight""": sd_megatron["""model.language_model.encoder.final_layernorm.weight"""],
         """transformer.ln_f.bias""": sd_megatron["""model.language_model.encoder.final_layernorm.bias"""],
     }

     lowerCAmelCase_       :							Tuple			 =		"""model.language_model.encoder.layers."""
     for i in range(_a      ):
          lowerCAmelCase_       :							Optional[int]			 =		torch.tril(torch.ones((n_positions, n_positions)       , dtype=torch.bool      )      )
          lowerCAmelCase_       :							Optional[Any]			 =		causal_mask.view(1       , 1       , _a       , _a      )
          lowerCAmelCase_       :							Dict			 =		causal_mask
          lowerCAmelCase_       :							Union[str, Any]			 =		torch.tensor(-1E4       , dtype=torch.bfloataa      )

          lowerCAmelCase_       :							str			 =		sd_megatron[F'{pf}{i}.input_layernorm.weight']
          lowerCAmelCase_       :							List[str]			 =		sd_megatron[F'{pf}{i}.input_layernorm.bias']

          lowerCAmelCase_       :							Any			 =		sd_megatron[F'{pf}{i}.self_attention.query_key_value.weight']
          lowerCAmelCase_       :							Any			 =		fix_query_key_value_ordering(_a       , 3       , _a       , _a      )
          lowerCAmelCase_       :							Union[str, Any]			 =		vala.transpose(0       , 1      ).contiguous()

          lowerCAmelCase_       :							str			 =		sd_megatron[F'{pf}{i}.self_attention.query_key_value.bias']
          lowerCAmelCase_       :							Any			 =		fix_query_key_value_ordering(_a       , 3       , _a       , _a      )
          lowerCAmelCase_       :							Optional[int]			 =		vala

          lowerCAmelCase_       :							Optional[int]			 =		sd_megatron[F'{pf}{i}.self_attention.dense.weight'].transpose(
              0       , 1      )
          lowerCAmelCase_       :							str			 =		sd_megatron[F'{pf}{i}.self_attention.dense.bias']
          lowerCAmelCase_       :							Tuple			 =		sd_megatron[F'{pf}{i}.post_attention_layernorm.weight']
          lowerCAmelCase_       :							Any			 =		sd_megatron[F'{pf}{i}.post_attention_layernorm.bias']
          lowerCAmelCase_       :							int			 =		sd_megatron[F'{pf}{i}.mlp.dense_h_to_4h.weight'].transpose(0       , 1      )
          lowerCAmelCase_       :							int			 =		sd_megatron[F'{pf}{i}.mlp.dense_h_to_4h.bias']
          lowerCAmelCase_       :							str			 =		sd_megatron[F'{pf}{i}.mlp.dense_4h_to_h.weight'].transpose(
              0       , 1      )
          lowerCAmelCase_       :							Union[str, Any]			 =		sd_megatron[F'{pf}{i}.mlp.dense_4h_to_h.bias']

     # For LM head, transformers' wants the matrix to weight embeddings.
     lowerCAmelCase_       :							List[str]			 =		word_embeddings

     return sd_hf






def 				_lowerCAmelCase			(	_a :       Tuple       , _a :       int      )						->      str:
     lowerCAmelCase_       :							List[str]			 =		6_40_00
     lowerCAmelCase_       :							List[str]			 =		config_megatron["""encoder_seq_length"""]
     lowerCAmelCase_       :							Any			 =		config_megatron["""hidden_size"""]
     lowerCAmelCase_       :							Union[str, Any]			 =		config_megatron["""num_layers"""]
     lowerCAmelCase_       :							Optional[Any]			 =		config_megatron["""num_attention_heads"""]
     lowerCAmelCase_       :							Union[str, Any]			 =		config_megatron["""ffn_hidden_size"""]
     lowerCAmelCase_       :							str			 =		"""gelu"""
     lowerCAmelCase_       :							str			 =		0.1
     lowerCAmelCase_       :							str			 =		0.1
     lowerCAmelCase_       :							Dict			 =		0.1
     lowerCAmelCase_       :							int			 =		config_megatron["""layernorm_epsilon"""]  # 1e-5
     lowerCAmelCase_       :							int			 =		config_megatron["""init_method_std"""]  # 0.02
     lowerCAmelCase_       :							int			 =		config_megatron["""apply_query_key_layer_scaling"""]  # True
     lowerCAmelCase_       :							List[Any]			 =		True
     lowerCAmelCase_       :							str			 =		True

     # This identifies the 6.7B (7B) model which uses a different tokenizer
     if config_megatron["hidden_size"] == 40_96:
          lowerCAmelCase_       :							List[Any]			 =		1  # <|endoftext|>
          lowerCAmelCase_       :							Dict			 =		1  # <|endoftext|>
          lowerCAmelCase_       :							Tuple			 =		0  # <unk>
     else:
          lowerCAmelCase_       :							Optional[Any]			 =		2  # <s>
          lowerCAmelCase_       :							List[Any]			 =		3  # <|endoftext|>
          lowerCAmelCase_       :							int			 =		0  # <pad>

     return config_hf






def 				_lowerCAmelCase			(	_a :       int      )						->      int:
     print(_a      )

     lowerCAmelCase_       :							Union[str, Any]			 =		args.checkpoint_path
     lowerCAmelCase_       :							Optional[Any]			 =		args.save_path
     if isfile(_a      ):
          raise FileNotFoundError(F'ERROR! could not find file {checkpoint_path}'      )

     # Load the model.
     lowerCAmelCase_       :							Union[str, Any]			 =		torch.load(_a       , map_location="""cpu"""      )

     # Load the config.
     lowerCAmelCase_       :							List[str]			 =		checkpoint["""hyper_parameters"""]["""cfg"""]
     lowerCAmelCase_       :							Dict			 =		GPTaConfig()
     lowerCAmelCase_       :							Any			 =		copy_config(config_hf=_a       , config_megatron=_a      )
     lowerCAmelCase_       :							Tuple			 =		["""GPT2LMHeadModel"""]

     lowerCAmelCase_       :							str			 =		checkpoint["""state_dict"""]

     # Convert.
     print("""Converting"""      )
     lowerCAmelCase_       :							Tuple			 =		convert_megatron_checkpoint(_a       , _a      )

     # Print the structure of converted state dict.
     if args.print_checkpoint_structure:
          recursive_print(_a       , _a      )

     lowerCAmelCase_       :							List[str]			 =		"""GPTSw3Tokenizer"""

     # Store the config to file.
     print("""Saving config"""      )
     config_hf.save_pretrained(_a      )

     # Store the state_dict to file.
     lowerCAmelCase_       :							List[str]			 =		os.path.join(_a       , """pytorch_model.bin"""      )
     print(F'Saving checkpoint to "{output_checkpoint_file}"'      )
     torch.save(_a       , _a      )


if __name__ == "__main__":
      UpperCAmelCase_       : Tuple							=    argparse.ArgumentParser()
      parser.add_argument(
          """--checkpoint_path""",
          type=str,
          required=True,
          help="""e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000""",
      )
      parser.add_argument("""--save_path""", type=str, required=True, help="""e.g. /home/user/gpt-sw3/hf""")
      parser.add_argument("""--print-checkpoint-structure""", action="""store_true""")
      UpperCAmelCase_       : Optional[int]							=    parser.parse_args()
      main(_args)







import numpy as np






def 				_lowerCAmelCase			(	_a :       Optional[int]       , _a :       str      )						->      Union[str, Any]:
     lowerCAmelCase_       :							int			 =		np.array(_a      )
     lowerCAmelCase_       :							str			 =		np.array(_a      )

     lowerCAmelCase_       :							Optional[Any]			 =		abs(predict - actual      )
     lowerCAmelCase_       :							Dict			 =		difference.mean()

     return score






def 				_lowerCAmelCase			(	_a :       Union[str, Any]       , _a :       int      )						->      Tuple:
     lowerCAmelCase_       :							int			 =		np.array(_a      )
     lowerCAmelCase_       :							int			 =		np.array(_a      )

     lowerCAmelCase_       :							Union[str, Any]			 =		predict - actual
     lowerCAmelCase_       :							Tuple			 =		np.square(_a      )

     lowerCAmelCase_       :							Tuple			 =		square_diff.mean()
     return score






def 				_lowerCAmelCase			(	_a :       Optional[Any]       , _a :       Tuple      )						->      Optional[int]:
     lowerCAmelCase_       :							Optional[Any]			 =		np.array(_a      )
     lowerCAmelCase_       :							Optional[Any]			 =		np.array(_a      )

     lowerCAmelCase_       :							int			 =		predict - actual
     lowerCAmelCase_       :							Any			 =		np.square(_a      )
     lowerCAmelCase_       :							Union[str, Any]			 =		square_diff.mean()
     lowerCAmelCase_       :							Optional[Any]			 =		np.sqrt(_a      )
     return score






def 				_lowerCAmelCase			(	_a :       Optional[Any]       , _a :       Optional[Any]      )						->      List[Any]:
     lowerCAmelCase_       :							Dict			 =		np.array(_a      )
     lowerCAmelCase_       :							List[Any]			 =		np.array(_a      )

     lowerCAmelCase_       :							str			 =		np.log(predict + 1      )
     lowerCAmelCase_       :							Any			 =		np.log(actual + 1      )

     lowerCAmelCase_       :							List[str]			 =		log_predict - log_actual
     lowerCAmelCase_       :							Dict			 =		np.square(_a      )
     lowerCAmelCase_       :							Optional[Any]			 =		square_diff.mean()

     lowerCAmelCase_       :							int			 =		np.sqrt(_a      )

     return score






def 				_lowerCAmelCase			(	_a :       str       , _a :       List[Any]      )						->      int:
     lowerCAmelCase_       :							Any			 =		np.array(_a      )
     lowerCAmelCase_       :							Optional[int]			 =		np.array(_a      )

     lowerCAmelCase_       :							Optional[int]			 =		predict - actual
     lowerCAmelCase_       :							str			 =		np.sum(_a      ) / len(_a      )
     lowerCAmelCase_       :							List[str]			 =		np.sum(_a      ) / len(_a      )
     # print(numerator, denumerator)
     lowerCAmelCase_       :							int			 =		float(_a      ) / denumerator * 1_00

     return score







def 				_lowerCAmelCase			(	_a :       Any       , _a :       Union[str, Any]      )						->      str:
     return np.mean(np.array(_a      ) == np.array(_a      )      )
