'''simple docstring'''

import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig


def        UpperCAmelCase__   (				_lowerCAmelCase    :			Dict,      _lowerCAmelCase    :			Optional[Any],      _lowerCAmelCase    :			Optional[Any]=0)    ->		Tuple:



						"""simple docstring"""

						if name is None:
												snake_case :       Dict				=    None
						else:
												snake_case :       Any				=    "." * max(0,      spaces - 2) + "# {:" + str(50 - spaces) + "s}"
												snake_case :       int				=    fmt.format(_lowerCAmelCase)

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


def        UpperCAmelCase__   (				_lowerCAmelCase    :			str,      _lowerCAmelCase    :			List[Any],      _lowerCAmelCase    :			Dict,      _lowerCAmelCase    :			Union[str, Any])    ->		int:



						"""simple docstring"""

						snake_case :       Union[str, Any]				=    param.size()
						# other versions store [num_heads * num_splits * hidden_size, :]
						snake_case :       int				=    (num_heads, num_splits, hidden_size) + input_shape[1:]
						snake_case :       Optional[int]				=    param.view(*_lowerCAmelCase)
						snake_case :       Union[str, Any]				=    param.transpose(0,      1).contiguous()
						snake_case :       Optional[Any]				=    param.view(*_lowerCAmelCase)
						return param


def        UpperCAmelCase__   (				_lowerCAmelCase    :			Dict,      _lowerCAmelCase    :			List[str])    ->		Dict:



						"""simple docstring"""

						snake_case :       Dict				=    config.n_positions
						snake_case :       Tuple				=    config.n_layer
						snake_case :       Optional[int]				=    config.vocab_size
						snake_case :       Tuple				=    config.n_head
						snake_case :       List[str]				=    config.n_embd // config.n_head

						snake_case :       Union[str, Any]				=    sd_megatron["model.language_model.embedding.word_embeddings.weight"][:vocab_size, :]
						snake_case :       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"],
						}

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

												snake_case :       Union[str, Any]				=    sd_megatron[f"""{pf}{i}.input_layernorm.weight"""]
												snake_case :       Dict				=    sd_megatron[f"""{pf}{i}.input_layernorm.bias"""]

												snake_case :       Any				=    sd_megatron[f"""{pf}{i}.self_attention.query_key_value.weight"""]
												snake_case :       int				=    fix_query_key_value_ordering(_lowerCAmelCase,      3,      _lowerCAmelCase,      _lowerCAmelCase)
												snake_case :       Any				=    vala.transpose(0,      1).contiguous()

												snake_case :       Union[str, Any]				=    sd_megatron[f"""{pf}{i}.self_attention.query_key_value.bias"""]
												snake_case :       Any				=    fix_query_key_value_ordering(_lowerCAmelCase,      3,      _lowerCAmelCase,      _lowerCAmelCase)
												snake_case :       List[str]				=    vala

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

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

						return sd_hf


def        UpperCAmelCase__   (				_lowerCAmelCase    :			List[Any],      _lowerCAmelCase    :			List[Any])    ->		Optional[int]:



						"""simple docstring"""

						snake_case :       Union[str, Any]				=    6_4000
						snake_case :       Union[str, Any]				=    config_megatron["encoder_seq_length"]
						snake_case :       int				=    config_megatron["hidden_size"]
						snake_case :       Optional[Any]				=    config_megatron["num_layers"]
						snake_case :       Dict				=    config_megatron["num_attention_heads"]
						snake_case :       Optional[int]				=    config_megatron["ffn_hidden_size"]
						snake_case :       str				=    "gelu"
						snake_case :       List[Any]				=    0.1
						snake_case :       Any				=    0.1
						snake_case :       List[Any]				=    0.1
						snake_case :       Tuple				=    config_megatron["layernorm_epsilon"]  # 1e-5
						snake_case :       List[Any]				=    config_megatron["init_method_std"]  # 0.02
						snake_case :       Optional[Any]				=    config_megatron["apply_query_key_layer_scaling"]  # True
						snake_case :       str				=    True
						snake_case :       str				=    True

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

						return config_hf


def        UpperCAmelCase__   (				_lowerCAmelCase    :			Any)    ->		List[str]:



						"""simple docstring"""

						print(_lowerCAmelCase)

						snake_case :       Dict				=    args.checkpoint_path
						snake_case :       Any				=    args.save_path
						if isfile(_lowerCAmelCase):
												raise FileNotFoundError(f"""ERROR! could not find file {checkpoint_path}""")

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

						# Load the config.
						snake_case :       List[Any]				=    checkpoint["hyper_parameters"]["cfg"]
						snake_case :       Dict				=    GPTaConfig()
						snake_case :       int				=    copy_config(config_hf=_lowerCAmelCase,      config_megatron=_lowerCAmelCase)
						snake_case :       Optional[Any]				=    ["GPT2LMHeadModel"]

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

						# Convert.
						print("Converting")
						snake_case :       Tuple				=    convert_megatron_checkpoint(_lowerCAmelCase,      _lowerCAmelCase)

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

						snake_case :       List[str]				=    "GPTSw3Tokenizer"

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

						# Store the state_dict to file.
						snake_case :       Union[str, Any]				=    os.path.join(_lowerCAmelCase,      "pytorch_model.bin")
						print(f"""Saving checkpoint to \"{output_checkpoint_file}\"""")
						torch.save(_lowerCAmelCase,      _lowerCAmelCase)


if __name__ == "__main__":
						A_	   = 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")
						A_	   = parser.parse_args()
						main(_args)






'''simple docstring'''

# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


A_	   = {
    "configuration_convnextv2": [
        "CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "ConvNextV2Config",
    ]
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A_	   = [
						    "CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "ConvNextV2ForImageClassification",
						    "ConvNextV2Model",
						    "ConvNextV2PreTrainedModel",
						    "ConvNextV2Backbone",
						]


if TYPE_CHECKING:
						from .configuration_convnextva import (
						    CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    ConvNextVaConfig,
						)

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_convnextva import (
												    CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST,
												    ConvNextVaBackbone,
												    ConvNextVaForImageClassification,
												    ConvNextVaModel,
												    ConvNextVaPreTrainedModel,
												)

else:
						import sys

						A_	   = _LazyModule(__name__, globals()["__file__"], _import_structure)





