





'''simple docstring'''




import argparse
from pathlib import Path

import fairseq
import torch
from fairseq.models.xmod import XMODModel as FairseqXmodModel
from packaging import version

from transformers import XmodConfig, XmodForMaskedLM, XmodForSequenceClassification
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse("""0.12.2"""):
  raise Exception("""requires fairseq >= 0.12.2""")
if version.parse(fairseq.__version__) > version.parse("""2"""):
  raise Exception("""requires fairseq < v2""")

logging.set_verbosity_info()
_A:					List[str]							=				logging.get_logger(__name__)

_A:					List[Any]							=				"""Hello, World!"""
_A:					Optional[Any]							=				"""en_XX"""







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase       ,    _lowerCAmelCase    )->     Optional[Any]:
 __UpperCAmelCase					=   Path('data_bin'    )
 __UpperCAmelCase					=   FairseqXmodModel.from_pretrained(
     model_name_or_path=str(Path(_lowerCAmelCase    ).parent    )       ,    checkpoint_file=Path(_lowerCAmelCase    ).name       ,    _name='xmod_base'       ,    arch='xmod_base'       ,    task='multilingual_masked_lm'       ,    data_name_or_path=str(_lowerCAmelCase    )       ,    bpe='sentencepiece'       ,    sentencepiece_model=str(Path(_lowerCAmelCase    ).parent / 'sentencepiece.bpe.model'    )       ,    src_dict=str(data_dir / 'dict.txt'    )       ,    )
 xmod.eval()  # disable dropout
 print(_lowerCAmelCase    )

 __UpperCAmelCase					=   xmod.model.encoder.sentence_encoder
 __UpperCAmelCase					=   XmodConfig(
     vocab_size=xmod_sent_encoder.embed_tokens.num_embeddings       ,    hidden_size=xmod.cfg.model.encoder_embed_dim       ,    num_hidden_layers=xmod.cfg.model.encoder_layers       ,    num_attention_heads=xmod.cfg.model.encoder_attention_heads       ,    intermediate_size=xmod.cfg.model.encoder_ffn_embed_dim       ,    max_position_embeddings=5_14       ,    type_vocab_size=1       ,    layer_norm_eps=1e-5       ,    pre_norm=xmod.cfg.model.encoder_normalize_before       ,    adapter_reduction_factor=getattr(xmod.cfg.model       ,    'bottleneck'       ,    2    )       ,    adapter_layer_norm=xmod.cfg.model.adapter_layer_norm       ,    adapter_reuse_layer_norm=xmod.cfg.model.adapter_reuse_layer_norm       ,    ln_before_adapter=xmod.cfg.model.ln_before_adapter       ,    languages=xmod.cfg.model.languages       ,    )
 if classification_head:
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'].out_proj.weight.shape[0]

 print('Our X-MOD config:'       ,    _lowerCAmelCase    )

 __UpperCAmelCase					=   XmodForSequenceClassification(_lowerCAmelCase    ) if classification_head else XmodForMaskedLM(_lowerCAmelCase    )
 model.eval()

 # Now let's copy all the weights.
 # Embeddings
 __UpperCAmelCase					=   xmod_sent_encoder.embed_tokens.weight
 __UpperCAmelCase					=   xmod_sent_encoder.embed_positions.weight
 __UpperCAmelCase					=   torch.zeros_like(
     model.roberta.embeddings.token_type_embeddings.weight    )  # just zero them out b/c xmod doesn't use them.

 __UpperCAmelCase					=   xmod_sent_encoder.layernorm_embedding.weight
 __UpperCAmelCase					=   xmod_sent_encoder.layernorm_embedding.bias

 for i in range(config.num_hidden_layers    ):
  # Encoder: start of layer
  __UpperCAmelCase					=   model.roberta.encoder.layer[i]
  __UpperCAmelCase					=   xmod_sent_encoder.layers[i]

  # self attention
  __UpperCAmelCase					=   layer.attention.self
  if not (
      xmod_layer.self_attn.k_proj.weight.data.shape
      == xmod_layer.self_attn.q_proj.weight.data.shape
      == xmod_layer.self_attn.v_proj.weight.data.shape
      == torch.Size((config.hidden_size, config.hidden_size)    )
  ):
   raise AssertionError('Dimensions of self-attention weights do not match.'    )

  __UpperCAmelCase					=   xmod_layer.self_attn.q_proj.weight
  __UpperCAmelCase					=   xmod_layer.self_attn.q_proj.bias
  __UpperCAmelCase					=   xmod_layer.self_attn.k_proj.weight
  __UpperCAmelCase					=   xmod_layer.self_attn.k_proj.bias
  __UpperCAmelCase					=   xmod_layer.self_attn.v_proj.weight
  __UpperCAmelCase					=   xmod_layer.self_attn.v_proj.bias

  # self-attention output
  __UpperCAmelCase					=   layer.attention.output
  if self_output.dense.weight.shape != xmod_layer.self_attn.out_proj.weight.shape:
   raise AssertionError('Dimensions of self-attention output weights do not match.'    )
  __UpperCAmelCase					=   xmod_layer.self_attn.out_proj.weight
  __UpperCAmelCase					=   xmod_layer.self_attn.out_proj.bias
  __UpperCAmelCase					=   xmod_layer.self_attn_layer_norm.weight
  __UpperCAmelCase					=   xmod_layer.self_attn_layer_norm.bias

  # intermediate
  __UpperCAmelCase					=   layer.intermediate
  if intermediate.dense.weight.shape != xmod_layer.fca.weight.shape:
   raise AssertionError('Dimensions of intermediate weights do not match.'    )
  __UpperCAmelCase					=   xmod_layer.fca.weight
  __UpperCAmelCase					=   xmod_layer.fca.bias

  # output
  __UpperCAmelCase					=   layer.output
  if bert_output.dense.weight.shape != xmod_layer.fca.weight.shape:
   raise AssertionError('Dimensions of feed-forward weights do not match.'    )
  __UpperCAmelCase					=   xmod_layer.fca.weight
  __UpperCAmelCase					=   xmod_layer.fca.bias
  __UpperCAmelCase					=   xmod_layer.final_layer_norm.weight
  __UpperCAmelCase					=   xmod_layer.final_layer_norm.bias
  if bert_output.adapter_layer_norm is not None:
   __UpperCAmelCase					=   xmod_layer.adapter_layer_norm.weight
   __UpperCAmelCase					=   xmod_layer.adapter_layer_norm.bias

  if sorted(bert_output.adapter_modules.keys()    ) != sorted(xmod_layer.adapter_modules.keys()    ):
   raise AssertionError('Lists of language adapters do not match.'    )
  for lang_code, adapter in xmod_layer.adapter_modules.items():
   __UpperCAmelCase					=   bert_output.adapter_modules[lang_code]
   __UpperCAmelCase					=   xmod_layer.adapter_modules[lang_code]
   __UpperCAmelCase					=   from_adapter.fca.weight
   __UpperCAmelCase					=   from_adapter.fca.bias
   __UpperCAmelCase					=   from_adapter.fca.weight
   __UpperCAmelCase					=   from_adapter.fca.bias

  # end of layer

 if xmod_sent_encoder.layer_norm is not None:
  __UpperCAmelCase					=   xmod_sent_encoder.layer_norm.weight
  __UpperCAmelCase					=   xmod_sent_encoder.layer_norm.bias

 if classification_head:
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'].dense.weight
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'].dense.bias
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'].out_proj.weight
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'].out_proj.bias
 else:
  # LM Head
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.dense.weight
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.dense.bias
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.layer_norm.weight
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.layer_norm.bias
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.weight
  __UpperCAmelCase					=   xmod.model.encoder.lm_head.bias

 # Let's check that we get the same results.
 __UpperCAmelCase					=   xmod.encode(_lowerCAmelCase    ).unsqueeze(0    )  # batch of size 1
 model.roberta.set_default_language(_lowerCAmelCase    )

 __UpperCAmelCase					=   model(_lowerCAmelCase    )[0]
 if classification_head:
  __UpperCAmelCase					=   xmod.model.classification_heads['mnli'](xmod.extract_features(_lowerCAmelCase    )    )
 else:
  __UpperCAmelCase					=   xmod.model(_lowerCAmelCase       ,    lang_id=[SAMPLE_LANGUAGE]    )[0]
 print(our_output.shape       ,    their_output.shape    )
 __UpperCAmelCase					=   torch.max(torch.abs(our_output - their_output    )    ).item()
 print(F'max_absolute_diff = {max_absolute_diff}'    )  # ~ 1e-7
 __UpperCAmelCase					=   torch.allclose(_lowerCAmelCase       ,    _lowerCAmelCase       ,    atol=1e-3    )
 print('Do both models output the same tensors?'       ,    '🔥' if success else '💩'    )
 if not success:
  raise Exception('Something went wRoNg'    )

 Path(_lowerCAmelCase    ).mkdir(parents=_lowerCAmelCase       ,    exist_ok=_lowerCAmelCase    )
 print(F'Saving model to {pytorch_dump_folder_path}'    )
 model.save_pretrained(_lowerCAmelCase    )


if __name__ == "__main__":
  _A:					Any							=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--xmod_checkpoint_path""", default=None, type=str, required=True, help="""Path the official PyTorch dump."""
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--classification_head""", action="""store_true""", help="""Whether to convert a final classification head."""
  )
  _A:					List[Any]							=				parser.parse_args()
  convert_xmod_checkpoint_to_pytorch(
      args.xmod_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
  )












'''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 UpperCAmelCase							:



    def __init__(  self     ,			__A     ,			__A=13     ,			__A=7     ,			__A=True     ,			__A=True     ,			__A=True     ,			__A=True     ,			__A=99     ,			__A=64     ,			__A=32     ,			__A=5     ,			__A=4     ,			__A=37     ,			__A="gelu"     ,			__A=0.1     ,			__A=0.1     ,			__A=512     ,			__A=16     ,			__A=2     ,			__A=0.0_2     ,			__A=3     ,			__A=4     ,			__A=None     ,			):
     __UpperCAmelCase					=   parent
     __UpperCAmelCase					=   batch_size
     __UpperCAmelCase					=   seq_length
     __UpperCAmelCase					=   is_training
     __UpperCAmelCase					=   use_input_mask
     __UpperCAmelCase					=   use_token_type_ids
     __UpperCAmelCase					=   use_labels
     __UpperCAmelCase					=   vocab_size
     __UpperCAmelCase					=   hidden_size
     __UpperCAmelCase					=   embedding_size
     __UpperCAmelCase					=   num_hidden_layers
     __UpperCAmelCase					=   num_attention_heads
     __UpperCAmelCase					=   intermediate_size
     __UpperCAmelCase					=   hidden_act
     __UpperCAmelCase					=   hidden_dropout_prob
     __UpperCAmelCase					=   attention_probs_dropout_prob
     __UpperCAmelCase					=   max_position_embeddings
     __UpperCAmelCase					=   type_vocab_size
     __UpperCAmelCase					=   type_sequence_label_size
     __UpperCAmelCase					=   initializer_range
     __UpperCAmelCase					=   num_labels
     __UpperCAmelCase					=   num_choices
     __UpperCAmelCase					=   scope



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size					)

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

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

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

     __UpperCAmelCase					=   self.get_config()

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



    def 					__lowerCamelCase						(  self					):
     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=__A     ,			initializer_range=self.initializer_range     ,			)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertModel(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			token_type_ids=__A					)
     __UpperCAmelCase					=   model(__A     ,			token_type_ids=__A					)
     __UpperCAmelCase					=   model(__A					)

     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 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertForMaskedLM(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertForCausalLM(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.vocab_size)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertForNextSentencePrediction(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(
         __A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A     ,			)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, 2)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertForPreTraining(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(
         __A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A     ,			next_sentence_label=__A     ,			)
     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 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   MegatronBertForQuestionAnswering(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(
         __A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			start_positions=__A     ,			end_positions=__A     ,			)
     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 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_labels
     __UpperCAmelCase					=   MegatronBertForSequenceClassification(__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_labels)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_labels
     __UpperCAmelCase					=   MegatronBertForTokenClassification(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   model(__A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A					)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.seq_length, self.num_labels)					)



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A					):
     __UpperCAmelCase					=   self.num_choices
     __UpperCAmelCase					=   MegatronBertForMultipleChoice(config=__A					)
     model.to(__A					)
     model.eval()
     __UpperCAmelCase					=   input_ids.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   token_type_ids.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   input_mask.unsqueeze(1					).expand(-1     ,			self.num_choices     ,			-1					).contiguous()
     __UpperCAmelCase					=   model(
         __A     ,			attention_mask=__A     ,			token_type_ids=__A     ,			labels=__A     ,			)
     self.parent.assertEqual(result.logits.shape     ,			(self.batch_size, self.num_choices)					)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.prepare_config_and_inputs()
     (
         (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    (
         __UpperCAmelCase
     )					,    
     )					=   config_and_inputs
     __UpperCAmelCase					=   {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict






@require_torch
class UpperCAmelCase							(		UpperCAmelCase_					,			UpperCAmelCase_					,			unittest.TestCase							):
    _A :			str						   =       (
        (
            MegatronBertModel,
            MegatronBertForMaskedLM,
            MegatronBertForCausalLM,
            MegatronBertForMultipleChoice,
            MegatronBertForNextSentencePrediction,
            MegatronBertForPreTraining,
            MegatronBertForQuestionAnswering,
            MegatronBertForSequenceClassification,
            MegatronBertForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    _A :			Dict						   =       (
        {
            """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 {}
    )
    _A :			List[str]						   =       True
    # test_resize_embeddings = False
    _A :			Tuple						   =       False



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A=False					):
     __UpperCAmelCase					=   super()._prepare_for_class(__A     ,			__A     ,			return_labels=__A					)

     if return_labels:
      if model_class in get_values(__A					):
       __UpperCAmelCase					=   torch.zeros(
           (self.model_tester.batch_size, self.model_tester.seq_length)     ,			dtype=torch.long     ,			device=__A					)
       __UpperCAmelCase					=   torch.zeros(
           self.model_tester.batch_size     ,			dtype=torch.long     ,			device=__A					)
     return inputs_dict



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   MegatronBertModelTester(self					)
     __UpperCAmelCase					=   ConfigTester(self     ,			config_class=__A     ,			hidden_size=37					)



    def 					__lowerCamelCase						(  self					):
     self.config_tester.run_common_tests()



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A					)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A					)







def _lowerCAmelCase     (      _lowerCAmelCase    )->     Dict:
 return torch.tensor(
     _lowerCAmelCase       ,    dtype=torch.long       ,    device=_lowerCAmelCase       ,    )


_A:					Any							=				1E-4






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



    @slow
    @unittest.skip('Model is not available.'					)
    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   'nvidia/megatron-bert-uncased-345m'
     if "MYDIR" in os.environ:
      __UpperCAmelCase					=   os.path.join(os.environ['MYDIR']     ,			__A					)
     __UpperCAmelCase					=   MegatronBertModel.from_pretrained(__A					)
     model.to(__A					)
     model.half()
     __UpperCAmelCase					=   _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]					)
     with torch.no_grad():
      __UpperCAmelCase					=   model(__A					)[0]
     __UpperCAmelCase					=   torch.Size((1, 9, 1_024)					)
     self.assertEqual(output.shape     ,			__A					)

     __UpperCAmelCase					=   [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
     for ii in range(3					):
      for jj in range(3					):
       __UpperCAmelCase					=   output[0, ii, jj]
       __UpperCAmelCase					=   expected[3 * ii + jj]
       __UpperCAmelCase					=   'ii={} jj={} a={} b={}'.format(__A     ,			__A     ,			__A     ,			__A					)
       self.assertTrue(math.isclose(__A     ,			__A     ,			rel_tol=__A     ,			abs_tol=__A					)     ,			msg=__A					)





