

'''simple docstring'''





import json
import os
import re
import shutil
import tempfile
import unittest
from typing import Tuple

from transformers import AddedToken, BatchEncoding, PerceiverTokenizer
from transformers.utils import cached_property, is_tf_available, is_torch_available

from ...test_tokenization_common import TokenizerTesterMixin


if is_torch_available():
							lowerCAmelCase__								=		"""pt"""
elif is_tf_available():
							lowerCAmelCase__								=		"""tf"""
else:
							lowerCAmelCase__								=		"""jax"""







class 	lowercase    (				_a    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  PerceiverTokenizer
		a__					       =  False


		def 		A__					(				self):
			super().setUp()
			_UpperCamelCase   :		Optional[Any]			=						PerceiverTokenizer()
			tokenizer.save_pretrained(self.tmpdirname)


		@cached_property
		def 		A__					(				self):
			return PerceiverTokenizer.from_pretrained('deepmind/language-perceiver')


		def 		A__					(				self	,       **__snake_case):
			return self.tokenizer_class.from_pretrained(self.tmpdirname	,       **_A)


		def 		A__					(				self	,       __snake_case	,       __snake_case=False	,       __snake_case=20	,       __snake_case=5):
			# XXX The default common tokenizer tests assume that every ID is decodable on its own.
			# This assumption is invalid for Perceiver because single bytes might not be
			# valid utf-8 (byte 128 for instance).
			# Here we're overriding the smallest possible method to provide
			# a clean sequence without making the same assumption.

			_UpperCamelCase   :		List[Any]			=						[]
			for i in range(len(_A)):
				try:
					_UpperCamelCase   :		Union[str, Any]			=						tokenizer.decode([i]	,       clean_up_tokenization_spaces=_A)
				except UnicodeDecodeError:
					pass
				toks.append((i, tok))

			_UpperCamelCase   :		List[str]			=						list(filter(lambda __snake_case: re.match(r'^[ a-zA-Z]+$'	,       t[1])	,       _A))
			_UpperCamelCase   :		Optional[Any]			=						list(filter(lambda __snake_case: [t[0]] == tokenizer.encode(t[1]	,       add_special_tokens=_A)	,       _A))
			if max_length is not None and len(_A) > max_length:
				_UpperCamelCase   :		Any			=						toks[:max_length]
			if min_length is not None and len(_A) < min_length and len(_A) > 0:
				while len(_A) < min_length:
					_UpperCamelCase   :		List[Any]			=						toks + toks
        # toks_str = [t[1] for t in toks]
			_UpperCamelCase   :		Union[str, Any]			=						[t[0] for t in toks]

			# Ensure consistency
			_UpperCamelCase   :		List[str]			=						tokenizer.decode(_A	,       clean_up_tokenization_spaces=_A)
			if " " not in output_txt and len(_A) > 1:
				_UpperCamelCase   :		int			=						(
				    tokenizer.decode([toks_ids[0]]	,       clean_up_tokenization_spaces=_A)
				    + ' '
				    + tokenizer.decode(toks_ids[1:]	,       clean_up_tokenization_spaces=_A)
				)
			if with_prefix_space:
				_UpperCamelCase   :		Dict			=						' ' + output_txt
			_UpperCamelCase   :		str			=						tokenizer.encode(_A	,       add_special_tokens=_A)
			return output_txt, output_ids


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						self.perceiver_tokenizer
			_UpperCamelCase   :		Tuple			=						'Unicode €.'
			_UpperCamelCase   :		Dict			=						tokenizer(_A)
			_UpperCamelCase   :		Union[str, Any]			=						[4, 91, 1_16, 1_11, 1_05, 1_17, 1_06, 1_07, 38, 2_32, 1_36, 1_78, 52, 5]
			self.assertEqual(encoded['input_ids']	,       _A)

			# decoding
			_UpperCamelCase   :		int			=						tokenizer.decode(_A)
			self.assertEqual(_A	,       '[CLS]Unicode €.[SEP]')

			_UpperCamelCase   :		str			=						tokenizer('e è é ê ë')
			_UpperCamelCase   :		List[str]			=						[4, 1_07, 38, 2_01, 1_74, 38, 2_01, 1_75, 38, 2_01, 1_76, 38, 2_01, 1_77, 5]
			self.assertEqual(encoded['input_ids']	,       _A)
			# decoding
			_UpperCamelCase   :		int			=						tokenizer.decode(_A)
			self.assertEqual(_A	,       '[CLS]e è é ê ë[SEP]')

			# encode/decode, but with `encode` instead of `__call__`
			self.assertEqual(tokenizer.decode(tokenizer.encode('e è é ê ë'))	,       '[CLS]e è é ê ë[SEP]')


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.perceiver_tokenizer
			_UpperCamelCase   :		str			=						['A long paragraph for summarization.', 'Another paragraph for summarization.']
			# fmt: off
			_UpperCamelCase   :		Optional[Any]			=						[4, 71, 38, 1_14, 1_17, 1_16, 1_09, 38, 1_18, 1_03, 1_20, 1_03, 1_09, 1_20, 1_03, 1_18, 1_10, 38, 1_08, 1_17, 1_20, 38, 1_21, 1_23, 1_15, 1_15, 1_03, 1_20, 1_11, 1_28, 1_03, 1_22, 1_11, 1_17, 1_16, 52, 5, 0]
			# fmt: on
			_UpperCamelCase   :		List[Any]			=						tokenizer(_A	,       padding=_A	,       return_tensors=_A)
			self.assertIsInstance(_A	,       _A)

			if FRAMEWORK != "jax":
				_UpperCamelCase   :		Optional[int]			=						list(batch.input_ids.numpy()[0])
			else:
				_UpperCamelCase   :		int			=						list(batch.input_ids.tolist()[0])

			self.assertListEqual(_A	,       _A)

			self.assertEqual((2, 38)	,       batch.input_ids.shape)
			self.assertEqual((2, 38)	,       batch.attention_mask.shape)


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						self.perceiver_tokenizer
			_UpperCamelCase   :		Any			=						['A long paragraph for summarization.', 'Another paragraph for summarization.']
			_UpperCamelCase   :		List[str]			=						tokenizer(_A	,       padding=_A	,       return_tensors=_A)
			# check if input_ids are returned and no decoder_input_ids
			self.assertIn('input_ids'	,       _A)
			self.assertIn('attention_mask'	,       _A)
			self.assertNotIn('decoder_input_ids'	,       _A)
			self.assertNotIn('decoder_attention_mask'	,       _A)


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						self.perceiver_tokenizer
			_UpperCamelCase   :		List[Any]			=						[
			    'Summary of the text.',
			    'Another summary.',
			]
			_UpperCamelCase   :		int			=						tokenizer(
			    text_target=_A	,       max_length=32	,       padding='max_length'	,       truncation=_A	,       return_tensors=_A)
			self.assertEqual(32	,       targets['input_ids'].shape[1])


		def 		A__					(				self):
			# safety check on max_len default value so we are sure the test works
			_UpperCamelCase   :		Optional[Any]			=						self.get_tokenizers()
			for tokenizer in tokenizers:
				with self.subTest(f'''{tokenizer.__class__.__name__}'''):
					self.assertNotEqual(tokenizer.model_max_length	,       42)

        # Now let's start the test
			_UpperCamelCase   :		List[str]			=						self.get_tokenizers()
			for tokenizer in tokenizers:
				with self.subTest(f'''{tokenizer.__class__.__name__}'''):
					# Isolate this from the other tests because we save additional tokens/etc
					_UpperCamelCase   :		Union[str, Any]			=						tempfile.mkdtemp()

					_UpperCamelCase   :		Optional[int]			=						' He is very happy, UNwant\u00E9d,running'
					_UpperCamelCase   :		List[Any]			=						tokenizer.encode(_A	,       add_special_tokens=_A)
					tokenizer.save_pretrained(_A)

					_UpperCamelCase   :		List[str]			=						tokenizer.__class__.from_pretrained(_A)
					_UpperCamelCase   :		Tuple			=						after_tokenizer.encode(_A	,       add_special_tokens=_A)
					self.assertListEqual(_A	,       _A)

					shutil.rmtree(_A)

			_UpperCamelCase   :		List[str]			=						self.get_tokenizers(model_max_length=42)
			for tokenizer in tokenizers:
				with self.subTest(f'''{tokenizer.__class__.__name__}'''):
					# Isolate this from the other tests because we save additional tokens/etc
					_UpperCamelCase   :		List[Any]			=						tempfile.mkdtemp()

					_UpperCamelCase   :		List[str]			=						' He is very happy, UNwant\u00E9d,running'
					tokenizer.add_tokens(['bim', 'bambam'])
					_UpperCamelCase   :		List[Any]			=						tokenizer.additional_special_tokens
					additional_special_tokens.append('new_additional_special_token')
					tokenizer.add_special_tokens({'additional_special_tokens': additional_special_tokens})
					_UpperCamelCase   :		Optional[Any]			=						tokenizer.encode(_A	,       add_special_tokens=_A)
					tokenizer.save_pretrained(_A)

					_UpperCamelCase   :		str			=						tokenizer.__class__.from_pretrained(_A)
					_UpperCamelCase   :		Optional[Any]			=						after_tokenizer.encode(_A	,       add_special_tokens=_A)
					self.assertListEqual(_A	,       _A)
					self.assertIn('new_additional_special_token'	,       after_tokenizer.additional_special_tokens)
					self.assertEqual(after_tokenizer.model_max_length	,       42)

					_UpperCamelCase   :		List[str]			=						tokenizer.__class__.from_pretrained(_A	,       model_max_length=43)
					self.assertEqual(tokenizer.model_max_length	,       43)

					shutil.rmtree(_A)


		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						[]
			if self.test_slow_tokenizer:
				tokenizer_list.append((self.tokenizer_class, self.get_tokenizer()))

			if self.test_rust_tokenizer:
				tokenizer_list.append((self.rust_tokenizer_class, self.get_rust_tokenizer()))

			for tokenizer_class, tokenizer_utils in tokenizer_list:
				with tempfile.TemporaryDirectory() as tmp_dir:
					tokenizer_utils.save_pretrained(_A)

					with open(os.path.join(_A	,       'special_tokens_map.json')	,       encoding='utf-8') as json_file:
						_UpperCamelCase   :		Any			=						json.load(_A)

					with open(os.path.join(_A	,       'tokenizer_config.json')	,       encoding='utf-8') as json_file:
						_UpperCamelCase   :		str			=						json.load(_A)

					_UpperCamelCase   :		List[str]			=						[f'''<extra_id_{i}>''' for i in range(1_25)]

					_UpperCamelCase   :		Union[str, Any]			=						added_tokens_extra_ids + [
					    'an_additional_special_token'
					]
					_UpperCamelCase   :		Any			=						added_tokens_extra_ids + [
					    'an_additional_special_token'
					]

					with open(os.path.join(_A	,       'special_tokens_map.json')	,       'w'	,       encoding='utf-8') as outfile:
						json.dump(_A	,       _A)
					with open(os.path.join(_A	,       'tokenizer_config.json')	,       'w'	,       encoding='utf-8') as outfile:
						json.dump(_A	,       _A)

					# the following checks allow us to verify that our test works as expected, i.e. that the tokenizer takes
					# into account the new value of additional_special_tokens given in the "tokenizer_config.json" and
					# "special_tokens_map.json" files
					_UpperCamelCase   :		int			=						tokenizer_class.from_pretrained(
					    _A	,       )
					self.assertIn(
					    'an_additional_special_token'	,       tokenizer_without_change_in_init.additional_special_tokens)
					self.assertEqual(
					    ['an_additional_special_token']	,       tokenizer_without_change_in_init.convert_ids_to_tokens(
					        tokenizer_without_change_in_init.convert_tokens_to_ids(['an_additional_special_token']))	,       )

					# Now we test that we can change the value of additional_special_tokens in the from_pretrained
					_UpperCamelCase   :		Dict			=						added_tokens_extra_ids + [AddedToken('a_new_additional_special_token'	,       lstrip=_A)]
					_UpperCamelCase   :		Optional[int]			=						tokenizer_class.from_pretrained(
					    _A	,       additional_special_tokens=_A	,       )

					self.assertIn('a_new_additional_special_token'	,       tokenizer.additional_special_tokens)
					self.assertEqual(
					    ['a_new_additional_special_token']	,       tokenizer.convert_ids_to_tokens(
					        tokenizer.convert_tokens_to_ids(['a_new_additional_special_token']))	,       )


		def 		A__					(				self):
			_UpperCamelCase   :		int			=						self.perceiver_tokenizer
			self.assertEqual(tokenizer.decode([1_78])	,       '�')


		def 		A__					(				self):
			pass


		def 		A__					(				self):
			pass


		def 		A__					(				self):
			pass


		def 		A__					(				self):
			pass






		def 		A__					(				self):
			# The default common tokenizer tests uses invalid tokens for Perceiver that can only accept one-character
			# strings and special added tokens as tokens
			_UpperCamelCase   :		List[str]			=						self.get_tokenizers(fast=_A	,       do_lower_case=_A)
			for tokenizer in tokenizers:
				with self.subTest(f'''{tokenizer.__class__.__name__}'''):
					_UpperCamelCase   :		Union[str, Any]			=						['[CLS]', 't', 'h', 'i', 's', ' ', 'i', 's', ' ', 'a', ' ', 't', 'e', 's', 't', '[SEP]']
					_UpperCamelCase   :		List[Any]			=						tokenizer.convert_tokens_to_string(_A)

					self.assertIsInstance(_A	,       _A)







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 	lowercase    :


		"""simple docstring"""




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


		def 		A__					(				self):
			_UpperCamelCase   :		List[str]			=						ids_tensor([self.batch_size, self.seq_length]	,       self.vocab_size)

			_UpperCamelCase   :		Any			=						None
			if self.use_input_mask:
				_UpperCamelCase   :		int			=						random_attention_mask([self.batch_size, self.seq_length])

			_UpperCamelCase   :		Optional[Any]			=						None
			if self.use_token_type_ids:
				_UpperCamelCase   :		Dict			=						ids_tensor([self.batch_size, self.seq_length]	,       self.type_vocab_size)

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

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

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


		def 		A__					(				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=__snake_case	,       initializer_range=self.initializer_range	,       )


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		List[str]			=						MegatronBertModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[int]			=						model(__snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case)
			_UpperCamelCase   :		Dict			=						model(__snake_case	,       token_type_ids=__snake_case)
			_UpperCamelCase   :		Optional[Any]			=						model(__snake_case)

			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 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		int			=						MegatronBertForMaskedLM(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Dict			=						model(__snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case	,       labels=__snake_case)
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.seq_length, self.vocab_size))


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		str			=						MegatronBertForCausalLM(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[int]			=						model(__snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case	,       labels=__snake_case)
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.seq_length, self.vocab_size))


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Tuple			=						MegatronBertForNextSentencePrediction(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[Any]			=						model(
			    __snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case	,       labels=__snake_case	,       )
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, 2))


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


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		Any			=						self.num_labels
			_UpperCamelCase   :		Optional[int]			=						MegatronBertForTokenClassification(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Tuple			=						model(__snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case	,       labels=__snake_case)
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.seq_length, self.num_labels))


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case):
			_UpperCamelCase   :		List[str]			=						self.num_choices
			_UpperCamelCase   :		Optional[int]			=						MegatronBertForMultipleChoice(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[Any]			=						input_ids.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		List[Any]			=						token_type_ids.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		Optional[Any]			=						input_mask.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		Union[str, Any]			=						model(
			    __snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case	,       labels=__snake_case	,       )
			self.parent.assertEqual(result.logits.shape	,       (self.batch_size, self.num_choices))






		def 		A__					(				self):
			_UpperCamelCase   :		Union[str, Any]			=						self.prepare_config_and_inputs()
			(
			    (
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		
			)   :		Optional[int]			=						config_and_inputs
			_UpperCamelCase   :		int			=						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
			return config, inputs_dict







@require_torch
class 	lowercase    (				_lowercase    ,							_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (
		    (
		        MegatronBertModel,
		        MegatronBertForMaskedLM,
		        MegatronBertForCausalLM,
		        MegatronBertForMultipleChoice,
		        MegatronBertForNextSentencePrediction,
		        MegatronBertForPreTraining,
		        MegatronBertForQuestionAnswering,
		        MegatronBertForSequenceClassification,
		        MegatronBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		a__					       =  (
		    {
		        "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__					       =  True
		# test_resize_embeddings = False
		a__					       =  False


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case=False):
			_UpperCamelCase   :		str			=						super()._prepare_for_class(__snake_case	,       __snake_case	,       return_labels=__snake_case)

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


		def 		A__					(				self):
			_UpperCamelCase   :		Any			=						MegatronBertModelTester(self)
			_UpperCamelCase   :		int			=						ConfigTester(self	,       config_class=__snake_case	,       hidden_size=37)


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


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_model(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		int			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		str			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_pretraining(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		List[Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_question_answering(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__snake_case)






		def 		A__					(				self):
			_UpperCamelCase   :		Any			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_token_classification(*__snake_case)






def 					lowerCamelCase_     (							UpperCAmelCase_							:     str	)						->				Optional[Any]:




	'''simple docstring'''




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


lowerCAmelCase__								=		1e-4







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


		"""simple docstring"""




		@slow
		@unittest.skip('Model is not available.')
		def 		A__					(				self):
			_UpperCamelCase   :		int			=						'nvidia/megatron-bert-uncased-345m'
			if "MYDIR" in os.environ:
				_UpperCamelCase   :		int			=						os.path.join(os.environ['MYDIR']	,       __snake_case)
			_UpperCamelCase   :		Optional[int]			=						MegatronBertModel.from_pretrained(__snake_case)
			model.to(__snake_case)
			model.half()
			_UpperCamelCase   :		Optional[Any]			=						_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():
				_UpperCamelCase   :		str			=						model(__snake_case)[0]
			_UpperCamelCase   :		Optional[int]			=						torch.Size((1, 9, 10_24))
			self.assertEqual(output.shape	,       __snake_case)

			_UpperCamelCase   :		Union[str, Any]			=						[-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   :		Optional[Any]			=						output[0, ii, jj]
					_UpperCamelCase   :		Dict			=						expected[3 * ii + jj]
					_UpperCamelCase   :		Optional[int]			=						'ii={} jj={} a={} b={}'.format(__snake_case	,       __snake_case	,       __snake_case	,       __snake_case)
					self.assertTrue(math.isclose(__snake_case	,       __snake_case	,       rel_tol=__snake_case	,       abs_tol=__snake_case)	,       msg=__snake_case)




