


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




							"""simple docstring"""


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

							def        __snake_case   (				self			:    List[str]							)      ->       int:
														__snake_case						:     str									=  ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size							)

														__snake_case						:     List[str]									=  None
														if self.use_input_mask:
																					__snake_case						:     Union[str, Any]									=  random_attention_mask([self.batch_size, self.seq_length]							)

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

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

														__snake_case						:     Optional[Any]									=  self.get_config()

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

							def        __snake_case   (				self			:    Any							)      ->       int:
														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=lowerCamelCase       ,      initializer_range=self.initializer_range       ,      )

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

														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        __snake_case   (				self			:    Optional[int]       ,      lowerCamelCase			:    Union[str, Any]       ,      lowerCamelCase			:    Union[str, Any]       ,      lowerCamelCase			:    Dict       ,      lowerCamelCase			:    int       ,      lowerCamelCase			:    List[Any]       ,      lowerCamelCase			:    List[Any]       ,      lowerCamelCase			:    List[Any]							)      ->       List[Any]:
														__snake_case						:     Union[str, Any]									=  MegatronBertForMaskedLM(config=lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.eval()
														__snake_case						:     Optional[int]									=  model(lowerCamelCase       ,      attention_mask=lowerCamelCase       ,      token_type_ids=lowerCamelCase       ,      labels=lowerCamelCase							)
														self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.vocab_size)							)

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

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

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

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

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






							def        __snake_case   (				self			:    Dict							)      ->       Tuple:
														__snake_case						:     Union[str, Any]									=  self.prepare_config_and_inputs()
														(
														    (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     (
														    __snake_case
														)					,     
														)						:     int									=  config_and_inputs
														__snake_case						:     Optional[Any]									=  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
														return config, inputs_dict

@require_torch
class 							a						(_lowerCAmelCase	,							_lowerCAmelCase	,							unittest.TestCase ):




							"""simple docstring"""


							__UpperCAmelCase       :				str	    =    (
							    (
							        MegatronBertModel,
							        MegatronBertForMaskedLM,
							        MegatronBertForCausalLM,
							        MegatronBertForMultipleChoice,
							        MegatronBertForNextSentencePrediction,
							        MegatronBertForPreTraining,
							        MegatronBertForQuestionAnswering,
							        MegatronBertForSequenceClassification,
							        MegatronBertForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							__UpperCAmelCase       :				Optional[Any]	    =    (
							    {
							        "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 {}
							)
							__UpperCAmelCase       :				List[str]	    =    True
							# test_resize_embeddings = False
							__UpperCAmelCase       :				Union[str, Any]	    =    False

							def        __snake_case   (				self			:    Any       ,      lowerCamelCase			:    Tuple       ,      lowerCamelCase			:    int       ,      lowerCamelCase			:    List[str]=False							)      ->       Any:
														__snake_case						:     Tuple									=  super()._prepare_for_class(lowerCamelCase       ,      lowerCamelCase       ,      return_labels=lowerCamelCase							)

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

							def        __snake_case   (				self			:    List[str]							)      ->       int:
														__snake_case						:     Any									=  MegatronBertModelTester(self							)
														__snake_case						:     Union[str, Any]									=  ConfigTester(self       ,      config_class=lowerCamelCase       ,      hidden_size=37							)

							def        __snake_case   (				self			:    str							)      ->       str:
														self.config_tester.run_common_tests()

							def        __snake_case   (				self			:    List[str]							)      ->       Optional[Any]:
														__snake_case						:     List[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase							)

							def        __snake_case   (				self			:    Optional[int]							)      ->       Optional[Any]:
														__snake_case						:     Optional[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase							)

							def        __snake_case   (				self			:    str							)      ->       List[Any]:
														__snake_case						:     Optional[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase							)

							def        __snake_case   (				self			:    List[Any]							)      ->       Optional[int]:
														__snake_case						:     Union[str, Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase							)

							def        __snake_case   (				self			:    List[str]							)      ->       str:
														__snake_case						:     Optional[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase							)

							def        __snake_case   (				self			:    List[Any]							)      ->       Any:
														__snake_case						:     Optional[Any]									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase							)

							def        __snake_case   (				self			:    int							)      ->       List[Any]:
														__snake_case						:     int									=  self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase							)






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







def 							lowerCAmelCase_	(   __lowerCamelCase      ):
							return torch.tensor(
							    __lowerCamelCase				,      dtype=torch.long				,      device=__lowerCamelCase				,      )


_snake_case	:						Optional[Any]         =							1E-4

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




							"""simple docstring"""


							@slow
							@unittest.skip("Model is not available."							)
							def        __snake_case   (				self			:    List[Any]							)      ->       Dict:
														__snake_case						:     List[str]									=  "nvidia/megatron-bert-uncased-345m"
														if "MYDIR" in os.environ:
																					__snake_case						:     str									=  os.path.join(os.environ["MYDIR"]       ,      lowerCamelCase							)
														__snake_case						:     Dict									=  MegatronBertModel.from_pretrained(lowerCamelCase							)
														model.to(lowerCamelCase							)
														model.half()
														__snake_case						:     int									=  _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]							)
														with torch.no_grad():
																					__snake_case						:     Optional[Any]									=  model(lowerCamelCase							)[0]
														__snake_case						:     str									=  torch.Size((1, 9, 1024)							)
														self.assertEqual(output.shape       ,      lowerCamelCase							)

														__snake_case						:     Any									=  [-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
														for ii in range(3							):
																					for jj in range(3							):
																												__snake_case						:     List[Any]									=  output[0, ii, jj]
																												__snake_case						:     Tuple									=  expected[3 * ii + jj]
																												__snake_case						:     List[str]									=  "ii={} jj={} a={} b={}".format(lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase							)
																												self.assertTrue(math.isclose(lowerCamelCase       ,      lowerCamelCase       ,      rel_tol=lowerCamelCase       ,      abs_tol=lowerCamelCase							)       ,      msg=lowerCamelCase							)



import functools
from typing import Any



def 		UpperCamelCase     (       _UpperCAmelCase  :   str ,			_UpperCAmelCase  :   list[str]  )   ->			bool:






						'''simple docstring'''

						if not isinstance(_UpperCAmelCase ,			_UpperCAmelCase  ) or len(_UpperCAmelCase  ) == 0:
												raise ValueError("the string should be not empty string"  )

						if not isinstance(_UpperCAmelCase ,			_UpperCAmelCase  ) or not all(
						    isinstance(_UpperCAmelCase ,			_UpperCAmelCase  ) and len(_UpperCAmelCase  ) > 0 for item in words  ):
												raise ValueError("the words should be a list of non-empty strings"  )

						# Build trie
						_lowercase       :	dict[str, Any]           =   {}
						_lowercase       :	List[Any]           =   "WORD_KEEPER"

						for word in words:
												_lowercase       :	Optional[Any]           =   trie
												for c in word:
																		if c not in trie_node:
																								_lowercase       :	List[Any]           =   {}

																		_lowercase       :	Optional[int]           =   trie_node[c]

												_lowercase       :	str           =   True

						_lowercase       :	int           =   len(_UpperCAmelCase  )

						# Dynamic programming method
						@functools.cache
						def is_breakable(_UpperCAmelCase  :   int  ) -> bool:
												if index == len_string:
																		return True

												_lowercase       :	Union[str, Any]           =   trie
												for i in range(_UpperCAmelCase ,			_UpperCAmelCase  ):
																		_lowercase       :	List[str]           =   trie_node.get(string[i] ,			_UpperCAmelCase  )

																		if trie_node is None:
																								return False

																		if trie_node.get(_UpperCAmelCase ,			_UpperCAmelCase  ) and is_breakable(i + 1  ):
																								return True

												return False

						return is_breakable(0  )


if __name__ == "__main__":
							import doctest

							doctest.testmod()




