


"""simple docstring"""


import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    RemBertConfig,
		    RemBertForCausalLM,
		    RemBertForMaskedLM,
		    RemBertForMultipleChoice,
		    RemBertForQuestionAnswering,
		    RemBertForSequenceClassification,
		    RemBertForTokenClassification,
		    RemBertModel,
		)
		from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST
class snake_case   :







		def __init__(       self   :							Optional[Any]     , a__   :							Union[str, Any]     , a__   :							Tuple=13     , a__   :							Union[str, Any]=7     , a__   :							Optional[Any]=True     , a__   :							Optional[Any]=True     , a__   :							Optional[int]=True     , a__   :							List[Any]=True     , a__   :							Any=99     , a__   :							int=32     , a__   :							Dict=18     , a__   :							Optional[int]=43     , a__   :							Dict=5     , a__   :							int=4     , a__   :							List[str]=37     , a__   :							Dict="gelu"     , a__   :							Dict=0.1     , a__   :							Optional[int]=0.1     , a__   :							List[Any]=5_12     , a__   :							Union[str, Any]=16     , a__   :							Union[str, Any]=2     , a__   :							Dict=0.0_2     , a__   :							int=3     , a__   :							List[Any]=4     , a__   :							Optional[int]=None     , )							->    List[str]:





						'''simple docstring'''







						_A       	=      parent
						_A       	=      batch_size
						_A       	=      seq_length
						_A       	=      is_training
						_A       	=      use_input_mask
						_A       	=      use_token_type_ids
						_A       	=      use_labels
						_A       	=      vocab_size
						_A       	=      hidden_size
						_A       	=      input_embedding_size
						_A       	=      output_embedding_size
						_A       	=      num_hidden_layers
						_A       	=      num_attention_heads
						_A       	=      intermediate_size
						_A       	=      hidden_act
						_A       	=      hidden_dropout_prob
						_A       	=      attention_probs_dropout_prob
						_A       	=      max_position_embeddings
						_A       	=      type_vocab_size
						_A       	=      type_sequence_label_size
						_A       	=      initializer_range
						_A       	=      num_labels
						_A       	=      num_choices
						_A       	=      scope







		def 			a_			(       self   :							List[str]		)							->    str:





						'''simple docstring'''







						_A       	=      ids_tensor([self.batch_size, self.seq_length]     , self.vocab_size		)

						_A       	=      None
						if self.use_input_mask:
										_A       	=      random_attention_mask([self.batch_size, self.seq_length]		)

						_A       	=      None
						if self.use_token_type_ids:
										_A       	=      ids_tensor([self.batch_size, self.seq_length]     , self.type_vocab_size		)

						_A       	=      None
						_A       	=      None
						_A       	=      None
						if self.use_labels:
										_A       	=      ids_tensor([self.batch_size]     , self.type_sequence_label_size		)
										_A       	=      ids_tensor([self.batch_size, self.seq_length]     , self.num_labels		)
										_A       	=      ids_tensor([self.batch_size]     , self.num_choices		)

						_A       	=      RemBertConfig(
						    vocab_size=self.vocab_size     , hidden_size=self.hidden_size     , input_embedding_size=self.input_embedding_size     , output_embedding_size=self.output_embedding_size     , num_hidden_layers=self.num_hidden_layers     , num_attention_heads=self.num_attention_heads     , intermediate_size=self.intermediate_size     , hidden_act=self.hidden_act     , hidden_dropout_prob=self.hidden_dropout_prob     , attention_probs_dropout_prob=self.attention_probs_dropout_prob     , max_position_embeddings=self.max_position_embeddings     , type_vocab_size=self.type_vocab_size     , is_decoder=__UpperCamelCase     , initializer_range=self.initializer_range     , )

						return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







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





						'''simple docstring'''







						(
						    (
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,			
						)       	=      self.prepare_config_and_inputs()

						_A       	=      True
						_A       	=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
						_A       	=      ids_tensor([self.batch_size, self.seq_length]     , vocab_size=2		)

						return (
						    config,
						    input_ids,
						    token_type_ids,
						    input_mask,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)







		def 			a_			(       self   :							Dict     , a__   :							Dict     , a__   :							Tuple     , a__   :							Optional[Any]     , a__   :							Tuple     , a__   :							Optional[Any]     , a__   :							Optional[Any]     , a__   :							str		)							->    Optional[int]:





						'''simple docstring'''







						_A       	=      RemBertModel(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase		)
						_A       	=      model(__UpperCamelCase     , token_type_ids=__UpperCamelCase		)
						_A       	=      model(__UpperCamelCase		)
						self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)		)







		def 			a_			(       self   :							Optional[Any]     , a__   :							Optional[int]     , a__   :							Union[str, Any]     , a__   :							Dict     , a__   :							Union[str, Any]     , a__   :							Any     , a__   :							Any     , a__   :							Tuple     , a__   :							Dict     , a__   :							Dict     , )							->    Tuple:





						'''simple docstring'''







						_A       	=      True
						_A       	=      RemBertModel(__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , encoder_hidden_states=__UpperCamelCase     , encoder_attention_mask=__UpperCamelCase     , )
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , encoder_hidden_states=__UpperCamelCase     , )
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase		)
						self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)		)







		def 			a_			(       self   :							Optional[int]     , a__   :							int     , a__   :							List[str]     , a__   :							Optional[Any]     , a__   :							Optional[Any]     , a__   :							Optional[Any]     , a__   :							List[Any]     , a__   :							Any     , a__   :							Union[str, Any]     , a__   :							str     , )							->    Optional[int]:





						'''simple docstring'''







						_A       	=      RemBertForCausalLM(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , labels=__UpperCamelCase		)
						self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.vocab_size)		)







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





						'''simple docstring'''







						_A       	=      RemBertForMaskedLM(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , labels=__UpperCamelCase		)
						self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.vocab_size)		)







		def 			a_			(       self   :							Optional[int]     , a__   :							Optional[int]     , a__   :							Union[str, Any]     , a__   :							Optional[int]     , a__   :							Optional[Any]     , a__   :							Tuple     , a__   :							Optional[int]     , a__   :							str     , a__   :							Any     , a__   :							Optional[Any]     , )							->    str:





						'''simple docstring'''







						_A       	=      True
						_A       	=      True
						_A       	=      RemBertForCausalLM(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()

						# first forward pass
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , encoder_hidden_states=__UpperCamelCase     , encoder_attention_mask=__UpperCamelCase     , use_cache=__UpperCamelCase     , )
						_A       	=      outputs.past_key_values

						# create hypothetical multiple next token and extent to next_input_ids
						_A       	=      ids_tensor((self.batch_size, 3)     , config.vocab_size		)
						_A       	=      ids_tensor((self.batch_size, 3)     , vocab_size=2		)

						# append to next input_ids and
						_A       	=      torch.cat([input_ids, next_tokens]     , dim=-1		)
						_A       	=      torch.cat([input_mask, next_mask]     , dim=-1		)

						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , encoder_hidden_states=__UpperCamelCase     , encoder_attention_mask=__UpperCamelCase     , output_hidden_states=__UpperCamelCase     , )["hidden_states"][0]
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , encoder_hidden_states=__UpperCamelCase     , encoder_attention_mask=__UpperCamelCase     , past_key_values=__UpperCamelCase     , output_hidden_states=__UpperCamelCase     , )["hidden_states"][0]

						# select random slice
						_A       	=      ids_tensor((1,)     , output_from_past.shape[-1]		).item()
						_A       	=      output_from_no_past[:, -3:, random_slice_idx].detach()
						_A       	=      output_from_past[:, :, random_slice_idx].detach()

						self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

						# test that outputs are equal for slice
						self.parent.assertTrue(torch.allclose(__UpperCamelCase     , __UpperCamelCase     , atol=1E-3		)		)







		def 			a_			(       self   :							int     , a__   :							List[Any]     , a__   :							Union[str, Any]     , a__   :							Union[str, Any]     , a__   :							Any     , a__   :							Dict     , a__   :							Dict     , a__   :							Optional[int]		)							->    str:





						'''simple docstring'''







						_A       	=      RemBertForQuestionAnswering(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , start_positions=__UpperCamelCase     , end_positions=__UpperCamelCase     , )
						self.parent.assertEqual(result.start_logits.shape     , (self.batch_size, self.seq_length)		)
						self.parent.assertEqual(result.end_logits.shape     , (self.batch_size, self.seq_length)		)







		def 			a_			(       self   :							int     , a__   :							int     , a__   :							str     , a__   :							Optional[Any]     , a__   :							Dict     , a__   :							List[str]     , a__   :							Optional[int]     , a__   :							Optional[Any]		)							->    Optional[int]:





						'''simple docstring'''







						_A       	=      self.num_labels
						_A       	=      RemBertForSequenceClassification(__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , labels=__UpperCamelCase		)
						self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.num_labels)		)







		def 			a_			(       self   :							int     , a__   :							Optional[Any]     , a__   :							Dict     , a__   :							str     , a__   :							str     , a__   :							List[Any]     , a__   :							int     , a__   :							Any		)							->    List[Any]:





						'''simple docstring'''







						_A       	=      self.num_labels
						_A       	=      RemBertForTokenClassification(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      model(__UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , labels=__UpperCamelCase		)
						self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.num_labels)		)







		def 			a_			(       self   :							Tuple     , a__   :							Union[str, Any]     , a__   :							List[Any]     , a__   :							Optional[int]     , a__   :							List[str]     , a__   :							Dict     , a__   :							str     , a__   :							str		)							->    Any:





						'''simple docstring'''







						_A       	=      self.num_choices
						_A       	=      RemBertForMultipleChoice(config=__UpperCamelCase		)
						model.to(__UpperCamelCase		)
						model.eval()
						_A       	=      input_ids.unsqueeze(1		).expand(-1     , self.num_choices     , -1		).contiguous()
						_A       	=      token_type_ids.unsqueeze(1		).expand(-1     , self.num_choices     , -1		).contiguous()
						_A       	=      input_mask.unsqueeze(1		).expand(-1     , self.num_choices     , -1		).contiguous()
						_A       	=      model(
						    __UpperCamelCase     , attention_mask=__UpperCamelCase     , token_type_ids=__UpperCamelCase     , labels=__UpperCamelCase     , )
						self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.num_choices)		)







		def 			a_			(       self   :							int		)							->    Optional[Any]:





						'''simple docstring'''







						_A       	=      self.prepare_config_and_inputs()
						(
						    (
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,			
						)       	=      config_and_inputs
						_A       	=      {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class snake_case   (	__snake_case  ,					__snake_case  ,					unittest.TestCase):
		__UpperCamelCase							     =      (
		    (
		        RemBertModel,
		        RemBertForMaskedLM,
		        RemBertForCausalLM,
		        RemBertForMultipleChoice,
		        RemBertForQuestionAnswering,
		        RemBertForSequenceClassification,
		        RemBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__UpperCamelCase							     =      (RemBertForCausalLM,) if is_torch_available() else ()
		__UpperCamelCase							     =      (
		    {
		        'feature-extraction': RemBertModel,
		        'fill-mask': RemBertForMaskedLM,
		        'question-answering': RemBertForQuestionAnswering,
		        'text-classification': RemBertForSequenceClassification,
		        'text-generation': RemBertForCausalLM,
		        'token-classification': RemBertForTokenClassification,
		        'zero-shot': RemBertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)







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





						'''simple docstring'''







						_A       	=      RemBertModelTester(self		)
						_A       	=      ConfigTester(self     , config_class=__UpperCamelCase     , hidden_size=37		)







		def 			a_			(       self   :							Optional[Any]		)							->    Optional[Any]:





						'''simple docstring'''







						self.config_tester.run_common_tests()







		def 			a_			(       self   :							Optional[int]		)							->    str:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__UpperCamelCase		)







		def 			a_			(       self   :							Optional[Any]		)							->    Any:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										_A       	=      type
										self.model_tester.create_and_check_model(*__UpperCamelCase		)







		def 			a_			(       self   :							Optional[int]		)							->    List[str]:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase		)







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





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase		)







		def 			a_			(       self   :							Dict		)							->    List[Any]:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_decoder_model_past_large_inputs(*__UpperCamelCase		)







		def 			a_			(       self   :							Optional[int]		)							->    Any:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase		)







		def 			a_			(       self   :							Optional[int]		)							->    Dict:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase		)







		def 			a_			(       self   :							Dict		)							->    Optional[Any]:





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase		)







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





						'''simple docstring'''







						_A       	=      self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*__UpperCamelCase		)







		def 			a_			(       self   :							Any		)							->    str:





						'''simple docstring'''







						(
						    (
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,		(
						    _A	
						)	,			
						)       	=      self.model_tester.prepare_config_and_inputs_for_decoder()

						_A       	=      None

						self.model_tester.create_and_check_model_as_decoder(
						    __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , __UpperCamelCase     , )







		@slow
		def 			a_			(       self   :							Dict		)							->    Tuple:





						'''simple docstring'''







						for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										_A       	=      RemBertModel.from_pretrained(__UpperCamelCase		)
										self.assertIsNotNone(__UpperCamelCase		)





@require_torch
class snake_case   (	unittest.TestCase):







		@slow
		def 			a_			(       self   :							Tuple		)							->    int:





						'''simple docstring'''







						_A       	=      RemBertModel.from_pretrained("google/rembert"		)
						_A       	=      torch.tensor([[3_12, 5_64_98, 3_13, 21_25, 3_13]]		)
						_A       	=      torch.tensor([[0, 0, 0, 1, 1]]		)
						with torch.no_grad():
										_A       	=      model(__UpperCamelCase     , token_type_ids=__UpperCamelCase     , output_hidden_states=__UpperCamelCase		)

						_A       	=      11_52

						_A       	=      torch.Size((1, 5, hidden_size)		)
						self.assertEqual(output["last_hidden_state"].shape     , __UpperCamelCase		)

						_A       	=      torch.tensor(
						    [
						        [
						            [0.0_7_5_4, -0.2_0_2_2, 0.1_9_0_4],
						            [-0.3_3_5_4, -0.3_6_9_2, -0.4_7_9_1],
						            [-0.2_3_1_4, -0.6_7_2_9, -0.0_7_4_9],
						            [-0.0_3_9_6, -0.3_1_0_5, -0.4_2_3_4],
						            [-0.1_5_7_1, -0.0_5_2_5, 0.5_3_5_3],
						        ]
						    ]		)

						# Running on the original tf implementation gives slightly different results here.
						# Not clear why this variations is present
						# TODO: Find reason for discrepancy
						# expected_original_implementation = [[
						#     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
						#     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
						#     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
						#     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
						#     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
						# ]]

						self.assertTrue(torch.allclose(output["last_hidden_state"][:, :, :3]     , __UpperCamelCase     , atol=1E-4		)		)

import argparse
import os
from os.path import isfile

import torch

from transformers import GPTaConfig





def        lowerCAmelCase       (				snake_case__   :      List[str]				,					snake_case__   :      List[str]				,					snake_case__   :      Optional[int]=0 )->			Union[str, Any]:
    # Format the message.
    if name is None:
        A_								=					None
    else:
        A_								=					"." * max(0				,					spaces - 2 ) + "# {:" + str(50 - spaces ) + "s}"
        A_								=					fmt.format(snake_case__ )

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





def        lowerCAmelCase       (				snake_case__   :      Any				,					snake_case__   :      List[Any]				,					snake_case__   :      str				,					snake_case__   :      str )->			Union[str, Any]:
    # 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.
    A_								=					param.size()
    # other versions store [num_heads * num_splits * hidden_size, :]
    A_								=					(num_heads, num_splits, hidden_size) + input_shape[1:]
    A_								=					param.view(*snake_case__ )
    A_								=					param.transpose(0				,					1 ).contiguous()
    A_								=					param.view(*snake_case__ )
    return param





def        lowerCAmelCase       (				snake_case__   :      Dict				,					snake_case__   :      str )->			Optional[int]:
    A_								=					config.n_positions
    A_								=					config.n_layer
    A_								=					config.vocab_size
    A_								=					config.n_head
    A_								=					config.n_embd // config.n_head

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

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

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

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

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

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

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

    return sd_hf





def        lowerCAmelCase       (				snake_case__   :      Optional[int]				,					snake_case__   :      List[Any] )->			Dict:
    A_								=					64000
    A_								=					config_megatron["encoder_seq_length"]
    A_								=					config_megatron["hidden_size"]
    A_								=					config_megatron["num_layers"]
    A_								=					config_megatron["num_attention_heads"]
    A_								=					config_megatron["ffn_hidden_size"]
    A_								=					"gelu"
    A_								=					0.1
    A_								=					0.1
    A_								=					0.1
    A_								=					config_megatron["layernorm_epsilon"]  # 1e-5
    A_								=					config_megatron["init_method_std"]  # 0.02
    A_								=					config_megatron["apply_query_key_layer_scaling"]  # True
    A_								=					True
    A_								=					True

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

    return config_hf





def        lowerCAmelCase       (				snake_case__   :      Tuple )->			Dict:
    print(snake_case__ )

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

    # Load the model.
    A_								=					torch.load(snake_case__				,					map_location="cpu" )

    # Load the config.
    A_								=					checkpoint["hyper_parameters"]["cfg"]
    A_								=					GPTaConfig()
    A_								=					copy_config(config_hf=snake_case__				,					config_megatron=snake_case__ )
    A_								=					["GPT2LMHeadModel"]

    A_								=					checkpoint["state_dict"]

    # Convert.
    print("Converting" )
    A_								=					convert_megatron_checkpoint(snake_case__				,					snake_case__ )

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

    A_								=					"GPTSw3Tokenizer"

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

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


if __name__ == "__main__":
     __magic_name__  :  Union[str, Any]     			=   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')
     __magic_name__  :  str     			=   parser.parse_args()
     main(_args)






