






'''simple docstring'''
import argparse
from pathlib import Path

from transformers import AutoConfig, AutoTokenizer, RagConfig, RagSequenceForGeneration, RagTokenForGeneration





def  _a(  UpperCamelCase__				:						List[str],	UpperCamelCase__				:						str,	UpperCamelCase__				:						str,	UpperCamelCase__				:						Path,	UpperCamelCase__				:						str = None,	UpperCamelCase__				:						str = None,	UpperCamelCase__				:						str = None,	):




    '''simple docstring'''






    if config_name_or_path is None:
        SCREAMING_SNAKE_CASE__     :				List[Any]									='''facebook/rag-token-base''' if model_type == '''rag_token''' else '''facebook/rag-sequence-base'''

    if generator_tokenizer_name_or_path is None:
        SCREAMING_SNAKE_CASE__     :				int									=generator_name_or_path

    if question_encoder_tokenizer_name_or_path is None:
        SCREAMING_SNAKE_CASE__     :				List[str]									=question_encoder_name_or_path

    SCREAMING_SNAKE_CASE__     :				int									=RagTokenForGeneration if model_type == '''rag_token''' else RagSequenceForGeneration

    # Save model.
    SCREAMING_SNAKE_CASE__     :				Dict									=RagConfig.from_pretrained(UpperCamelCase__							)
    SCREAMING_SNAKE_CASE__     :				List[str]									=AutoConfig.from_pretrained(UpperCamelCase__							)
    SCREAMING_SNAKE_CASE__     :				List[str]									=AutoConfig.from_pretrained(UpperCamelCase__							)

    SCREAMING_SNAKE_CASE__     :				Optional[int]									=gen_config
    SCREAMING_SNAKE_CASE__     :				int									=question_encoder_config

    SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model_class.from_pretrained_question_encoder_generator(
        UpperCamelCase__,	UpperCamelCase__,	config=UpperCamelCase__							)
    rag_model.save_pretrained(UpperCamelCase__							)

    # Sanity check.
    model_class.from_pretrained(UpperCamelCase__							)

    # Save tokenizers.
    SCREAMING_SNAKE_CASE__     :				List[str]									=AutoTokenizer.from_pretrained(UpperCamelCase__							)
    gen_tokenizer.save_pretrained(dest_dir / '''generator_tokenizer/'''							)
    SCREAMING_SNAKE_CASE__     :				int									=AutoTokenizer.from_pretrained(UpperCamelCase__							)
    question_encoder_tokenizer.save_pretrained(dest_dir / '''question_encoder_tokenizer/'''							)


if __name__ == "__main__":
  a_					= argparse.ArgumentParser()
  parser.add_argument(
      '--model_type',
      choices=['rag_sequence', 'rag_token'],
      required=True,
      type=str,
      help='RAG model type: rag_sequence, rag_token',
  )
  parser.add_argument('--dest', type=str, required=True, help='Path to the output checkpoint directory.')
  parser.add_argument('--generator_name_or_path', type=str, required=True, help='Generator model identifier')
  parser.add_argument(
      '--question_encoder_name_or_path', type=str, required=True, help='Question encoder model identifier'
  )

  parser.add_argument(
      '--generator_tokenizer_name_or_path',
      type=str,
      help='Generator tokenizer identifier, if not specified, resolves to ``generator_name_or_path``',
  )
  parser.add_argument(
      '--question_encoder_tokenizer_name_or_path',
      type=str,
      help='Question encoder tokenizer identifier, if not specified, resolves to ``question_encoder_name_or_path``',
  )
  parser.add_argument(
      '--config_name_or_path',
      type=str,
      help=(
          'Identifier of the model config to use, if not provided, resolves to a base config for a given'
          ' ``model_type``'
      ),
  )

  a_					= parser.parse_args()

  a_					= Path(args.dest)
  dest_dir.mkdir(exist_ok=True)

  consolidate(
      args.model_type,
      args.generator_name_or_path,
      args.question_encoder_name_or_path,
      dest_dir,
      args.config_name_or_path,
      args.generator_tokenizer_name_or_path,
      args.question_encoder_tokenizer_name_or_path,
  )







'''simple docstring'''
import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )






class      __SCREAMING_SNAKE_CASE      :




 def __init__(      self	:							Optional[int]					,      __lowercase	:							Optional[Any]					,      __lowercase	:							List[str]=13					,      __lowercase	:							int=7					,      __lowercase	:							Tuple=True					,      __lowercase	:							Optional[Any]=True					,      __lowercase	:							List[Any]=True					,      __lowercase	:							str=True					,      __lowercase	:							Tuple=99					,      __lowercase	:							Union[str, Any]=64					,      __lowercase	:							Tuple=32					,      __lowercase	:							Optional[Any]=5					,      __lowercase	:							Tuple=4					,      __lowercase	:							Optional[int]=37					,      __lowercase	:							int="gelu"					,      __lowercase	:							Union[str, Any]=0.1					,      __lowercase	:							List[Any]=0.1					,      __lowercase	:							List[Any]=5_12					,      __lowercase	:							int=16					,      __lowercase	:							Optional[int]=2					,      __lowercase	:							Tuple=0.02					,      __lowercase	:							List[str]=3					,      __lowercase	:							List[str]=4					,      __lowercase	:							List[str]=None					,      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=parent
       SCREAMING_SNAKE_CASE__     :				Any									=batch_size
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=seq_length
       SCREAMING_SNAKE_CASE__     :				Dict									=is_training
       SCREAMING_SNAKE_CASE__     :				List[Any]									=use_input_mask
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=use_token_type_ids
       SCREAMING_SNAKE_CASE__     :				List[Any]									=use_labels
       SCREAMING_SNAKE_CASE__     :				int									=vocab_size
       SCREAMING_SNAKE_CASE__     :				str									=hidden_size
       SCREAMING_SNAKE_CASE__     :				Any									=embedding_size
       SCREAMING_SNAKE_CASE__     :				Tuple									=num_hidden_layers
       SCREAMING_SNAKE_CASE__     :				str									=num_attention_heads
       SCREAMING_SNAKE_CASE__     :				Tuple									=intermediate_size
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=hidden_act
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=hidden_dropout_prob
       SCREAMING_SNAKE_CASE__     :				str									=attention_probs_dropout_prob
       SCREAMING_SNAKE_CASE__     :				Any									=max_position_embeddings
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=type_vocab_size
       SCREAMING_SNAKE_CASE__     :				Any									=type_sequence_label_size
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=initializer_range
       SCREAMING_SNAKE_CASE__     :				str									=num_labels
       SCREAMING_SNAKE_CASE__     :				List[str]									=num_choices
       SCREAMING_SNAKE_CASE__     :				List[str]									=scope




 def      __magic_name__	(      self	:							str      )	->						Optional[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size      )

       SCREAMING_SNAKE_CASE__     :				int									=None
       if self.use_input_mask:
             SCREAMING_SNAKE_CASE__     :				List[str]									=random_attention_mask([self.batch_size, self.seq_length]      )

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

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=None
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=None
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=None
       if self.use_labels:
             SCREAMING_SNAKE_CASE__     :				Tuple									=ids_tensor([self.batch_size]					,      self.type_sequence_label_size      )
             SCREAMING_SNAKE_CASE__     :				List[Any]									=ids_tensor([self.batch_size, self.seq_length]					,      self.num_labels      )
             SCREAMING_SNAKE_CASE__     :				Any									=ids_tensor([self.batch_size]					,      self.num_choices      )

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.get_config()

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




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       return MegatronBertConfig(
           vocab_size=self.vocab_size					,      hidden_size=self.hidden_size					,      num_hidden_layers=self.num_hidden_layers					,      num_attention_heads=self.num_attention_heads					,      intermediate_size=self.intermediate_size					,      embedding_size=self.embedding_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=__lowercase					,      initializer_range=self.initializer_range					,      )




 def      __magic_name__	(      self	:							Any					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]      )	->						int:
       SCREAMING_SNAKE_CASE__     :				List[str]									=MegatronBertModel(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[Any]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model(__lowercase					,      token_type_ids=__lowercase      )
       SCREAMING_SNAKE_CASE__     :				str									=model(__lowercase      )

       self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, self.seq_length, self.hidden_size)      )
       self.parent.assertEqual(result.pooler_output.shape					,      (self.batch_size, self.hidden_size)      )




 def      __magic_name__	(      self	:							Dict					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[str]					,      __lowercase	:							str					,      __lowercase	:							Tuple					,      __lowercase	:							Any					,      __lowercase	:							Optional[int]					,      __lowercase	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				str									=MegatronBertForMaskedLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Tuple									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[Any]					,      __lowercase	:							Any					,      __lowercase	:							str					,      __lowercase	:							Any					,      __lowercase	:							str      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForCausalLM(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )




 def      __magic_name__	(      self	:							Union[str, Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							Dict					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=MegatronBertForNextSentencePrediction(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, 2)      )




 def      __magic_name__	(      self	:							List[str]					,      __lowercase	:							Optional[int]					,      __lowercase	:							str					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[Any]					,      __lowercase	:							str					,      __lowercase	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=MegatronBertForPreTraining(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      next_sentence_label=__lowercase					,      )
       self.parent.assertEqual(result.prediction_logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)      )
       self.parent.assertEqual(result.seq_relationship_logits.shape					,      (self.batch_size, 2)      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[int]					,      __lowercase	:							Union[str, Any]					,      __lowercase	:							int					,      __lowercase	:							Any					,      __lowercase	:							Any					,      __lowercase	:							List[str]      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=MegatronBertForQuestionAnswering(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Dict									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      start_positions=__lowercase					,      end_positions=__lowercase					,      )
       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      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Optional[int]					,      __lowercase	:							List[Any]					,      __lowercase	:							List[Any]					,      __lowercase	:							Optional[int]					,      __lowercase	:							int					,      __lowercase	:							List[str]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.num_labels
       SCREAMING_SNAKE_CASE__     :				Dict									=MegatronBertForSequenceClassification(__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_labels)      )




 def      __magic_name__	(      self	:							Optional[int]					,      __lowercase	:							List[Any]					,      __lowercase	:							str					,      __lowercase	:							Optional[Any]					,      __lowercase	:							Any					,      __lowercase	:							Tuple					,      __lowercase	:							Optional[Any]					,      __lowercase	:							int      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.num_labels
       SCREAMING_SNAKE_CASE__     :				int									=MegatronBertForTokenClassification(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				List[str]									=model(__lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.num_labels)      )




 def      __magic_name__	(      self	:							Dict					,      __lowercase	:							List[str]					,      __lowercase	:							Optional[int]					,      __lowercase	:							Dict					,      __lowercase	:							Any					,      __lowercase	:							List[Any]					,      __lowercase	:							int					,      __lowercase	:							Optional[int]      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				int									=self.num_choices
       SCREAMING_SNAKE_CASE__     :				List[str]									=MegatronBertForMultipleChoice(config=__lowercase      )
       model.to(__lowercase      )
       model.eval()
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=input_ids.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=token_type_ids.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Tuple									=input_mask.unsqueeze(1      ).expand(-1					,      self.num_choices					,      -1      ).contiguous()
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=model(
           __lowercase					,      attention_mask=__lowercase					,      token_type_ids=__lowercase					,      labels=__lowercase					,      )
       self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_choices)      )




 def      __magic_name__	(      self	:							str      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				Tuple									=self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	(
           SCREAMING_SNAKE_CASE__
       )		,	
       )     :				str									=config_and_inputs
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_torch
class      __SCREAMING_SNAKE_CASE      (			lowerCamelCase	,				lowerCamelCase	,				unittest.TestCase				):
 snake_case_  							=							(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 snake_case_  							=							(
     {
         """feature-extraction""": MegatronBertModel,
         """fill-mask""": MegatronBertForMaskedLM,
         """question-answering""": MegatronBertForQuestionAnswering,
         """text-classification""": MegatronBertForSequenceClassification,
         """text-generation""": MegatronBertForCausalLM,
         """token-classification""": MegatronBertForTokenClassification,
         """zero-shot""": MegatronBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 snake_case_  							=							True
 # test_resize_embeddings = False
 snake_case_  							=							False




 def      __magic_name__	(      self	:							List[Any]					,      __lowercase	:							List[str]					,      __lowercase	:							Tuple					,      __lowercase	:							Tuple=False      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=super()._prepare_for_class(__lowercase					,      __lowercase					,      return_labels=__lowercase      )

       if return_labels:
             if model_class in get_values(__lowercase      ):
                   SCREAMING_SNAKE_CASE__     :				List[str]									=torch.zeros(
                       (self.model_tester.batch_size, self.model_tester.seq_length)					,      dtype=torch.long					,      device=__lowercase      )
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=torch.zeros(
                       self.model_tester.batch_size					,      dtype=torch.long					,      device=__lowercase      )
       return inputs_dict




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=MegatronBertModelTester(self      )
       SCREAMING_SNAKE_CASE__     :				Tuple									=ConfigTester(self					,      config_class=__lowercase					,      hidden_size=37      )




 def      __magic_name__	(      self	:							str      )	->						Dict:
       self.config_tester.run_common_tests()




 def      __magic_name__	(      self	:							Tuple      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				str									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*__lowercase      )




 def      __magic_name__	(      self	:							List[str]      )	->						Any:
       SCREAMING_SNAKE_CASE__     :				int									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				Dict									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__lowercase      )




 def      __magic_name__	(      self	:							Optional[int]      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__lowercase      )




 def      __magic_name__	(      self	:							Dict      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*__lowercase      )




 def      __magic_name__	(      self	:							int      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Optional[int]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*__lowercase      )




 def      __magic_name__	(      self	:							List[Any]      )	->						List[Any]:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__lowercase      )




 def      __magic_name__	(      self	:							Union[str, Any]      )	->						Union[str, Any]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*__lowercase      )





def  _a(  UpperCamelCase__				:						List[str]							):




      '''simple docstring'''






      return torch.tensor(
          UpperCamelCase__,	dtype=torch.long,	device=UpperCamelCase__,	)


a_					= 1E-4







@require_torch
@require_sentencepiece
@require_tokenizers
class      __SCREAMING_SNAKE_CASE      (			unittest.TestCase				):




 @slow
 @unittest.skip('''Model is not available.'''      )
 def      __magic_name__	(      self	:							Any      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				Any									='''nvidia/megatron-bert-uncased-345m'''
       if "MYDIR" in os.environ:
             SCREAMING_SNAKE_CASE__     :				Optional[int]									=os.path.join(os.environ['''MYDIR''']					,      __lowercase      )
       SCREAMING_SNAKE_CASE__     :				Any									=MegatronBertModel.from_pretrained(__lowercase      )
       model.to(__lowercase      )
       model.half()
       SCREAMING_SNAKE_CASE__     :				Dict									=_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
       with torch.no_grad():
             SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=model(__lowercase      )[0]
       SCREAMING_SNAKE_CASE__     :				Dict									=torch.Size((1, 9, 10_24)      )
       self.assertEqual(output.shape					,      __lowercase      )

       SCREAMING_SNAKE_CASE__     :				str									=[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
       for ii in range(3      ):
             for jj in range(3      ):
                   SCREAMING_SNAKE_CASE__     :				List[Any]									=output[0, ii, jj]
                   SCREAMING_SNAKE_CASE__     :				Tuple									=expected[3 * ii + jj]
                   SCREAMING_SNAKE_CASE__     :				List[str]									='''ii={} jj={} a={} b={}'''.format(__lowercase					,      __lowercase					,      __lowercase					,      __lowercase      )
                   self.assertTrue(math.isclose(__lowercase					,      __lowercase					,      rel_tol=__lowercase					,      abs_tol=__lowercase      )					,      msg=__lowercase      )