





'''simple docstring'''




import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_A:					int							=				get_tests_dir("""fixtures/test_sentencepiece.model""")

if is_sentencepiece_available():
    import sentencepiece as sp


_A:					Tuple							=				5
_A:					Optional[int]							=				10






@require_sentencepiece
@require_tokenizers
class UpperCAmelCase							(		UpperCAmelCase_					,			unittest.TestCase							):
    _A :			Dict						   =       SpeechaTextTokenizer
    _A :			str						   =       False
    _A :			Union[str, Any]						   =       True



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

     __UpperCAmelCase					=   sp.SentencePieceProcessor()
     spm_model.Load(__A					)
     __UpperCAmelCase					=   ['<s>', '<pad>', '</s>', '<unk>']

     vocab += [spm_model.IdToPiece(id_					) for id_ in range(len(__A					)					)]
     __UpperCAmelCase					=   dict(zip(__A     ,			range(len(__A					)					)					)					)

     __UpperCAmelCase					=   Path(self.tmpdirname					)
     save_json(__A     ,			save_dir / VOCAB_FILES_NAMES['vocab_file']					)
     if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
      copyfile(__A     ,			save_dir / VOCAB_FILES_NAMES['spm_file']					)

     __UpperCAmelCase					=   SpeechaTextTokenizer.from_pretrained(self.tmpdirname					)
     tokenizer.save_pretrained(self.tmpdirname					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   '<pad>'
     __UpperCAmelCase					=   1

     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					):
     __UpperCAmelCase					=   list(self.get_tokenizer().get_vocab().keys()					)

     self.assertEqual(vocab_keys[0]     ,			'<s>'					)
     self.assertEqual(vocab_keys[1]     ,			'<pad>'					)
     self.assertEqual(vocab_keys[-1]     ,			'j'					)
     self.assertEqual(len(__A					)     ,			1_001					)



    def 					__lowerCamelCase						(  self					):
     self.assertEqual(self.get_tokenizer().vocab_size     ,			1_001					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   SpeechaTextTokenizer.from_pretrained(self.tmpdirname					)

     __UpperCAmelCase					=   tokenizer.tokenize('This is a test'					)
     self.assertListEqual(__A     ,			['▁This', '▁is', '▁a', '▁t', 'est']					)

     self.assertListEqual(
         tokenizer.convert_tokens_to_ids(__A					)     ,			[289, 50, 14, 174, 386]     ,			)

     __UpperCAmelCase					=   tokenizer.tokenize('I was born in 92000, and this is falsé.'					)
     self.assertListEqual(
         __A     ,			[SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '9', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', 'é', '.']     ,			)
     __UpperCAmelCase					=   tokenizer.convert_tokens_to_ids(__A					)
     self.assertListEqual(__A     ,			[12, 25, 88, 59, 28, 23, 11, 4, 606, 351, 351, 351, 7, 16, 70, 50, 76, 84, 10, 4, 8]					)

     __UpperCAmelCase					=   tokenizer.convert_ids_to_tokens(__A					)
     self.assertListEqual(
         __A     ,			[SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '<unk>', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', '<unk>', '.']     ,			)




    @slow
    def 					__lowerCamelCase						(  self					):
     # fmt: off
     __UpperCAmelCase					=   {'input_ids': [[3_791, 797, 31, 11, 64, 797, 31, 2_429, 433, 12, 1_176, 12, 20, 786, 915, 142, 2_413, 240, 37, 3_238, 797, 31, 11, 35, 93, 915, 142, 2_413, 240, 37, 5_540, 567, 1_276, 93, 37, 610, 40, 62, 455, 657, 1_042, 123, 780, 177, 37, 309, 241, 1_298, 514, 20, 292, 2_737, 114, 2_469, 241, 85, 64, 302, 548, 528, 423, 4, 509, 406, 423, 37, 601, 4, 777, 302, 548, 528, 423, 284, 4, 3_388, 511, 459, 4, 3_555, 40, 321, 302, 705, 4, 3_388, 511, 583, 326, 5, 5, 5, 62, 3_310, 560, 177, 2_680, 217, 1_508, 32, 31, 853, 418, 64, 583, 511, 1_605, 62, 35, 93, 560, 177, 2_680, 217, 1_508, 1_521, 64, 583, 511, 519, 62, 20, 1_515, 764, 20, 149, 261, 5_625, 7_972, 20, 5_540, 567, 1_276, 93, 3_925, 1_675, 11, 15, 802, 7_972, 576, 217, 1_508, 11, 35, 93, 1_253, 2_441, 15, 289, 652, 31, 416, 321, 3_842, 115, 40, 911, 8, 476, 619, 4, 380, 142, 423, 335, 240, 35, 93, 264, 8, 11, 335, 569, 420, 163, 5, 2], [260, 548, 528, 423, 20, 451, 20, 2_681, 1_153, 3_434, 20, 5_540, 37, 567, 126, 1_253, 2_441, 3_376, 449, 210, 431, 1_563, 177, 767, 5_540, 11, 1_203, 472, 11, 2_953, 685, 285, 364, 706, 1_153, 20, 6_799, 20, 2_869, 20, 4_464, 126, 40, 2_429, 20, 1_040, 866, 2_664, 418, 20, 318, 20, 1_726, 186, 20, 265, 522, 35, 93, 2_191, 4_634, 20, 1_040, 12, 6_799, 15, 228, 2_356, 142, 31, 11, 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, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2_575, 2_666, 684, 1_582, 1_176, 12, 627, 149, 619, 20, 4_902, 563, 11, 20, 149, 261, 3_420, 2_356, 174, 142, 4_714, 131, 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, 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]], 'attention_mask': [[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, 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, 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, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [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, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
     # fmt: on

     self.tokenizer_integration_test_util(
         expected_encoding=__A     ,			model_name='facebook/s2t-small-mustc-en-de-st'     ,			revision='a14f04cf0776c02f62a8cb800cf7909e15ea23ad'     ,			)






@require_sentencepiece
class UpperCAmelCase							(		unittest.TestCase							):
    _A :			Dict						   =       """valhalla/s2t_mustc_multilinguial_medium"""

    _A :			Any						   =       """C'est trop cool"""
    _A :			List[Any]						   =       """Esto es genial"""



    @classmethod
    def 					__lowerCamelCase						(  cls					):
     __UpperCAmelCase					=   SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name					)
     return cls



    def 					__lowerCamelCase						(  self					):
     self.assertEqual(self.tokenizer.lang_code_to_id['pt']     ,			4					)
     self.assertEqual(self.tokenizer.lang_code_to_id['ru']     ,			6					)
     self.assertEqual(self.tokenizer.lang_code_to_id['it']     ,			9					)
     self.assertEqual(self.tokenizer.lang_code_to_id['de']     ,			11					)



    def 					__lowerCamelCase						(  self					):
     self.assertEqual(self.tokenizer.vocab_size     ,			10_000					)



    def 					__lowerCamelCase						(  self					):
     self.assertIn(__A     ,			self.tokenizer.all_special_ids					)
     __UpperCAmelCase					=   [ES_CODE, 4, 1_601, 47, 7_647, 2]
     __UpperCAmelCase					=   self.tokenizer.decode(__A     ,			skip_special_tokens=__A					)
     __UpperCAmelCase					=   self.tokenizer.decode(generated_ids[1:]     ,			skip_special_tokens=__A					)
     self.assertEqual(__A     ,			__A					)
     self.assertNotIn(self.tokenizer.eos_token     ,			__A					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   'fr'
     __UpperCAmelCase					=   self.tokenizer(self.french_text					).input_ids
     self.assertEqual(encoded[0]     ,			__A					)
     self.assertEqual(encoded[-1]     ,			self.tokenizer.eos_token_id					)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   'fr'
     self.assertListEqual(self.tokenizer.prefix_tokens     ,			[FR_CODE]					)

     __UpperCAmelCase					=   'es'
     self.assertListEqual(self.tokenizer.prefix_tokens     ,			[ES_CODE]					)












'''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					)





