





'''simple docstring'''

import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase,       _UpperCAmelCase=0  ):
		# Format the message.
		if name is None:
				__UpperCAmelCase	:							Optional[Any]     =     None
		else:
				__UpperCAmelCase	:							Union[str, Any]     =     "." * max(0,       spaces - 2  ) + "# {:" + str(50 - spaces  ) + "s}"
				__UpperCAmelCase	:							Any     =     fmt.format(_UpperCAmelCase  )

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


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase,       _UpperCAmelCase,       _UpperCAmelCase  ):
		# 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.
		__UpperCAmelCase	:							Optional[Any]     =     param.size()
		# other versions store [num_heads * num_splits * hidden_size, :]
		__UpperCAmelCase	:							Tuple     =     (num_heads, num_splits, hidden_size) + input_shape[1:]
		__UpperCAmelCase	:							Tuple     =     param.view(*_UpperCAmelCase  )
		__UpperCAmelCase	:							str     =     param.transpose(0,       1  ).contiguous()
		__UpperCAmelCase	:							str     =     param.view(*_UpperCAmelCase  )
		return param


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase  ):
		__UpperCAmelCase	:							Dict     =     config.n_positions
		__UpperCAmelCase	:							int     =     config.n_layer
		__UpperCAmelCase	:							str     =     config.vocab_size
		__UpperCAmelCase	:							Any     =     config.n_head
		__UpperCAmelCase	:							List[Any]     =     config.n_embd // config.n_head

		__UpperCAmelCase	:							Tuple     =     sd_megatron["model.language_model.embedding.word_embeddings.weight"][:vocab_size, :]
		__UpperCAmelCase	:							List[str]     =     {
		    "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"],
		}

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

				__UpperCAmelCase	:							Dict     =     sd_megatron[F"{pf}{i}.input_layernorm.weight"]
				__UpperCAmelCase	:							str     =     sd_megatron[F"{pf}{i}.input_layernorm.bias"]

				__UpperCAmelCase	:							Tuple     =     sd_megatron[F"{pf}{i}.self_attention.query_key_value.weight"]
				__UpperCAmelCase	:							Any     =     fix_query_key_value_ordering(_UpperCAmelCase,       3,       _UpperCAmelCase,       _UpperCAmelCase  )
				__UpperCAmelCase	:							Dict     =     vala.transpose(0,       1  ).contiguous()

				__UpperCAmelCase	:							Dict     =     sd_megatron[F"{pf}{i}.self_attention.query_key_value.bias"]
				__UpperCAmelCase	:							int     =     fix_query_key_value_ordering(_UpperCAmelCase,       3,       _UpperCAmelCase,       _UpperCAmelCase  )
				__UpperCAmelCase	:							List[Any]     =     vala

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

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

		return sd_hf


def 						__UpperCamelCase						(			_UpperCAmelCase,       _UpperCAmelCase  ):
		__UpperCAmelCase	:							Optional[Any]     =     64000
		__UpperCAmelCase	:							str     =     config_megatron["encoder_seq_length"]
		__UpperCAmelCase	:							Optional[int]     =     config_megatron["hidden_size"]
		__UpperCAmelCase	:							Optional[Any]     =     config_megatron["num_layers"]
		__UpperCAmelCase	:							List[str]     =     config_megatron["num_attention_heads"]
		__UpperCAmelCase	:							List[str]     =     config_megatron["ffn_hidden_size"]
		__UpperCAmelCase	:							int     =     "gelu"
		__UpperCAmelCase	:							Dict     =     0.1
		__UpperCAmelCase	:							List[Any]     =     0.1
		__UpperCAmelCase	:							Any     =     0.1
		__UpperCAmelCase	:							List[Any]     =     config_megatron["layernorm_epsilon"]  # 1e-5
		__UpperCAmelCase	:							int     =     config_megatron["init_method_std"]  # 0.02
		__UpperCAmelCase	:							int     =     config_megatron["apply_query_key_layer_scaling"]  # True
		__UpperCAmelCase	:							str     =     True
		__UpperCAmelCase	:							str     =     True

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

		return config_hf


def 						__UpperCamelCase						(			_UpperCAmelCase  ):
		print(_UpperCAmelCase  )

		__UpperCAmelCase	:							Optional[Any]     =     args.checkpoint_path
		__UpperCAmelCase	:							List[str]     =     args.save_path
		if isfile(_UpperCAmelCase  ):
				raise FileNotFoundError(F"ERROR! could not find file {checkpoint_path}"  )

		# Load the model.
		__UpperCAmelCase	:							Optional[Any]     =     torch.load(_UpperCAmelCase,       map_location="cpu"  )

		# Load the config.
		__UpperCAmelCase	:							Optional[Any]     =     checkpoint["hyper_parameters"]["cfg"]
		__UpperCAmelCase	:							Optional[Any]     =     GPTaConfig()
		__UpperCAmelCase	:							str     =     copy_config(config_hf=_UpperCAmelCase,       config_megatron=_UpperCAmelCase  )
		__UpperCAmelCase	:							str     =     ["GPT2LMHeadModel"]

		__UpperCAmelCase	:							Optional[int]     =     checkpoint["state_dict"]

		# Convert.
		print("Converting"  )
		__UpperCAmelCase	:							int     =     convert_megatron_checkpoint(_UpperCAmelCase,       _UpperCAmelCase  )

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

		__UpperCAmelCase	:							Union[str, Any]     =     "GPTSw3Tokenizer"

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

		# Store the state_dict to file.
		__UpperCAmelCase	:							Optional[int]     =     os.path.join(_UpperCAmelCase,       "pytorch_model.bin"  )
		print(F"Saving checkpoint to \"{output_checkpoint_file}\""  )
		torch.save(_UpperCAmelCase,       _UpperCAmelCase  )


if __name__ == "__main__":
			lowerCAmelCase__      :   List[Any]      	=			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")
			lowerCAmelCase__      :   Union[str, Any]      	=			parser.parse_args()
			main(_args)











'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCAmelCase__      :   Any      	=			{
    "configuration_open_llama": ["OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenLlamaConfig"],
}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__      :   Optional[int]      	=			["LlamaTokenizer"]

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__      :   int      	=			["LlamaTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__      :   int      	=			[
			    "OpenLlamaForCausalLM",
			    "OpenLlamaModel",
			    "OpenLlamaPreTrainedModel",
			    "OpenLlamaForSequenceClassification",
			]


if TYPE_CHECKING:
			from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from transformers import LlamaTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from transformers import LlamaTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_open_llama import (
						    OpenLlamaForCausalLM,
						    OpenLlamaForSequenceClassification,
						    OpenLlamaModel,
						    OpenLlamaPreTrainedModel,
						)


else:
			import sys

			lowerCAmelCase__      :   Dict      	=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




