code: 
  
  
  
  
  
  import argparse
  import os
  import re
  import zipfile
  
  import torch
  
  from transformers import MegatronBertConfig
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Tuple   ,     lowercase       :    str   ,     lowercase       :    Dict=0				) ->    Dict:
  
  
    if name is None:
      __snake_case   :				Any        =       None
    else:
      __snake_case   :				List[str]        =       '''.''' * max(0   ,     spaces - 2				) + '''# {:''' + str(50 - spaces				) + '''s}'''
      __snake_case   :				Dict        =       fmt.format(SCREAMING_SNAKE_CASE__				)
  
    # Print and recurse (if needed).
    if isinstance(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				):
      if msg is not None:
        print(SCREAMING_SNAKE_CASE__				)
      for k in val.keys():
        recursive_print(SCREAMING_SNAKE_CASE__   ,     val[k]   ,     spaces + 2				)
    elif isinstance(SCREAMING_SNAKE_CASE__   ,     torch.Tensor				):
      print(SCREAMING_SNAKE_CASE__   ,     ":"   ,     val.size()				)
    else:
      print(SCREAMING_SNAKE_CASE__   ,     ":"   ,     SCREAMING_SNAKE_CASE__				)
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    int   ,     lowercase       :    Optional[int]   ,     lowercase       :    Union[str, Any]   ,     lowercase       :    Tuple   ,     lowercase       :    Any				) ->    Dict:
    __snake_case   :				List[Any]        =       param.size()
    if checkpoint_version == 1.0:
      # version 1.0 stores [num_heads * hidden_size * num_splits, :]
      __snake_case   :				int        =       (num_heads, hidden_size, num_splits) + input_shape[1:]
      __snake_case   :				Union[str, Any]        =       param.view(*SCREAMING_SNAKE_CASE__				)
      __snake_case   :				List[Any]        =       param.transpose(0   ,     2				)
      __snake_case   :				Optional[int]        =       param.transpose(1   ,     2				).contiguous()
    elif checkpoint_version >= 2.0:
      # other versions store [num_heads * num_splits * hidden_size, :]
      __snake_case   :				Any        =       (num_heads, num_splits, hidden_size) + input_shape[1:]
      __snake_case   :				Tuple        =       param.view(*SCREAMING_SNAKE_CASE__				)
      __snake_case   :				Any        =       param.transpose(0   ,     1				).contiguous()
    __snake_case   :				Any        =       param.view(*SCREAMING_SNAKE_CASE__				)
    return param
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Optional[int]   ,     lowercase       :    Any   ,     lowercase       :    Tuple				) ->    Tuple:
    __snake_case   :				Optional[Any]        =       {}
  
    # old versions did not store training args
    __snake_case   :				List[Any]        =       input_state_dict.get("args"   ,     SCREAMING_SNAKE_CASE__				)
    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))
  
      __snake_case   :				Any        =       ds_args.tokenizer_type
      __snake_case   :				Optional[int]        =       ds_args.padded_vocab_size
      __snake_case   :				Tuple        =       ds_args.max_position_embeddings
      __snake_case   :				str        =       ds_args.hidden_size
      __snake_case   :				Any        =       ds_args.num_layers
      __snake_case   :				List[Any]        =       ds_args.num_attention_heads
      __snake_case   :				str        =       ds_args.ffn_hidden_size if '''ffn_hidden_size''' in ds_args else 4 * ds_args.hidden_size
      # pprint(config)
  
    # The number of heads.
    __snake_case   :				Optional[int]        =       config.num_attention_heads
    # The hidden_size per head.
    __snake_case   :				str        =       config.hidden_size // heads
    # Megatron-LM checkpoint version
    if "checkpoint_version" in input_state_dict.keys():
      __snake_case   :				Optional[Any]        =       input_state_dict['''checkpoint_version''']
    else:
      __snake_case   :				str        =       0.0
  
    # The model.
    __snake_case   :				Dict        =       input_state_dict['''model''']
    # The language model.
    __snake_case   :				Optional[Any]        =       model['''language_model''']
    # The embeddings.
    __snake_case   :				Tuple        =       lm['''embedding''']
  
    # The word embeddings.
    __snake_case   :				Optional[Any]        =       embeddings['''word_embeddings''']['''weight''']
    # Truncate the embedding table to vocab_size rows.
    __snake_case   :				int        =       word_embeddings[: config.vocab_size, :]
    # Store the word embeddings.
    __snake_case   :				Optional[Any]        =       word_embeddings
  
    # The position embeddings.
    __snake_case   :				Union[str, Any]        =       embeddings['''position_embeddings''']['''weight''']
    assert pos_embeddings.size(0				) == config.max_position_embeddings and pos_embeddings.size(1				) == config.hidden_size
    # Store the position embeddings.
    __snake_case   :				List[str]        =       pos_embeddings
  
    # The token-type embeddings.
    __snake_case   :				Optional[int]        =       embeddings['''tokentype_embeddings''']['''weight''']
    # Store the position embeddings.
    __snake_case   :				Union[str, Any]        =       tokentype_embeddings
  
    # The transformer.
    __snake_case   :				int        =       lm['''transformer'''] if '''transformer''' in lm.keys() else lm['''encoder''']
  
    # The regex to extract layer names.
    __snake_case   :				Optional[int]        =       re.compile(R"layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)"				)
  
    # The simple map of names for "automated" rules.
    __snake_case   :				Tuple        =       {
        '''attention.dense''': '''.attention.output.dense.''',
        '''self_attention.dense''': '''.attention.output.dense.''',
        '''mlp.dense_h_to_4h''': '''.intermediate.dense.''',
        '''mlp.dense_4h_to_h''': '''.output.dense.''',
    }
  
    # Keep track of the attention/query/value tensor.
    __snake_case   :				Optional[Any]        =       None
  
    # Extract the layers.
    for key, val in transformer.items():
      # Match the name.
      __snake_case   :				str        =       layer_re.match(SCREAMING_SNAKE_CASE__				)
  
      # Stop if that's not a layer
      if m is None:
        break
  
      # The index of the layer.
      __snake_case   :				int        =       int(m.group(1				)				)
      # The name of the operation.
      __snake_case   :				Dict        =       m.group(2				)
      # Is it a weight or a bias?
      __snake_case   :				Any        =       m.group(3				)
  
      # The name of the layer.
      __snake_case   :				Dict        =       f"""bert.encoder.layer.{layer_idx}"""
  
      # For layernorm(s), simply store the layer norm.
      if op_name.endswith("layernorm"				):
        __snake_case   :				Optional[int]        =       '''attention.ln''' if op_name.startswith("input"				) else '''ln'''
        __snake_case   :				int        =       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":
        # Make sure the QKV pointer is nil.
        assert attention_qkv_weight is None, ""
  
        __snake_case   :				Union[str, Any]        =       fix_query_key_value_ordering(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__   ,     3   ,     SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				)
        # Store the tensor as we need the bias as well to interleave QKV and biases.
        __snake_case   :				Tuple        =       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":
        # Make sure we read the weight tensor.
        assert attention_qkv_weight is not None, ""
  
        # Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved.
        __snake_case   :				Tuple        =       attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
        __snake_case   :				str        =       attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
        __snake_case   :				Union[str, Any]        =       attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]
  
        __snake_case   :				Union[str, Any]        =       fix_query_key_value_ordering(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__   ,     3   ,     SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				)
        # Split the bias.
        __snake_case   :				Optional[int]        =       out_val[0 * config.hidden_size : 1 * config.hidden_size]
        __snake_case   :				List[Any]        =       out_val[1 * config.hidden_size : 2 * config.hidden_size]
        __snake_case   :				str        =       out_val[2 * config.hidden_size : 3 * config.hidden_size]
  
        # Store.
        __snake_case   :				Union[str, Any]        =       q
        __snake_case   :				int        =       q_bias
        __snake_case   :				Union[str, Any]        =       k
        __snake_case   :				Optional[Any]        =       k_bias
        __snake_case   :				Tuple        =       v
        __snake_case   :				Tuple        =       v_bias
  
        # Clear the stored tensor.
        __snake_case   :				Optional[Any]        =       None
  
      # Copy weights and biases as is.
      elif weight_or_bias in ["weight", "bias"]:
        __snake_case   :				Union[str, Any]        =       megatron_to_transformers[op_name]
        __snake_case   :				Optional[Any]        =       val
  
      # The final layernorm.
    __snake_case   :				Optional[int]        =       transformer['''final_layernorm.weight''']
    __snake_case   :				int        =       transformer['''final_layernorm.bias''']
  
    # The pooler.
    __snake_case   :				List[str]        =       lm['''pooler''']
  
    # Store the matrix and the bias.
    __snake_case   :				int        =       pooler['''dense.weight''']
    __snake_case   :				Tuple        =       pooler['''dense.bias''']
  
    # The LM head from Megatron (for RACE).
    __snake_case   :				Optional[Any]        =       model['''lm_head''']
  
    # The transform matrix.
    __snake_case   :				Optional[int]        =       lm_head['''dense.weight''']
    __snake_case   :				Dict        =       lm_head['''dense.bias''']
  
    # The transform LN.
    __snake_case   :				Optional[int]        =       lm_head['''layernorm.weight''']
    __snake_case   :				int        =       lm_head['''layernorm.bias''']
  
    # For the decoder, we replicate the weights.
    __snake_case   :				Optional[int]        =       word_embeddings
    __snake_case   :				List[str]        =       lm_head['''bias''']
  
    # The classifier from Megatron (for MLNI).
    __snake_case   :				int        =       model['''binary_head''']
  
    # Store the classifier.
    __snake_case   :				List[Any]        =       binary_head['''weight''']
    __snake_case   :				Tuple        =       binary_head['''bias''']
  
    # It should be done!
    return output_state_dict
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	) ->    int:
    __snake_case   :				Optional[Any]        =       argparse.ArgumentParser()
    parser.add_argument("--print-checkpoint-structure"   ,     action="store_true"				)
    parser.add_argument("path_to_checkpoint"   ,     type=SCREAMING_SNAKE_CASE__   ,     help="Path to the ZIP file containing the checkpoint"				)
    parser.add_argument(
        "--config_file"   ,     default=""   ,     type=SCREAMING_SNAKE_CASE__   ,     help="An optional config json file describing the pre-trained model."   ,     )
    __snake_case   :				Tuple        =       parser.parse_args()
  
    # Extract the basename.
    __snake_case   :				Dict        =       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:
          __snake_case   :				List[str]        =       torch.load(SCREAMING_SNAKE_CASE__   ,     map_location="cpu"				)
    else:
      __snake_case   :				str        =       torch.load(args.path_to_checkpoint   ,     map_location="cpu"				)
  
    if args.config_file == "":
      # Default config of megatron-bert 345m
      __snake_case   :				List[str]        =       MegatronBertConfig()
  
      # different megatron-bert-*-345m models have different vocab sizes, so override the default
      # config (which is for megatron-bert-cased-345m) with the actual vocab dimension
      __snake_case   :				Optional[int]        =       input_state_dict['''model''']['''lm_head''']['''bias'''].numel()
    else:
      __snake_case   :				List[str]        =       MegatronBertConfig.from_json_file(args.config_file				)
  
    # Convert.
    print("Converting"				)
    __snake_case   :				List[Any]        =       convert_megatron_checkpoint(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				)
  
    # Print the structure of converted state dict.
    if args.print_checkpoint_structure:
      recursive_print(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				)
  
    # Store the config to file.
    print("Saving config"				)
    config.save_pretrained(SCREAMING_SNAKE_CASE__				)
  
    # Store the state_dict to file.
    __snake_case   :				Optional[int]        =       os.path.join(SCREAMING_SNAKE_CASE__   ,     "pytorch_model.bin"				)
    print(f"""Saving checkpoint to \"{output_checkpoint_file}\""""				)
    torch.save(SCREAMING_SNAKE_CASE__   ,     SCREAMING_SNAKE_CASE__				)
  
  
  ####################################################################################################
  
  if __name__ == "__main__":
        main()
  
  ####################################################################################################
  
  
style_context: 
  
  
  
  
  
  import subprocess
  from argparse import ArgumentParser
  from typing import List, Union
  
  from huggingface_hub.hf_api import HfFolder, create_repo, whoami
  from requests.exceptions import HTTPError
  
  from . import BaseTransformersCLICommand
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        @staticmethod
        def 						UpperCAmelCase		( UpperCAmelCase   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Dict        =       parser.add_parser("login"				,						help="Log in using the same credentials as on huggingface.co"   )
          login_parser.set_defaults(func=lambda UpperCAmelCase   : LoginCommand(UpperCAmelCase   )   )
          __snake_case   :				int        =       parser.add_parser("whoami"				,						help="Find out which huggingface.co account you are logged in as."   )
          whoami_parser.set_defaults(func=lambda UpperCAmelCase   : WhoamiCommand(UpperCAmelCase   )   )
          __snake_case   :				int        =       parser.add_parser("logout"				,						help="Log out"   )
          logout_parser.set_defaults(func=lambda UpperCAmelCase   : LogoutCommand(UpperCAmelCase   )   )
  
          # new system: git-based repo system
          __snake_case   :				Union[str, Any]        =       parser.add_parser(
              "repo"				,						help="Deprecated: use `huggingface-cli` instead. Commands to interact with your huggingface.co repos."				,						)
          __snake_case   :				List[str]        =       repo_parser.add_subparsers(
              help="Deprecated: use `huggingface-cli` instead. huggingface.co repos related commands"   )
          __snake_case   :				int        =       repo_subparsers.add_parser(
              "create"				,						help="Deprecated: use `huggingface-cli` instead. Create a new repo on huggingface.co"   )
          repo_create_parser.add_argument(
              "name"				,						type=UpperCAmelCase				,						help="Name for your model's repo. Will be namespaced under your username to build the model id."				,						)
          repo_create_parser.add_argument("--organization"				,						type=UpperCAmelCase				,						help="Optional: organization namespace."   )
          repo_create_parser.add_argument("-y"				,						"--yes"				,						action="store_true"				,						help="Optional: answer Yes to the prompt"   )
          repo_create_parser.set_defaults(func=lambda UpperCAmelCase   : RepoCreateCommand(UpperCAmelCase   )   )
  
  
  
  
  class 				_lowerCamelCase :
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        UpperCAmelCase_						:      Tuple       ="\u001b[1m"
        UpperCAmelCase_						:      int       ="\u001b[31m"
        UpperCAmelCase_						:      Optional[Any]       ="\u001b[90m"
        UpperCAmelCase_						:      Optional[Any]       ="\u001b[0m"
  
  
  
  
  
  
  
        @classmethod
        def 						UpperCAmelCase		( cls				,						UpperCAmelCase   )       ->   Union[str, Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return F"""{cls._bold}{s}{cls._reset}"""
  
  
  
  
  
  
  
        @classmethod
        def 						UpperCAmelCase		( cls				,						UpperCAmelCase   )       ->   Optional[int]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return F"""{cls._bold}{cls._red}{s}{cls._reset}"""
  
  
  
  
  
  
  
        @classmethod
        def 						UpperCAmelCase		( cls				,						UpperCAmelCase   )       ->   List[str]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          return F"""{cls._gray}{s}{cls._reset}"""
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    List[List[Union[str, int]]]   ,     lowercase       :    List[str]				) ->    str:
    __snake_case   :				Dict        =       [max(len(str(lowercase				)				) for x in col				) for col in zip(*lowercase   ,     lowercase				)]
    __snake_case   :				str        =       ("{{:{}}} " * len(lowercase				)).format(*lowercase				)
    __snake_case   :				Tuple        =       []
    lines.append(row_format.format(*lowercase				)				)
    lines.append(row_format.format(*["-" * w for w in col_widths]				)				)
    for row in rows:
      lines.append(row_format.format(*lowercase				)				)
    return "\n".join(lowercase				)
  
  
  
  
  class 				_lowerCamelCase :
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def __init__( self				,						UpperCAmelCase   )       ->   Dict:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				List[str]        =       args
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Union[str, Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          print(
              ANSI.red(
                  "ERROR! `huggingface-cli login` uses an outdated login mechanism "
                  "that is not compatible with the Hugging Face Hub backend anymore. "
                  "Please use `huggingface-cli login instead."   )   )
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   str:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          print(
              ANSI.red(
                  "WARNING! `transformers-cli whoami` is deprecated and will be removed in v5. Please use "
                  "`huggingface-cli whoami` instead."   )   )
          __snake_case   :				Optional[Any]        =       HfFolder.get_token()
          if token is None:
            print("Not logged in"   )
            exit()
          try:
            __snake_case    ,   __snake_case   :				Tuple        =       whoami(UpperCAmelCase   )
            print(UpperCAmelCase   )
            if orgs:
              print(ANSI.bold("orgs: "   )				,						",".join(UpperCAmelCase   )   )
          except HTTPError as e:
            print(UpperCAmelCase   )
            print(ANSI.red(e.response.text   )   )
            exit(1   )
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   List[str]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          print(
              ANSI.red(
                  "ERROR! `transformers-cli logout` uses an outdated logout mechanism "
                  "that is not compatible with the Hugging Face Hub backend anymore. "
                  "Please use `huggingface-cli logout instead."   )   )
  
  
  
  
  class 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Tuple:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          print(
              ANSI.red(
                  "WARNING! Managing repositories through transformers-cli is deprecated. "
                  "Please use `huggingface-cli` instead."   )   )
          __snake_case   :				Dict        =       HfFolder.get_token()
          if token is None:
            print("Not logged in"   )
            exit(1   )
          try:
            __snake_case   :				Optional[int]        =       subprocess.check_output(["git", "--version"]   ).decode("utf-8"   )
            print(ANSI.gray(stdout.strip()   )   )
          except FileNotFoundError:
            print("Looks like you do not have git installed, please install."   )
  
          try:
            __snake_case   :				Union[str, Any]        =       subprocess.check_output(["git-lfs", "--version"]   ).decode("utf-8"   )
            print(ANSI.gray(stdout.strip()   )   )
          except FileNotFoundError:
            print(
                ANSI.red(
                    "Looks like you do not have git-lfs installed, please install."
                    " You can install from https://git-lfs.github.com/."
                    " Then run `git lfs install` (you only have to do this once)."   )   )
          print(""   )
  
          __snake_case    ,   __snake_case   :				Any        =       whoami(UpperCAmelCase   )
          __snake_case   :				List[str]        =       self.args.organization if self.args.organization is not None else user
          __snake_case   :				Any        =       F"""{namespace}/{self.args.name}"""
          print(F"""You are about to create {ANSI.bold(UpperCAmelCase   )}"""   )
  
          if not self.args.yes:
            __snake_case   :				Optional[int]        =       input("Proceed? [Y/n] "   ).lower()
            if not (choice == "" or choice == "y" or choice == "yes"):
              print("Abort"   )
              exit()
          try:
            __snake_case   :				List[str]        =       create_repo(UpperCAmelCase				,						name=self.args.name				,						organization=self.args.organization   )
          except HTTPError as e:
            print(UpperCAmelCase   )
            print(ANSI.red(e.response.text   )   )
            exit(1   )
          print("\nYour repo now lives at:"   )
          print(F"""  {ANSI.bold(UpperCAmelCase   )}"""   )
          print("\nYou can clone it locally with the command below, and commit/push as usual."   )
          print(F"""\n  git clone {url}"""   )
          print(""   )
  
  