

"""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        snake_case      :

				def __init__(      self				:							List[str]	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							Dict=1_3	,    UpperCamelCase__				:							Optional[Any]=7	,    UpperCamelCase__				:							Dict=True	,    UpperCamelCase__				:							List[Any]=True	,    UpperCamelCase__				:							str=True	,    UpperCamelCase__				:							List[Any]=True	,    UpperCamelCase__				:							List[Any]=9_9	,    UpperCamelCase__				:							int=6_4	,    UpperCamelCase__				:							Dict=3_2	,    UpperCamelCase__				:							str=5	,    UpperCamelCase__				:							List[str]=4	,    UpperCamelCase__				:							List[str]=3_7	,    UpperCamelCase__				:							Optional[int]="gelu"	,    UpperCamelCase__				:							List[Any]=0.1	,    UpperCamelCase__				:							Any=0.1	,    UpperCamelCase__				:							Optional[Any]=5_1_2	,    UpperCamelCase__				:							Dict=1_6	,    UpperCamelCase__				:							Tuple=2	,    UpperCamelCase__				:							Optional[int]=0.02	,    UpperCamelCase__				:							Optional[int]=3	,    UpperCamelCase__				:							Tuple=4	,    UpperCamelCase__				:							Tuple=None	,    )->  Dict:



						'''simple docstring'''


						__lowerCAmelCase:			Tuple        =     parent
						__lowerCAmelCase:			Optional[int]        =     batch_size
						__lowerCAmelCase:			Optional[Any]        =     seq_length
						__lowerCAmelCase:			Optional[int]        =     is_training
						__lowerCAmelCase:			Tuple        =     use_input_mask
						__lowerCAmelCase:			List[str]        =     use_token_type_ids
						__lowerCAmelCase:			int        =     use_labels
						__lowerCAmelCase:			Any        =     vocab_size
						__lowerCAmelCase:			Optional[Any]        =     hidden_size
						__lowerCAmelCase:			Optional[int]        =     embedding_size
						__lowerCAmelCase:			Optional[Any]        =     num_hidden_layers
						__lowerCAmelCase:			Dict        =     num_attention_heads
						__lowerCAmelCase:			str        =     intermediate_size
						__lowerCAmelCase:			List[Any]        =     hidden_act
						__lowerCAmelCase:			str        =     hidden_dropout_prob
						__lowerCAmelCase:			Tuple        =     attention_probs_dropout_prob
						__lowerCAmelCase:			Optional[int]        =     max_position_embeddings
						__lowerCAmelCase:			Dict        =     type_vocab_size
						__lowerCAmelCase:			Optional[Any]        =     type_sequence_label_size
						__lowerCAmelCase:			Any        =     initializer_range
						__lowerCAmelCase:			str        =     num_labels
						__lowerCAmelCase:			Union[str, Any]        =     num_choices
						__lowerCAmelCase:			Optional[int]        =     scope

				def 				lowercase_	(      self				:							Tuple)->  Tuple:



						'''simple docstring'''


						__lowerCAmelCase:			Union[str, Any]        =     ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size)

						__lowerCAmelCase:			Dict        =     None
						if self.use_input_mask:
								__lowerCAmelCase:			List[Any]        =     random_attention_mask([self.batch_size, self.seq_length])

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

						__lowerCAmelCase:			Tuple        =     None
						__lowerCAmelCase:			Any        =     None
						__lowerCAmelCase:			Tuple        =     None
						if self.use_labels:
								__lowerCAmelCase:			Tuple        =     ids_tensor([self.batch_size]	,    self.type_sequence_label_size)
								__lowerCAmelCase:			Dict        =     ids_tensor([self.batch_size, self.seq_length]	,    self.num_labels)
								__lowerCAmelCase:			List[Any]        =     ids_tensor([self.batch_size]	,    self.num_choices)

						__lowerCAmelCase:			List[str]        =     self.get_config()

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

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



						'''simple docstring'''


						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=UpperCamelCase__	,    initializer_range=self.initializer_range	,    )

				def 				lowercase_	(      self				:							Optional[int]	,    UpperCamelCase__				:							List[Any]	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Optional[Any]	,    UpperCamelCase__				:							Optional[Any]	,    UpperCamelCase__				:							Tuple	,    UpperCamelCase__				:							Any	,    UpperCamelCase__				:							Dict)->  Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:			Optional[Any]        =     MegatronBertModel(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			str        =     model(UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__)
						__lowerCAmelCase:			List[Any]        =     model(UpperCamelCase__	,    token_type_ids=UpperCamelCase__)
						__lowerCAmelCase:			List[Any]        =     model(UpperCamelCase__)

						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 				lowercase_	(      self				:							List[str]	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							Dict	,    UpperCamelCase__				:							Dict	,    UpperCamelCase__				:							Tuple	,    UpperCamelCase__				:							Optional[Any]	,    UpperCamelCase__				:							Dict)->  List[str]:



						'''simple docstring'''


						__lowerCAmelCase:			int        =     MegatronBertForMaskedLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			str        =     model(UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size))

				def 				lowercase_	(      self				:							Dict	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Optional[int]	,    UpperCamelCase__				:							Optional[int]	,    UpperCamelCase__				:							int	,    UpperCamelCase__				:							int	,    UpperCamelCase__				:							Optional[int]	,    UpperCamelCase__				:							Any)->  Tuple:



						'''simple docstring'''


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

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



						'''simple docstring'''


						__lowerCAmelCase:			List[Any]        =     MegatronBertForNextSentencePrediction(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			Any        =     model(
						    UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    labels=UpperCamelCase__	,    )
						self.parent.assertEqual(result.logits.shape	,    (self.batch_size, 2))

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



						'''simple docstring'''


						__lowerCAmelCase:			List[Any]        =     MegatronBertForPreTraining(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			Any        =     model(
						    UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    labels=UpperCamelCase__	,    next_sentence_label=UpperCamelCase__	,    )
						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 				lowercase_	(      self				:							List[Any]	,    UpperCamelCase__				:							Union[str, Any]	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Tuple	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Union[str, Any]	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Tuple)->  Tuple:



						'''simple docstring'''


						__lowerCAmelCase:			str        =     MegatronBertForQuestionAnswering(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			Any        =     model(
						    UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    start_positions=UpperCamelCase__	,    end_positions=UpperCamelCase__	,    )
						self.parent.assertEqual(result.start_logits.shape	,    (self.batch_size, self.seq_length))
						self.parent.assertEqual(result.end_logits.shape	,    (self.batch_size, self.seq_length))

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



						'''simple docstring'''


						__lowerCAmelCase:			List[str]        =     self.num_labels
						__lowerCAmelCase:			str        =     MegatronBertForSequenceClassification(UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			Union[str, Any]        =     model(UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels))

				def 				lowercase_	(      self				:							Any	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							Union[str, Any]	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							Tuple	,    UpperCamelCase__				:							Dict	,    UpperCamelCase__				:							Tuple)->  List[Any]:



						'''simple docstring'''


						__lowerCAmelCase:			Optional[Any]        =     self.num_labels
						__lowerCAmelCase:			Any        =     MegatronBertForTokenClassification(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:			List[str]        =     model(UpperCamelCase__	,    attention_mask=UpperCamelCase__	,    token_type_ids=UpperCamelCase__	,    labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.num_labels))

				def 				lowercase_	(      self				:							Union[str, Any]	,    UpperCamelCase__				:							Tuple	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							Optional[Any]	,    UpperCamelCase__				:							List[Any]	,    UpperCamelCase__				:							str	,    UpperCamelCase__				:							Union[str, Any]	,    UpperCamelCase__				:							int)->  Optional[int]:



						'''simple docstring'''


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




				def 				lowercase_	(      self				:							List[str])->  Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:			Optional[int]        =     self.prepare_config_and_inputs()
						(
						    (
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,					(
						    __lowerCAmelCase						
						) ,											
						):			Any        =     config_and_inputs
						__lowerCAmelCase:			Tuple        =     {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class        snake_case      (							__snake_case,				__snake_case,				unittest.TestCase    ):
				SCREAMING_SNAKE_CASE_  : Optional[int]   =    (
				    (
				        MegatronBertModel,
				        MegatronBertForMaskedLM,
				        MegatronBertForCausalLM,
				        MegatronBertForMultipleChoice,
				        MegatronBertForNextSentencePrediction,
				        MegatronBertForPreTraining,
				        MegatronBertForQuestionAnswering,
				        MegatronBertForSequenceClassification,
				        MegatronBertForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE_  : str   =    (
				    {
				        "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 {}
				)
				SCREAMING_SNAKE_CASE_  : Optional[Any]   =    True
				# test_resize_embeddings = False
				SCREAMING_SNAKE_CASE_  : Union[str, Any]   =    False

				def 				lowercase_	(      self				:							List[Any]	,    UpperCamelCase__				:							List[Any]	,    UpperCamelCase__				:							List[str]	,    UpperCamelCase__				:							int=False)->  Any:



						'''simple docstring'''


						__lowerCAmelCase:			int        =     super()._prepare_for_class(UpperCamelCase__	,    UpperCamelCase__	,    return_labels=UpperCamelCase__)

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

				def 				lowercase_	(      self				:							str)->  Any:



						'''simple docstring'''


						__lowerCAmelCase:			List[Any]        =     MegatronBertModelTester(self)
						__lowerCAmelCase:			Union[str, Any]        =     ConfigTester(self	,    config_class=UpperCamelCase__	,    hidden_size=3_7)

				def 				lowercase_	(      self				:							Optional[int])->  int:



						'''simple docstring'''


						self.config_tester.run_common_tests()

				def 				lowercase_	(      self				:							Any)->  List[str]:



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase__)

				def 				lowercase_	(      self				:							Any)->  Union[str, Any]:



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase:			Optional[Any]        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase__)

				def 				lowercase_	(      self				:							int)->  Any:



						'''simple docstring'''


						__lowerCAmelCase:			str        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase__)




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



						'''simple docstring'''


						__lowerCAmelCase:			Any        =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase__)






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


__A		  =     1E-4





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

				@slow
				@unittest.skip("Model is not available.")
				def 				lowercase_	(      self				:							Union[str, Any])->  Union[str, Any]:



						'''simple docstring'''


						__lowerCAmelCase:			int        =     "nvidia/megatron-bert-uncased-345m"
						if "MYDIR" in os.environ:
								__lowerCAmelCase:			Tuple        =     os.path.join(os.environ["MYDIR"]	,    UpperCamelCase__)
						__lowerCAmelCase:			Optional[int]        =     MegatronBertModel.from_pretrained(UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.half()
						__lowerCAmelCase:			int        =     _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]])
						with torch.no_grad():
								__lowerCAmelCase:			str        =     model(UpperCamelCase__)[0]
						__lowerCAmelCase:			Dict        =     torch.Size((1, 9, 1_0_2_4))
						self.assertEqual(output.shape	,    UpperCamelCase__)

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







from __future__ import annotations

def 	_lowerCAmelCase (    A__    ,	A__ = None						):
      lowercase__  =	word_bank or []
      # create a table
      lowercase__  =	len(A__						) + 1

      lowercase__  =	[]
      for _ in range(A__						):
            table.append([]						)
      # seed value
      lowercase__  =	[[]]  # because empty string has empty combination

      # iterate through the indices
      for i in range(A__						):
            # condition
            if table[i] != []:
                  for word in word_bank:
                        # slice condition
                        if target[i : i + len(A__						)] == word:
                              lowercase__  =	[
                                  [word, *way] for way in table[i]
                              ]
                              # adds the word to every combination the current position holds
                              # now,push that combination to the table[i+len(word)]
                              table[i + len(A__						)] += new_combinations

    # combinations are in reverse order so reverse for better output
      for combination in table[len(A__						)]:
            combination.reverse()

      return table[len(A__						)]


if __name__ == "__main__":
       print(all_construct("jwajalapa", ["jwa", "j", "w", "a", "la", "lapa"]))
       print(all_construct("rajamati", ["s", "raj", "amat", "raja", "ma", "i", "t"]))
       print(
           all_construct(
               "hexagonosaurus",
               ["h", "ex", "hex", "ag", "ago", "ru", "auru", "rus", "go", "no", "o", "s"],
           )
       )
