





import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig





def      _UpperCamelCase  (a__ :Optional[Any] ,      a__ :List[str] ,      a__ :Any=0							):






							"""simple docstring"""





							if name is None:
														UpperCamelCase__					  =      None
							else:
														UpperCamelCase__					  =      """.""" * max(0 ,      spaces - 2							) + """# {:""" + str(50 - spaces							) + """s}"""
														UpperCamelCase__					  =      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      _UpperCamelCase  (a__ :Optional[Any] ,      a__ :Union[str, Any] ,      a__ :Dict ,      a__ :Optional[int] ,      a__ :int							):






							"""simple docstring"""





							UpperCamelCase__					  =      param.size()
							if checkpoint_version == 1.0:
														# version 1.0 stores [num_heads * hidden_size * num_splits, :]
														UpperCamelCase__					  =      (num_heads, hidden_size, num_splits) + input_shape[1:]
														UpperCamelCase__					  =      param.view(*a__							)
														UpperCamelCase__					  =      param.transpose(0 ,      2							)
														UpperCamelCase__					  =      param.transpose(1 ,      2							).contiguous()
							elif checkpoint_version >= 2.0:
														# other versions store [num_heads * num_splits * hidden_size, :]
														UpperCamelCase__					  =      (num_heads, num_splits, hidden_size) + input_shape[1:]
														UpperCamelCase__					  =      param.view(*a__							)
														UpperCamelCase__					  =      param.transpose(0 ,      1							).contiguous()
							UpperCamelCase__					  =      param.view(*a__							)
							return param





def      _UpperCamelCase  (a__ :List[str] ,      a__ :Optional[int] ,      a__ :List[str]							):






							"""simple docstring"""





							UpperCamelCase__					  =      {}

							# old versions did not store training args
							UpperCamelCase__					  =      input_state_dict.get("""args""" ,      a__							)
							if ds_args is not None:
														# do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
														# from pprint import pprint
														# pprint(vars(ds_args))

														UpperCamelCase__					  =      ds_args.padded_vocab_size
														UpperCamelCase__					  =      ds_args.max_position_embeddings
														UpperCamelCase__					  =      ds_args.hidden_size
														UpperCamelCase__					  =      ds_args.num_layers
														UpperCamelCase__					  =      ds_args.num_attention_heads
														UpperCamelCase__					  =      ds_args.ffn_hidden_size
														# pprint(config)

							# The number of heads.
							UpperCamelCase__					  =      config.n_head
							# The hidden_size per head.
							UpperCamelCase__					  =      config.n_embd // config.n_head
							# Megatron-LM checkpoint version
							if "checkpoint_version" in input_state_dict.keys():
														UpperCamelCase__					  =      input_state_dict["""checkpoint_version"""]
							else:
														UpperCamelCase__					  =      0.0

							# The model.
							UpperCamelCase__					  =      input_state_dict["""model"""]
							# The language model.
							UpperCamelCase__					  =      model["""language_model"""]
							# The embeddings.
							UpperCamelCase__					  =      lm["""embedding"""]

							# The word embeddings.
							UpperCamelCase__					  =      embeddings["""word_embeddings"""]["""weight"""]
							# Truncate the embedding table to vocab_size rows.
							UpperCamelCase__					  =      word_embeddings[: config.vocab_size, :]
							UpperCamelCase__					  =      word_embeddings

							# The position embeddings.
							UpperCamelCase__					  =      embeddings["""position_embeddings"""]["""weight"""]
							# Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
							UpperCamelCase__					  =      pos_embeddings.size(0							)
							if n_positions != config.n_positions:
														raise ValueError(
														    f"""pos_embeddings.max_sequence_length={n_positions} and config.n_positions={config.n_positions} don\'t match"""							)
							# Store the position embeddings.
							UpperCamelCase__					  =      pos_embeddings

							# The transformer.
							UpperCamelCase__					  =      lm["""transformer"""] if """transformer""" in lm.keys() else lm["""encoder"""]

							# The regex to extract layer names.
							UpperCamelCase__					  =      re.compile(r"""layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"""							)

							# The simple map of names for "automated" rules.
							UpperCamelCase__					  =      {
							    """attention.dense""": """.attn.c_proj.""",
							    """self_attention.dense""": """.attn.c_proj.""",
							    """mlp.dense_h_to_4h""": """.mlp.c_fc.""",
							    """mlp.dense_4h_to_h""": """.mlp.c_proj.""",
							}

							# Extract the layers.
							for key, val in transformer.items():
														# Match the name.
														UpperCamelCase__					  =      layer_re.match(a__							)

														# Stop if that's not a layer
														if m is None:
																					break

														# The index of the layer.
														UpperCamelCase__					  =      int(m.group(1							)							)
														# The name of the operation.
														UpperCamelCase__					  =      m.group(2							)
														# Is it a weight or a bias?
														UpperCamelCase__					  =      m.group(3							)

														# The name of the layer.
														UpperCamelCase__					  =      f"""transformer.h.{layer_idx}"""

														# For layernorm(s), simply store the layer norm.
														if op_name.endswith("""layernorm"""							):
																					UpperCamelCase__					  =      """ln_1""" if op_name.startswith("""input"""							) else """ln_2"""
																					UpperCamelCase__					  =      val

														# Transpose the QKV matrix.
														elif (
														    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
														) and weight_or_bias == "weight":
																					# Insert a tensor of 1x1xDxD bias.
																					UpperCamelCase__					  =      torch.tril(torch.ones((n_positions, n_positions) ,      dtype=torch.floataa							)							).view(
																					    1 ,      1 ,      a__ ,      a__							)
																					UpperCamelCase__					  =      causal_mask

																					# Insert a "dummy" tensor for masked_bias.
																					UpperCamelCase__					  =      torch.tensor(-1e4 ,      dtype=torch.floataa							)
																					UpperCamelCase__					  =      masked_bias

																					UpperCamelCase__					  =      fix_query_key_value_ordering(a__ ,      a__ ,      3 ,      a__ ,      a__							)
																					# Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
																					UpperCamelCase__					  =      out_val.transpose(0 ,      1							).contiguous()
																					# Store.
																					UpperCamelCase__					  =      out_val

														# Transpose the bias.
														elif (
														    op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
														) and weight_or_bias == "bias":
																					UpperCamelCase__					  =      fix_query_key_value_ordering(a__ ,      a__ ,      3 ,      a__ ,      a__							)
																					# Store. No change of shape.
																					UpperCamelCase__					  =      out_val

														# Transpose the weights.
														elif weight_or_bias == "weight":
																					UpperCamelCase__					  =      megatron_to_transformers[op_name]
																					UpperCamelCase__					  =      val.transpose(0 ,      1							)

														# Copy the bias.
														elif weight_or_bias == "bias":
																					UpperCamelCase__					  =      megatron_to_transformers[op_name]
																					UpperCamelCase__					  =      val

    # DEBUG.
							assert config.n_layer == layer_idx + 1

							# The final layernorm.
							UpperCamelCase__					  =      transformer["""final_layernorm.weight"""]
							UpperCamelCase__					  =      transformer["""final_layernorm.bias"""]

							# For LM head, transformers' wants the matrix to weight embeddings.
							UpperCamelCase__					  =      word_embeddings

							# It should be done!
							return output_state_dict





def      _UpperCamelCase  ():






							"""simple docstring"""





							UpperCamelCase__					  =      argparse.ArgumentParser()
							parser.add_argument("""--print-checkpoint-structure""" ,      action="""store_true"""							)
							parser.add_argument(
							    """path_to_checkpoint""" ,      type=a__ ,      help="""Path to the checkpoint file (.zip archive or direct .pt file)""" ,      )
							parser.add_argument(
							    """--config_file""" ,      default="""""" ,      type=a__ ,      help="""An optional config json file describing the pre-trained model.""" ,      )
							UpperCamelCase__					  =      parser.parse_args()

							# Extract the basename.
							UpperCamelCase__					  =      os.path.dirname(args.path_to_checkpoint							)

							# Load the model.
							# the .zip is very optional, let's keep it for backward compatibility
							print(f"""Extracting PyTorch state dictionary from {args.path_to_checkpoint}"""							)
							if args.path_to_checkpoint.endswith(""".zip"""							):
														with zipfile.ZipFile(args.path_to_checkpoint ,      """r"""							) as checkpoint:
																					with checkpoint.open("""release/mp_rank_00/model_optim_rng.pt"""							) as pytorch_dict:
																												UpperCamelCase__					  =      torch.load(a__ ,      map_location="""cpu"""							)
							else:
														UpperCamelCase__					  =      torch.load(args.path_to_checkpoint ,      map_location="""cpu"""							)

							UpperCamelCase__					  =      input_state_dict.get("""args""" ,      a__							)

							# Read the config, or default to the model released by NVIDIA.
							if args.config_file == "":
														if ds_args is not None:
																					if ds_args.bias_gelu_fusion:
																												UpperCamelCase__					  =      """gelu_fast"""
																					elif ds_args.openai_gelu:
																												UpperCamelCase__					  =      """gelu_new"""
																					else:
																												UpperCamelCase__					  =      """gelu"""
														else:
																					# in the very early days this used to be "gelu_new"
																					UpperCamelCase__					  =      """gelu_new"""

														# Spell out all parameters in case the defaults change.
														UpperCamelCase__					  =      GPTaConfig(
														    vocab_size=5_0257 ,      n_positions=1024 ,      n_embd=1024 ,      n_layer=24 ,      n_head=16 ,      n_inner=4096 ,      activation_function=a__ ,      resid_pdrop=0.1 ,      embd_pdrop=0.1 ,      attn_pdrop=0.1 ,      layer_norm_epsilon=1e-5 ,      initializer_range=0.02 ,      summary_type="""cls_index""" ,      summary_use_proj=a__ ,      summary_activation=a__ ,      summary_proj_to_labels=a__ ,      summary_first_dropout=0.1 ,      scale_attn_weights=a__ ,      use_cache=a__ ,      bos_token_id=5_0256 ,      eos_token_id=5_0256 ,      )
							else:
														UpperCamelCase__					  =      GPTaConfig.from_json_file(args.config_file							)

							UpperCamelCase__					  =      ["""GPT2LMHeadModel"""]

							# Convert.
							print("""Converting"""							)
							UpperCamelCase__					  =      convert_megatron_checkpoint(a__ ,      a__ ,      a__							)

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

							# Add tokenizer class info to config
							# see https://github.com/huggingface/transformers/issues/13906)
							if ds_args is not None:
														UpperCamelCase__					  =      ds_args.tokenizer_type
														if tokenizer_type == "GPT2BPETokenizer":
																					UpperCamelCase__					  =      """gpt2"""
														elif tokenizer_type == "PretrainedFromHF":
																					UpperCamelCase__					  =      ds_args.tokenizer_name_or_path
														else:
																					raise ValueError(f"""Unrecognized tokenizer_type {tokenizer_type}"""							)
							else:
														UpperCamelCase__					  =      """gpt2"""

							UpperCamelCase__					  =      AutoTokenizer.from_pretrained(a__							)
							UpperCamelCase__					  =      type(a__							).__name__
							UpperCamelCase__					  =      tokenizer_class

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

							# Save tokenizer based on args
							print(f"""Adding {tokenizer_class} tokenizer files"""							)
							tokenizer.save_pretrained(a__							)

							# Store the state_dict to file.
							UpperCamelCase__					  =      os.path.join(a__ ,      """pytorch_model.bin"""							)
							print(f"""Saving checkpoint to \"{output_checkpoint_file}\""""							)
							torch.save(a__ ,      a__							)


####################################################################################################

if __name__ == "__main__":
					main()

####################################################################################################











"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE__   =  {
    "configuration_roberta_prelayernorm": [
        "ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "RobertaPreLayerNormConfig",
        "RobertaPreLayerNormOnnxConfig",
    ],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "RobertaPreLayerNormForCausalLM",
           "RobertaPreLayerNormForMaskedLM",
           "RobertaPreLayerNormForMultipleChoice",
           "RobertaPreLayerNormForQuestionAnswering",
           "RobertaPreLayerNormForSequenceClassification",
           "RobertaPreLayerNormForTokenClassification",
           "RobertaPreLayerNormModel",
           "RobertaPreLayerNormPreTrainedModel",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFRobertaPreLayerNormForCausalLM",
           "TFRobertaPreLayerNormForMaskedLM",
           "TFRobertaPreLayerNormForMultipleChoice",
           "TFRobertaPreLayerNormForQuestionAnswering",
           "TFRobertaPreLayerNormForSequenceClassification",
           "TFRobertaPreLayerNormForTokenClassification",
           "TFRobertaPreLayerNormMainLayer",
           "TFRobertaPreLayerNormModel",
           "TFRobertaPreLayerNormPreTrainedModel",
       ]

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__   =  [
           "FlaxRobertaPreLayerNormForCausalLM",
           "FlaxRobertaPreLayerNormForMaskedLM",
           "FlaxRobertaPreLayerNormForMultipleChoice",
           "FlaxRobertaPreLayerNormForQuestionAnswering",
           "FlaxRobertaPreLayerNormForSequenceClassification",
           "FlaxRobertaPreLayerNormForTokenClassification",
           "FlaxRobertaPreLayerNormModel",
           "FlaxRobertaPreLayerNormPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_roberta_prelayernorm import (
           ROBERTA_PRELAYERNORM_PRETRAINED_CONFIG_ARCHIVE_MAP,
           RobertaPreLayerNormConfig,
           RobertaPreLayerNormOnnxConfig,
       )

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_roberta_prelayernorm import (
                  ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  RobertaPreLayerNormForCausalLM,
                  RobertaPreLayerNormForMaskedLM,
                  RobertaPreLayerNormForMultipleChoice,
                  RobertaPreLayerNormForQuestionAnswering,
                  RobertaPreLayerNormForSequenceClassification,
                  RobertaPreLayerNormForTokenClassification,
                  RobertaPreLayerNormModel,
                  RobertaPreLayerNormPreTrainedModel,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_roberta_prelayernorm import (
                  TF_ROBERTA_PRELAYERNORM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFRobertaPreLayerNormForCausalLM,
                  TFRobertaPreLayerNormForMaskedLM,
                  TFRobertaPreLayerNormForMultipleChoice,
                  TFRobertaPreLayerNormForQuestionAnswering,
                  TFRobertaPreLayerNormForSequenceClassification,
                  TFRobertaPreLayerNormForTokenClassification,
                  TFRobertaPreLayerNormMainLayer,
                  TFRobertaPreLayerNormModel,
                  TFRobertaPreLayerNormPreTrainedModel,
              )

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_roberta_prelayernorm import (
                  FlaxRobertaPreLayerNormForCausalLM,
                  FlaxRobertaPreLayerNormForMaskedLM,
                  FlaxRobertaPreLayerNormForMultipleChoice,
                  FlaxRobertaPreLayerNormForQuestionAnswering,
                  FlaxRobertaPreLayerNormForSequenceClassification,
                  FlaxRobertaPreLayerNormForTokenClassification,
                  FlaxRobertaPreLayerNormModel,
                  FlaxRobertaPreLayerNormPreTrainedModel,
              )

else:
       import sys

       SCREAMING_SNAKE_CASE__   =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






