






"""simple docstring"""





import argparse
import os
import re
import zipfile

import torch

from transformers import AutoTokenizer, GPTaConfig
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__=0				):





      if name is None:
            UpperCAmelCase_      			=							None
      else:
            UpperCAmelCase_      			=							"." * max(0     ,      spaces - 2				) + "# {:" + str(50 - spaces				) + "s}"
            UpperCAmelCase_      			=							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        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):
      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(*_SCREAMING_SNAKE_CASE				)
            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(*_SCREAMING_SNAKE_CASE				)
            UpperCAmelCase_      			=							param.transpose(0     ,      1				).contiguous()
      UpperCAmelCase_      			=							param.view(*_SCREAMING_SNAKE_CASE				)
      return param
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):
      UpperCAmelCase_      			=							{}

      # old versions did not store training args
      UpperCAmelCase_      			=							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))

            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(_SCREAMING_SNAKE_CASE				)

            # 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     ,      _SCREAMING_SNAKE_CASE     ,      _SCREAMING_SNAKE_CASE				)
                  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(_SCREAMING_SNAKE_CASE     ,      _SCREAMING_SNAKE_CASE     ,      3     ,      _SCREAMING_SNAKE_CASE     ,      _SCREAMING_SNAKE_CASE				)
                  # 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(_SCREAMING_SNAKE_CASE     ,      _SCREAMING_SNAKE_CASE     ,      3     ,      _SCREAMING_SNAKE_CASE     ,      _SCREAMING_SNAKE_CASE				)
                  # 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        a__     (				):
      UpperCAmelCase_      			=							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 checkpoint file (.zip archive or direct .pt file)"     ,      )
      parser.add_argument(
          "--config_file"     ,      default=""     ,      type=_SCREAMING_SNAKE_CASE     ,      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(_SCREAMING_SNAKE_CASE     ,      map_location="cpu"				)
      else:
            UpperCAmelCase_      			=							torch.load(args.path_to_checkpoint     ,      map_location="cpu"				)

      UpperCAmelCase_      			=							input_state_dict.get("args"     ,      _SCREAMING_SNAKE_CASE				)

      # 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=50257     ,      n_positions=1024     ,      n_embd=1024     ,      n_layer=24     ,      n_head=16     ,      n_inner=4096     ,      activation_function=_SCREAMING_SNAKE_CASE     ,      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=_SCREAMING_SNAKE_CASE     ,      summary_activation=_SCREAMING_SNAKE_CASE     ,      summary_proj_to_labels=_SCREAMING_SNAKE_CASE     ,      summary_first_dropout=0.1     ,      scale_attn_weights=_SCREAMING_SNAKE_CASE     ,      use_cache=_SCREAMING_SNAKE_CASE     ,      bos_token_id=50256     ,      eos_token_id=50256     ,      )
      else:
            UpperCAmelCase_      			=							GPTaConfig.from_json_file(args.config_file				)

      UpperCAmelCase_      			=							["GPT2LMHeadModel"]

      # Convert.
      print("Converting"				)
      UpperCAmelCase_      			=							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				)

      # 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(_SCREAMING_SNAKE_CASE				)
      UpperCAmelCase_      			=							type(_SCREAMING_SNAKE_CASE				).__name__
      UpperCAmelCase_      			=							tokenizer_class

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

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

      # Store the state_dict to file.
      UpperCAmelCase_      			=							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()

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









"""simple docstring"""





import warnings

from ...utils import logging
from .image_processing_flava import FlavaImageProcessor


lowerCamelCase       	=						logging.get_logger(__name__)
class 				lowercase__       (   SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       def __init__(							self      :  Union[str, Any]		,    *_UpperCAmelCase      :  List[Any]		,    **_UpperCAmelCase      :  Union[str, Any]   )			->						None:


             '''simple docstring'''




             warnings.warn(
                 "The class FlavaFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
                 " use FlavaImageProcessor instead."		,    _UpperCAmelCase		,    )
             super().__init__(*_UpperCAmelCase		,    **_UpperCAmelCase   )




