'''simple docstring'''




# Copyright (c) 2021-, NVIDIA CORPORATION.  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.

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

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig






def __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__=0				)	->    str:
    # Format the message.
    if name is None:
        __lowerCamelCase		=  None
    else:
        __lowerCamelCase		=  '''.''' * max(0				,  spaces - 2				) + '''# {:''' + str(50 - spaces				) + '''s}'''
        __lowerCamelCase		=  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 __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				)	->    Union[str, Any]:
    # 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		=  param.size()
    if checkpoint_version == 1.0:
        # version 1.0 stores [num_heads * hidden_size * num_splits, :]
        __lowerCamelCase		=  (num_heads, hidden_size, num_splits) + input_shape[1:]
        __lowerCamelCase		=  param.view(*UpperCamelCase__				)
        __lowerCamelCase		=  param.transpose(0				,  2				)
        __lowerCamelCase		=  param.transpose(1				,  2				).contiguous()
    elif checkpoint_version >= 2.0:
        # other versions store [num_heads * num_splits * hidden_size, :]
        __lowerCamelCase		=  (num_heads, num_splits, hidden_size) + input_shape[1:]
        __lowerCamelCase		=  param.view(*UpperCamelCase__				)
        __lowerCamelCase		=  param.transpose(0				,  1				).contiguous()
    __lowerCamelCase		=  param.view(*UpperCamelCase__				)
    return param






def __lowerCAmelCase (							UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				)	->    Any:
    # The converted output model.
    __lowerCamelCase		=  {}

    # old versions did not store training args
    __lowerCamelCase		=  input_state_dict.get('''args'''				,  UpperCamelCase__				)
    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))

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

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

    # The model.
    __lowerCamelCase		=  input_state_dict['''model''']
    # The language model.
    __lowerCamelCase		=  model['''language_model''']
    # The embeddings.
    __lowerCamelCase		=  lm['''embedding''']

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

    # The position embeddings.
    __lowerCamelCase		=  embeddings['''position_embeddings''']['''weight''']
    # Read the causal mask dimension (seqlen). [max_sequence_length, hidden_size]
    __lowerCamelCase		=  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.
    __lowerCamelCase		=  pos_embeddings

    # The transformer.
    __lowerCamelCase		=  lm['''transformer'''] if '''transformer''' in lm.keys() else lm['''encoder''']

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

    # The simple map of names for "automated" rules.
    __lowerCamelCase		=  {
        '''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.
        __lowerCamelCase		=  layer_re.match(UpperCamelCase__				)

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

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

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

        # For layernorm(s), simply store the layer norm.
        if op_name.endswith('''layernorm'''				):
            __lowerCamelCase		=  '''ln_1''' if op_name.startswith('''input'''				) else '''ln_2'''
            __lowerCamelCase		=  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.
            __lowerCamelCase		=  torch.tril(torch.ones((n_positions, n_positions)				,  dtype=torch.floataa				)				).view(
                1				,  1				,  UpperCamelCase__				,  UpperCamelCase__				)
            __lowerCamelCase		=  causal_mask

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

            __lowerCamelCase		=  fix_query_key_value_ordering(UpperCamelCase__				,  UpperCamelCase__				,  3				,  UpperCamelCase__				,  UpperCamelCase__				)
            # Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
            __lowerCamelCase		=  out_val.transpose(0				,  1				).contiguous()
            # Store.
            __lowerCamelCase		=  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":
            __lowerCamelCase		=  fix_query_key_value_ordering(UpperCamelCase__				,  UpperCamelCase__				,  3				,  UpperCamelCase__				,  UpperCamelCase__				)
            # Store. No change of shape.
            __lowerCamelCase		=  out_val

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

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

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

    # The final layernorm.
    __lowerCamelCase		=  transformer['''final_layernorm.weight''']
    __lowerCamelCase		=  transformer['''final_layernorm.bias''']

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

    # It should be done!
    return output_state_dict






def __lowerCAmelCase (							)	->    Optional[int]:
    # Create the argument parser.
    __lowerCamelCase		=  argparse.ArgumentParser()
    parser.add_argument('''--print-checkpoint-structure'''				,  action='''store_true'''				)
    parser.add_argument(
        '''path_to_checkpoint'''				,  type=UpperCamelCase__				,  help='''Path to the checkpoint file (.zip archive or direct .pt file)'''				,  )
    parser.add_argument(
        '''--config_file'''				,  default=''''''				,  type=UpperCamelCase__				,  help='''An optional config json file describing the pre-trained model.'''				,  )
    __lowerCamelCase		=  parser.parse_args()

    # Extract the basename.
    __lowerCamelCase		=  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:
                __lowerCamelCase		=  torch.load(UpperCamelCase__				,  map_location='''cpu'''				)
    else:
        __lowerCamelCase		=  torch.load(args.path_to_checkpoint				,  map_location='''cpu'''				)

    __lowerCamelCase		=  input_state_dict.get('''args'''				,  UpperCamelCase__				)

    # 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:
                __lowerCamelCase		=  '''gelu_fast'''
            elif ds_args.openai_gelu:
                __lowerCamelCase		=  '''gelu_new'''
            else:
                __lowerCamelCase		=  '''gelu'''
        else:
            # in the very early days this used to be "gelu_new"
            __lowerCamelCase		=  '''gelu_new'''

        # Spell out all parameters in case the defaults change.
        __lowerCamelCase		=  GPTaConfig(
            vocab_size=5_02_57				,  n_positions=10_24				,  n_embd=10_24				,  n_layer=24				,  n_head=16				,  n_inner=40_96				,  activation_function=UpperCamelCase__				,  resid_pdrop=0.1				,  embd_pdrop=0.1				,  attn_pdrop=0.1				,  layer_norm_epsilon=1E-5				,  initializer_range=0.0_2				,  summary_type='''cls_index'''				,  summary_use_proj=UpperCamelCase__				,  summary_activation=UpperCamelCase__				,  summary_proj_to_labels=UpperCamelCase__				,  summary_first_dropout=0.1				,  scale_attn_weights=UpperCamelCase__				,  use_cache=UpperCamelCase__				,  bos_token_id=5_02_56				,  eos_token_id=5_02_56				,  )
    else:
        __lowerCamelCase		=  GPTaConfig.from_json_file(args.config_file				)

    __lowerCamelCase		=  ['''GPT2LMHeadModel''']

    # Convert.
    print('''Converting'''				)
    __lowerCamelCase		=  convert_megatron_checkpoint(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				)

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

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

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

    # Store the config to file.
    print('''Saving config'''				)
    config.save_pretrained(UpperCamelCase__				)

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

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


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

if __name__ == "__main__":
  main()

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





"""simple docstring"""







from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate







@dataclass(frozen=lowercase_  )
class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     field(default="""summarization""",					metadata={"""include_in_asdict_even_if_is_default""": True}  )
      _SCREAMING_SNAKE_CASE  				=     Features({"""text""": Value("""string"""  )}  )
      _SCREAMING_SNAKE_CASE  				=     Features({"""summary""": Value("""string"""  )}  )
      _SCREAMING_SNAKE_CASE  				=     "text"
      _SCREAMING_SNAKE_CASE  				=     "summary"



      @property
      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        return {self.text_column: "text", self.summary_column: "summary"}





