





import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402



class 							_UpperCamelCase					(	unittest.TestCase     ):

       def lowercase (		self: Tuple		)       -> Dict:




              """simple docstring"""


              UpperCamelCase_          =  0

       def lowercase (		self: Dict		)       -> Tuple:




              """simple docstring"""


              UpperCamelCase_          =  AutoImageProcessor.from_pretrained("openai/clip-vit-base-patch32"		)
              self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Any		)       -> List[Any]:




              """simple docstring"""


              with tempfile.TemporaryDirectory() as tmpdirname:
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "preprocessor_config.json"
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "config.json"
                     json.dump(
                         {"image_processor_type": "CLIPImageProcessor", "processor_class": "CLIPProcessor"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)					,   )
                     json.dump({"model_type": "clip"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)		)

                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)
                     self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> Dict:




              """simple docstring"""


              with tempfile.TemporaryDirectory() as tmpdirname:
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "preprocessor_config.json"
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "config.json"
                     json.dump(
                         {"feature_extractor_type": "CLIPFeatureExtractor", "processor_class": "CLIPProcessor"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)					,   )
                     json.dump({"model_type": "clip"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)		)

                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)
                     self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Optional[int]		)       -> Union[str, Any]:




              """simple docstring"""


              with tempfile.TemporaryDirectory() as tmpdirname:
                     UpperCamelCase_          =  CLIPConfig()

                     # Create a dummy config file with image_proceesor_type
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "preprocessor_config.json"
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "config.json"
                     json.dump(
                         {"image_processor_type": "CLIPImageProcessor", "processor_class": "CLIPProcessor"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)					,   )
                     json.dump({"model_type": "clip"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)		)

                     # remove image_processor_type to make sure config.json alone is enough to load image processor locally
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		).to_dict()

                     config_dict.pop("image_processor_type"		)
                     UpperCamelCase_          =  CLIPImageProcessor(**_SCREAMING_SNAKE_CASE		)

                     # save in new folder
                     model_config.save_pretrained(_SCREAMING_SNAKE_CASE		)
                     config.save_pretrained(_SCREAMING_SNAKE_CASE		)

                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)

                     # make sure private variable is not incorrectly saved
                     UpperCamelCase_          =  json.loads(config.to_json_string()		)
                     self.assertTrue("_processor_class" not in dict_as_saved		)

              self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> List[str]:




              """simple docstring"""


              with tempfile.TemporaryDirectory() as tmpdirname:
                     UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "preprocessor_config.json"
                     json.dump(
                         {"image_processor_type": "CLIPImageProcessor", "processor_class": "CLIPProcessor"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)					,   )

                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)
                     self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Dict		)       -> List[str]:




              """simple docstring"""


              with self.assertRaisesRegex(
                  _SCREAMING_SNAKE_CASE					,   "clip-base is not a local folder and is not a valid model identifier"		):
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained("clip-base"		)

       def lowercase (		self: Optional[int]		)       -> Tuple:




              """simple docstring"""


              with self.assertRaisesRegex(
                  _SCREAMING_SNAKE_CASE					,   R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"		):
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE					,   revision="aaaaaa"		)

       def lowercase (		self: List[str]		)       -> Union[str, Any]:




              """simple docstring"""


              with self.assertRaisesRegex(
                  _SCREAMING_SNAKE_CASE					,   "hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."					,   ):
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained("hf-internal-testing/config-no-model"		)

       def lowercase (		self: List[Any]		)       -> str:




              """simple docstring"""


              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained("hf-internal-testing/test_dynamic_image_processor"		)
              # If remote code is disabled, we can't load this config.
              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(
                         "hf-internal-testing/test_dynamic_image_processor"					,   trust_remote_code=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  AutoImageProcessor.from_pretrained(
                  "hf-internal-testing/test_dynamic_image_processor"					,   trust_remote_code=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(image_processor.__class__.__name__					,   "NewImageProcessor"		)

              # Test image processor can be reloaded.
              with tempfile.TemporaryDirectory() as tmp_dir:
                     image_processor.save_pretrained(_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE					,   trust_remote_code=_SCREAMING_SNAKE_CASE		)
              self.assertEqual(reloaded_image_processor.__class__.__name__					,   "NewImageProcessor"		)

       def lowercase (		self: Tuple		)       -> Union[str, Any]:




              """simple docstring"""


              try:
                     AutoConfig.register("custom"					,   _SCREAMING_SNAKE_CASE		)
                     AutoImageProcessor.register(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
                     # Trying to register something existing in the Transformers library will raise an error
                     with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                            AutoImageProcessor.register(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "preprocessor_config.json"
                            UpperCamelCase_          =  Path(_SCREAMING_SNAKE_CASE		) / "config.json"
                            json.dump(
                                {"feature_extractor_type": "CLIPFeatureExtractor", "processor_class": "CLIPProcessor"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)					,   )
                            json.dump({"model_type": "clip"}					,   open(_SCREAMING_SNAKE_CASE					,   "w"		)		)

                            UpperCamelCase_          =  CustomImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)

                     # Now that the config is registered, it can be used as any other config with the auto-API
                     with tempfile.TemporaryDirectory() as tmp_dir:
                            image_processor.save_pretrained(_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  AutoImageProcessor.from_pretrained(_SCREAMING_SNAKE_CASE		)
                            self.assertIsInstance(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)

              finally:
                     if "custom" in CONFIG_MAPPING._extra_content:
                            del CONFIG_MAPPING._extra_content["custom"]
                     if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                            del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]



       def lowercase (		self: Dict		)       -> Dict:




              """simple docstring"""


              class 							_UpperCamelCase					(	lowerCAmelCase_     ):
                     _UpperCamelCase				:					Any         =   True

              try:
                     AutoConfig.register("custom"					,   _SCREAMING_SNAKE_CASE		)
                     AutoImageProcessor.register(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
                     # If remote code is not set, the default is to use local
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained("hf-internal-testing/test_dynamic_image_processor"		)
                     self.assertEqual(image_processor.__class__.__name__					,   "NewImageProcessor"		)
                     self.assertTrue(image_processor.is_local		)

                     # If remote code is disabled, we load the local one.
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(
                         "hf-internal-testing/test_dynamic_image_processor"					,   trust_remote_code=_SCREAMING_SNAKE_CASE		)
                     self.assertEqual(image_processor.__class__.__name__					,   "NewImageProcessor"		)
                     self.assertTrue(image_processor.is_local		)

                     # If remote is enabled, we load from the Hub
                     UpperCamelCase_          =  AutoImageProcessor.from_pretrained(
                         "hf-internal-testing/test_dynamic_image_processor"					,   trust_remote_code=_SCREAMING_SNAKE_CASE		)
                     self.assertEqual(image_processor.__class__.__name__					,   "NewImageProcessor"		)
                     self.assertTrue(not hasattr(_SCREAMING_SNAKE_CASE					,   "is_local"		)		)

              finally:
                     if "custom" in CONFIG_MAPPING._extra_content:
                            del CONFIG_MAPPING._extra_content["custom"]
                     if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                            del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]








# 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  )     ->		int:
       # Format the message.
       if name is None:
              UpperCamelCase_          =  None
       else:
              UpperCamelCase_          =  "." * max(0					,		spaces - 2  ) + "# {:" + str(50 - spaces  ) + "s}"
              UpperCamelCase_          =  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_  )     ->		int:
       # 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.
       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(*UpperCamelCase_  )
              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(*UpperCamelCase_  )
              UpperCamelCase_          =  param.transpose(0					,		1  ).contiguous()
       UpperCamelCase_          =  param.view(*UpperCamelCase_  )
       return param




def  lowerCAmelCase_ (							UpperCamelCase_					,		UpperCamelCase_					,		UpperCamelCase_  )     ->		int:
       # The converted output model.
       UpperCamelCase_          =  {}

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

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

              # 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					,		UpperCamelCase_					,		UpperCamelCase_  )
                     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(UpperCamelCase_					,		UpperCamelCase_					,		3					,		UpperCamelCase_					,		UpperCamelCase_  )
                     # 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(UpperCamelCase_					,		UpperCamelCase_					,		3					,		UpperCamelCase_					,		UpperCamelCase_  )
                     # 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  lowerCAmelCase_ (							)     ->		Union[str, Any]:
       # Create the argument parser.
       UpperCamelCase_          =  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."					,		)
       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(UpperCamelCase_					,		map_location="cpu"  )
       else:
              UpperCamelCase_          =  torch.load(args.path_to_checkpoint					,		map_location="cpu"  )

       UpperCamelCase_          =  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:
                            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=UpperCamelCase_					,		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=UpperCamelCase_					,		summary_activation=UpperCamelCase_					,		summary_proj_to_labels=UpperCamelCase_					,		summary_first_dropout=0.1					,		scale_attn_weights=UpperCamelCase_					,		use_cache=UpperCamelCase_					,		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(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:
              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(UpperCamelCase_  )
       UpperCamelCase_          =  type(UpperCamelCase_  ).__name__
       UpperCamelCase_          =  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.
       UpperCamelCase_          =  os.path.join(UpperCamelCase_					,		"pytorch_model.bin"  )
       print(F'''Saving checkpoint to "{output_checkpoint_file}"'''  )
       torch.save(UpperCamelCase_					,		UpperCamelCase_  )


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

if __name__ == "__main__":
     main()

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

