






import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase=0   ):


    '''simple docstring'''



    if name is None:
        _lowercase     			=						None
    else:
        _lowercase     			=						'.' * max(0,   spaces - 2   ) + '# {:' + str(50 - spaces   ) + 's}'
        _lowercase     			=						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 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



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


    '''simple docstring'''



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

    _lowercase     			=						sd_megatron['model.language_model.embedding.word_embeddings.weight'][:vocab_size, :]
    _lowercase     			=						{
        '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'],
    }

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

        _lowercase     			=						sd_megatron[f'''{pf}{i}.input_layernorm.weight''']
        _lowercase     			=						sd_megatron[f'''{pf}{i}.input_layernorm.bias''']

        _lowercase     			=						sd_megatron[f'''{pf}{i}.self_attention.query_key_value.weight''']
        _lowercase     			=						fix_query_key_value_ordering(_SCREAMING_SNAKE_CASE,   3,   _SCREAMING_SNAKE_CASE,   _SCREAMING_SNAKE_CASE   )
        _lowercase     			=						vala.transpose(0,   1   ).contiguous()

        _lowercase     			=						sd_megatron[f'''{pf}{i}.self_attention.query_key_value.bias''']
        _lowercase     			=						fix_query_key_value_ordering(_SCREAMING_SNAKE_CASE,   3,   _SCREAMING_SNAKE_CASE,   _SCREAMING_SNAKE_CASE   )
        _lowercase     			=						vala

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

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

    return sd_hf
def 							UpperCAmelCase   ( _lowercase,   _lowercase   ):


    '''simple docstring'''



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

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

    return config_hf
def 							UpperCAmelCase   ( _lowercase   ):


    '''simple docstring'''



    print(_SCREAMING_SNAKE_CASE   )

    _lowercase     			=						args.checkpoint_path
    _lowercase     			=						args.save_path
    if isfile(_SCREAMING_SNAKE_CASE   ):
        raise FileNotFoundError(f'''ERROR! could not find file {checkpoint_path}'''   )

    # Load the model.
    _lowercase     			=						torch.load(_SCREAMING_SNAKE_CASE,   map_location='cpu'   )

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

    _lowercase     			=						checkpoint['state_dict']

    # Convert.
    print('Converting'   )
    _lowercase     			=						convert_megatron_checkpoint(_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   )

    _lowercase     			=						'GPTSw3Tokenizer'

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

    # Store the state_dict to file.
    _lowercase     			=						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__":
   lowerCamelCase					    =	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''')
   lowerCamelCase					    =	parser.parse_args()
   main(_args)



import unittest

from transformers import ErnieMTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


UpperCamelCase__	:   List[str]								=     get_tests_dir("fixtures/spiece.model")



@require_sentencepiece
@require_tokenizers
class    __snake_case	(     lowerCAmelCase__  ,			unittest.TestCase							):
  __lowerCAmelCase  :						Union[str, Any]    			=	ErnieMTokenizer
  __lowerCAmelCase  :						Union[str, Any]    			=	False
  __lowerCAmelCase  :						Optional[int]    			=	True
  __lowerCAmelCase  :						List[str]    			=	False
  __lowerCAmelCase  :						List[str]    			=	False

  def     lowerCAmelCase__  (       self):
      super().setUp()

      # We have a SentencePiece fixture for testing
      SCREAMING_SNAKE_CASE_						=    ErnieMTokenizer(_A    ,		unk_token='<unk>'    ,		pad_token='<pad>')
      tokenizer.save_pretrained(self.tmpdirname)

  def     lowerCAmelCase__  (       self    ,		_A):
      SCREAMING_SNAKE_CASE_						=    'this is a test'
      SCREAMING_SNAKE_CASE_						=    'this is a test'
      return input_text, output_text

  def     lowerCAmelCase__  (       self):
      SCREAMING_SNAKE_CASE_						=    '<pad>'
      SCREAMING_SNAKE_CASE_						=    0

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(_A)    ,		_A)
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(_A)    ,		_A)

  def     lowerCAmelCase__  (       self):
      SCREAMING_SNAKE_CASE_						=    list(self.get_tokenizer().get_vocab().keys())

      self.assertEqual(vocab_keys[0]    ,		'<pad>')
      self.assertEqual(vocab_keys[1]    ,		'<unk>')
      self.assertEqual(vocab_keys[-1]    ,		'▁eloquent')
      self.assertEqual(len(_A)    ,		30000)

  def     lowerCAmelCase__  (       self):
      self.assertEqual(self.get_tokenizer().vocab_size    ,		30000)

  def     lowerCAmelCase__  (       self):
      if not self.test_rust_tokenizer:
          return

      SCREAMING_SNAKE_CASE_						=    self.get_tokenizer()
      SCREAMING_SNAKE_CASE_						=    self.get_rust_tokenizer()

      SCREAMING_SNAKE_CASE_						=    'I was born in 92000, and this is falsé.'

      SCREAMING_SNAKE_CASE_						=    tokenizer.tokenize(_A)
      SCREAMING_SNAKE_CASE_						=    rust_tokenizer.tokenize(_A)
      self.assertListEqual(_A    ,		_A)

      SCREAMING_SNAKE_CASE_						=    tokenizer.encode(_A    ,		add_special_tokens=_A)
      SCREAMING_SNAKE_CASE_						=    rust_tokenizer.encode(_A    ,		add_special_tokens=_A)
      self.assertListEqual(_A    ,		_A)

      SCREAMING_SNAKE_CASE_						=    self.get_rust_tokenizer()
      SCREAMING_SNAKE_CASE_						=    tokenizer.encode(_A)
      SCREAMING_SNAKE_CASE_						=    rust_tokenizer.encode(_A)
      self.assertListEqual(_A    ,		_A)

  def     lowerCAmelCase__  (       self):
      SCREAMING_SNAKE_CASE_						=    ErnieMTokenizer(_A    ,		do_lower_case=_A    ,		unk_token='<unk>'    ,		pad_token='<pad>')

      SCREAMING_SNAKE_CASE_						=    tokenizer.tokenize('This is a test')
      self.assertListEqual(_A    ,		['▁this', '▁is', '▁a', '▁test'])

      self.assertListEqual(tokenizer.convert_tokens_to_ids(_A)    ,		[48, 25, 21, 1289])

      SCREAMING_SNAKE_CASE_						=    tokenizer.tokenize('I was born in 92000, and this is falsé.')
      # ErnieMTokenizer(paddlenlp implementation) outputs '9' instead of '_9' so to mimic that '_9' is changed to '9'
      self.assertListEqual(
          _A    ,		['▁i', '▁was', '▁born', '▁in', '9', '2000', ',', '▁and', '▁this', '▁is', '▁fal', 's', 'é', '.'])
      SCREAMING_SNAKE_CASE_						=    tokenizer.convert_tokens_to_ids(_A)
      self.assertListEqual(_A    ,		[31, 23, 386, 19, 518, 3050, 15, 17, 48, 25, 8256, 18, 1, 9])

      SCREAMING_SNAKE_CASE_						=    tokenizer.convert_ids_to_tokens(_A)
      self.assertListEqual(
          _A    ,		['▁i', '▁was', '▁born', '▁in', '9', '2000', ',', '▁and', '▁this', '▁is', '▁fal', 's', '<unk>', '.']    ,		)

  def     lowerCAmelCase__  (       self):
      SCREAMING_SNAKE_CASE_						=    ErnieMTokenizer(_A    ,		unk_token='<unk>'    ,		pad_token='<pad>')

      SCREAMING_SNAKE_CASE_						=    tokenizer.encode('sequence builders')
      SCREAMING_SNAKE_CASE_						=    tokenizer.encode('multi-sequence build')

      SCREAMING_SNAKE_CASE_						=    tokenizer.build_inputs_with_special_tokens(_A)
      SCREAMING_SNAKE_CASE_						=    tokenizer.build_inputs_with_special_tokens(_A    ,		_A)

      assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
      assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + [
          tokenizer.sep_token_id
      ] + text_a + [tokenizer.sep_token_id]


  @slow
  def     lowerCAmelCase__  (       self):
      # fmt: off
      SCREAMING_SNAKE_CASE_						=    {'input_ids': [[0, 11062, 82772, 7, 15, 82772, 538, 51529, 237, 17198, 1290, 206, 9, 215175, 1314, 136, 17198, 1290, 206, 9, 56359, 42, 122009, 9, 16466, 16, 87344, 4537, 9, 4717, 78381, 6, 159958, 7, 15, 24480, 618, 4, 527, 22693, 9, 304, 4, 2777, 24480, 9874, 4, 43523, 594, 4, 803, 18392, 33189, 18, 4, 43523, 24447, 5, 5, 5, 16, 100, 24955, 83658, 9626, 144057, 15, 839, 22335, 16, 136, 24955, 83658, 83479, 15, 39102, 724, 16, 678, 645, 6460, 1328, 4589, 42, 122009, 115774, 23, 3559, 1328, 46876, 7, 136, 53894, 1940, 42227, 41159, 17721, 823, 425, 4, 27512, 98722, 206, 136, 5531, 4970, 919, 17336, 5, 2], [0, 20080, 618, 83, 82775, 47, 479, 9, 1517, 73, 53894, 333, 80581, 110117, 18811, 5256, 1295, 51, 152526, 297, 7986, 390, 124416, 538, 35431, 214, 98, 15044, 25737, 136, 7108, 43701, 23, 756, 135355, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 581, 63773, 119455, 6, 147797, 88203, 7, 645, 70, 21, 3285, 10269, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=_A    ,		model_name='susnato/ernie-m-base_pytorch'    ,		sequences=[
              'Transformers (formerly known as pytorch-transformers and pytorch-pretrained-bert) provides '
              'general-purpose architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet...) for Natural '
              'Language Understanding (NLU) and Natural Language Generation (NLG) with over32+ pretrained '
              'models in100+ languages and deep interoperability between Jax, PyTorch and TensorFlow.',
              'BERT is designed to pre-train deep bidirectional representations from unlabeled text by jointly '
              'conditioning on both left and right context in all layers.',
              'The quick brown fox jumps over the lazy dog.',
          ]    ,		)


