




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE     : Dict      =logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE     : List[str]      ={
    """openbmb/cpm-ant-10b""": """https://huggingface.co/openbmb/cpm-ant-10b/blob/main/config.json"""
    # See all CPMAnt models at https://huggingface.co/models?filter=cpmant
}






class       A      (		_lowerCAmelCase		):
     """simple docstring"""

     _snake_case:						Union[str, Any]  ="cpmant"




     def __init__(     self	: int						,	A__	: int = 3_0_7_2_0						,	A__	: int = 4_0_9_6						,	A__	: int = 3_2						,	A__	: int = 1_2_8						,	A__	: int = 1_0_2_4_0						,	A__	: int = 4_8						,	A__	: int = 0.0						,	A__	: int = 5_1_2						,	A__	: int = 2_0_4_8						,	A__	: int = 1E-6						,	A__	: float = 1.0						,	A__	: int = 3_2						,	A__	: int = 3_2						,	A__	: int = 3_2						,	A__	: bool = True						,	**A__	: Optional[int]						,	)					->							int:
           super().__init__(**A__      )
           lowerCamelCase__   	=					prompt_types
           lowerCamelCase__   	=					prompt_length
           lowerCamelCase__   	=					segment_types
           lowerCamelCase__   	=					hidden_size
           lowerCamelCase__   	=					num_attention_heads
           lowerCamelCase__   	=					dim_head
           lowerCamelCase__   	=					dim_ff
           lowerCamelCase__   	=					num_hidden_layers
           lowerCamelCase__   	=					position_bias_num_buckets
           lowerCamelCase__   	=					position_bias_max_distance
           lowerCamelCase__   	=					dropout_p
           lowerCamelCase__   	=					eps
           lowerCamelCase__   	=					use_cache
           lowerCamelCase__   	=					vocab_size
           lowerCamelCase__   	=					init_std





import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig

def       lowerCAmelCase_      (	lowerCamelCase:							str  ,					lowerCamelCase:							List[str]  ,					lowerCamelCase:							List[Any]=0    )       ->      Dict:


      '''simple docstring'''




      if name is None:
            lowerCamelCase__   	=					None
      else:
            lowerCamelCase__   	=					'.' * max(0  ,					spaces - 2    ) + '# {:' + str(50 - spaces    ) + 's}'
            lowerCamelCase__   	=					fmt.format(lowerCamelCase    )

      # Print and recurse (if needed).
      if isinstance(lowerCamelCase  ,					lowerCamelCase    ):
            if msg is not None:
                  print(lowerCamelCase    )
            for k in val.keys():
                  recursive_print(lowerCamelCase  ,					val[k]  ,					spaces + 2    )
      elif isinstance(lowerCamelCase  ,					torch.Tensor    ):
            print(lowerCamelCase  ,					':'  ,					val.size()    )
      else:
            print(lowerCamelCase  ,					':'  ,					lowerCamelCase    )

def       lowerCAmelCase_      (	lowerCamelCase:							Tuple  ,					lowerCamelCase:							Dict  ,					lowerCamelCase:							List[str]  ,					lowerCamelCase:							Any    )       ->      Optional[int]:


      '''simple docstring'''




      lowerCamelCase__   	=					param.size()
      # other versions store [num_heads * num_splits * hidden_size, :]
      lowerCamelCase__   	=					(num_heads, num_splits, hidden_size) + input_shape[1:]
      lowerCamelCase__   	=					param.view(*lowerCamelCase    )
      lowerCamelCase__   	=					param.transpose(0  ,					1    ).contiguous()
      lowerCamelCase__   	=					param.view(*lowerCamelCase    )
      return param

def       lowerCAmelCase_      (	lowerCamelCase:							List[str]  ,					lowerCamelCase:							int    )       ->      int:


      '''simple docstring'''




      lowerCamelCase__   	=					config.n_positions
      lowerCamelCase__   	=					config.n_layer
      lowerCamelCase__   	=					config.vocab_size
      lowerCamelCase__   	=					config.n_head
      lowerCamelCase__   	=					config.n_embd // config.n_head

      lowerCamelCase__   	=					sd_megatron['model.language_model.embedding.word_embeddings.weight'][:vocab_size, :]
      lowerCamelCase__   	=					{
          'transformer.wte.weight': word_embeddings,
          'transformer.wpe.weight': sd_megatron['model.language_model.embedding.position_embeddings.weight'],
          'transformer.ln_f.weight': sd_megatron['model.language_model.encoder.final_layernorm.weight'],
          'transformer.ln_f.bias': sd_megatron['model.language_model.encoder.final_layernorm.bias'],
      }

      lowerCamelCase__   	=					'model.language_model.encoder.layers.'
      for i in range(lowerCamelCase    ):
            lowerCamelCase__   	=					torch.tril(torch.ones((n_positions, n_positions)  ,					dtype=torch.bool    )    )
            lowerCamelCase__   	=					causal_mask.view(1  ,					1  ,					lowerCamelCase  ,					lowerCamelCase    )
            lowerCamelCase__   	=					causal_mask
            lowerCamelCase__   	=					torch.tensor(-1e4  ,					dtype=torch.bfloataa    )

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.input_layernorm.weight"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.input_layernorm.bias"""]

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.query_key_value.weight"""]
            lowerCamelCase__   	=					fix_query_key_value_ordering(lowerCamelCase  ,					3  ,					lowerCamelCase  ,					lowerCamelCase    )
            lowerCamelCase__   	=					vala.transpose(0  ,					1    ).contiguous()

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.query_key_value.bias"""]
            lowerCamelCase__   	=					fix_query_key_value_ordering(lowerCamelCase  ,					3  ,					lowerCamelCase  ,					lowerCamelCase    )
            lowerCamelCase__   	=					vala

            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.dense.weight"""].transpose(
                0  ,					1    )
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.self_attention.dense.bias"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.post_attention_layernorm.weight"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.post_attention_layernorm.bias"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.mlp.dense_h_to_4h.weight"""].transpose(0  ,					1    )
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.mlp.dense_h_to_4h.bias"""]
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.mlp.dense_4h_to_h.weight"""].transpose(
                0  ,					1    )
            lowerCamelCase__   	=					sd_megatron[f"""{pf}{i}.mlp.dense_4h_to_h.bias"""]

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

      return sd_hf

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[int]  ,					lowerCamelCase:							List[Any]    )       ->      Tuple:


      '''simple docstring'''




      lowerCamelCase__   	=					6_40_00
      lowerCamelCase__   	=					config_megatron['encoder_seq_length']
      lowerCamelCase__   	=					config_megatron['hidden_size']
      lowerCamelCase__   	=					config_megatron['num_layers']
      lowerCamelCase__   	=					config_megatron['num_attention_heads']
      lowerCamelCase__   	=					config_megatron['ffn_hidden_size']
      lowerCamelCase__   	=					'gelu'
      lowerCamelCase__   	=					0.1
      lowerCamelCase__   	=					0.1
      lowerCamelCase__   	=					0.1
      lowerCamelCase__   	=					config_megatron['layernorm_epsilon']  # 1e-5
      lowerCamelCase__   	=					config_megatron['init_method_std']  # 0.02
      lowerCamelCase__   	=					config_megatron['apply_query_key_layer_scaling']  # True
      lowerCamelCase__   	=					True
      lowerCamelCase__   	=					True

      # This identifies the 6.7B (7B) model which uses a different tokenizer
      if config_megatron["hidden_size"] == 40_96:
            lowerCamelCase__   	=					1  # <|endoftext|>
            lowerCamelCase__   	=					1  # <|endoftext|>
            lowerCamelCase__   	=					0  # <unk>
      else:
            lowerCamelCase__   	=					2  # <s>
            lowerCamelCase__   	=					3  # <|endoftext|>
            lowerCamelCase__   	=					0  # <pad>

      return config_hf

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[Any]    )       ->      Dict:


      '''simple docstring'''




      print(lowerCamelCase    )

      lowerCamelCase__   	=					args.checkpoint_path
      lowerCamelCase__   	=					args.save_path
      if isfile(lowerCamelCase    ):
            raise FileNotFoundError(f"""ERROR! could not find file {checkpoint_path}"""    )

      # Load the model.
      lowerCamelCase__   	=					torch.load(lowerCamelCase  ,					map_location='cpu'    )

      # Load the config.
      lowerCamelCase__   	=					checkpoint['hyper_parameters']['cfg']
      lowerCamelCase__   	=					GPTaConfig()
      lowerCamelCase__   	=					copy_config(config_hf=lowerCamelCase  ,					config_megatron=lowerCamelCase    )
      lowerCamelCase__   	=					['GPT2LMHeadModel']

      lowerCamelCase__   	=					checkpoint['state_dict']

      # Convert.
      print('Converting'    )
      lowerCamelCase__   	=					convert_megatron_checkpoint(lowerCamelCase  ,					lowerCamelCase    )

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

      lowerCamelCase__   	=					'GPTSw3Tokenizer'

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

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


if __name__ == "__main__":
      __SCREAMING_SNAKE_CASE     : Optional[int]      =argparse.ArgumentParser()
      parser.add_argument(
          """--checkpoint_path""",
          type=str,
          required=True,
          help="""e.g. megatron_gpt--val_loss=2.42-step=38000-consumed_samples=54720000""",
      )
      parser.add_argument("""--save_path""", type=str, required=True, help="""e.g. /home/user/gpt-sw3/hf""")
      parser.add_argument("""--print-checkpoint-structure""", action="""store_true""")
      __SCREAMING_SNAKE_CASE     : Tuple      =parser.parse_args()
      main(_args)