



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


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


				def __UpperCamelCase						(							self      :  Dict       )		->     Optional[int]:
											A      						=    ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size       )

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

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

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

											A      						=    self.get_config()

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


				def __UpperCamelCase						(							self      :  List[Any]       )		->     List[Any]:
											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 __UpperCamelCase						(							self      :  Tuple						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  Optional[Any]       )		->     Optional[int]:
											A      						=    MegatronBertModel(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    model(__UpperCamelCase						, attention_mask=__UpperCamelCase						, token_type_ids=__UpperCamelCase       )
											A      						=    model(__UpperCamelCase						, token_type_ids=__UpperCamelCase       )
											A      						=    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 __UpperCamelCase						(							self      :  Optional[int]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  Tuple						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  Tuple       )		->     Union[str, Any]:
											A      						=    MegatronBertForMaskedLM(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    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 __UpperCamelCase						(							self      :  Optional[Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  Optional[Any]       )		->     int:
											A      						=    MegatronBertForCausalLM(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    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 __UpperCamelCase						(							self      :  List[Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Tuple						, __UpperCamelCase      :  List[str]       )		->     Any:
											A      						=    MegatronBertForNextSentencePrediction(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    model(
											    __UpperCamelCase						, attention_mask=__UpperCamelCase						, token_type_ids=__UpperCamelCase						, labels=__UpperCamelCase						, )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, 2)       )


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


				def __UpperCamelCase						(							self      :  Tuple						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  int						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  List[Any]       )		->     Tuple:
											A      						=    self.num_labels
											A      						=    MegatronBertForTokenClassification(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    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 __UpperCamelCase						(							self      :  Optional[int]						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  List[Any]       )		->     List[Any]:
											A      						=    self.num_choices
											A      						=    MegatronBertForMultipleChoice(config=__UpperCamelCase       )
											model.to(__UpperCamelCase       )
											model.eval()
											A      						=    input_ids.unsqueeze(1       ).expand(-1						, self.num_choices						, -1       ).contiguous()
											A      						=    token_type_ids.unsqueeze(1       ).expand(-1						, self.num_choices						, -1       ).contiguous()
											A      						=    input_mask.unsqueeze(1       ).expand(-1						, self.num_choices						, -1       ).contiguous()
											A      						=    model(
											    __UpperCamelCase						, attention_mask=__UpperCamelCase						, token_type_ids=__UpperCamelCase						, labels=__UpperCamelCase						, )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_choices)       )


				def __UpperCamelCase						(							self      :  Union[str, Any]       )		->     Dict:
											A      						=    self.prepare_config_and_inputs()
											(
											    (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   
											)      						=    config_and_inputs
											A      						=    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
											return config, inputs_dict





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


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

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


				def __UpperCamelCase						(							self      :  str       )		->     str:
											A      						=    MegatronBertModelTester(self       )
											A      						=    ConfigTester(self						, config_class=__UpperCamelCase						, hidden_size=37       )


				def __UpperCamelCase						(							self      :  Union[str, Any]       )		->     Tuple:
											self.config_tester.run_common_tests()


				def __UpperCamelCase						(							self      :  int       )		->     List[str]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_model(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     Union[str, Any]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Optional[Any]       )		->     str:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  List[Any]       )		->     List[str]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  List[str]       )		->     Tuple:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Any       )		->     Dict:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Dict       )		->     List[Any]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCamelCase       )


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






def  lowerCamelCase_	(      lowerCAmelCase__	:	Union[str, Any]				)       ->							Optional[int]:

							'''simple docstring'''






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


__snake_case     :Optional[Any]			       =1E-4



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


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

											A      						=    [-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       ):
																									A      						=    output[0, ii, jj]
																									A      						=    expected[3 * ii + jj]
																									A      						=    'ii={} jj={} a={} b={}'.format(__UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase       )
																									self.assertTrue(math.isclose(__UpperCamelCase						, __UpperCamelCase						, rel_tol=__UpperCamelCase						, abs_tol=__UpperCamelCase       )						, msg=__UpperCamelCase       )






import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel


class       __lowerCAmelCase					:


	'''simple docstring'''





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

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

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

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

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

			lowercase__					:      Tuple      =		self.get_config()

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

	def    snake_case__(						self:						Tuple			):
			return OpenLlamaConfig(
			    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,					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,					use_stable_embedding=lowerCamelCase_,					)

	def    snake_case__(						self:						Optional[int],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						List[Any],					lowerCamelCase_:						List[str],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						Dict,					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						str			):
			lowercase__					:      Union[str, Any]      =		OpenLlamaModel(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Union[str, Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_			)
			lowercase__					:      List[Any]      =		model(lowerCamelCase_			)
			self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)			)

	def    snake_case__(						self:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						str,					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Tuple,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						List[str],					lowerCamelCase_:						Dict,					lowerCamelCase_:						Union[str, Any],					):
			lowercase__					:      Tuple      =		True
			lowercase__					:      int      =		OpenLlamaModel(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					)
			lowercase__					:      Optional[int]      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					)
			lowercase__					:      str      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_			)
			self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)			)

	def    snake_case__(						self:						List[Any],					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Any,					lowerCamelCase_:						str,					lowerCamelCase_:						List[str],					lowerCamelCase_:						Any,					lowerCamelCase_:						Dict,					lowerCamelCase_:						int,					lowerCamelCase_:						Any,					lowerCamelCase_:						str,					):
			lowercase__					:      Optional[Any]      =		OpenLlamaForCausalLM(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.vocab_size)			)

	def    snake_case__(						self:						Tuple,					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Optional[int],					lowerCamelCase_:						str,					lowerCamelCase_:						List[Any],					lowerCamelCase_:						List[Any],					lowerCamelCase_:						str,					lowerCamelCase_:						Optional[Any],					lowerCamelCase_:						Union[str, Any],					):
			lowercase__					:      Optional[int]      =		True
			lowercase__					:      Optional[int]      =		True
			lowercase__					:      Tuple      =		OpenLlamaForCausalLM(config=lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()

			# first forward pass
			lowercase__					:      Any      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					use_cache=lowerCamelCase_,					)
			lowercase__					:      str      =		outputs.past_key_values

			# create hypothetical multiple next token and extent to next_input_ids
			lowercase__					:      Optional[Any]      =		ids_tensor((self.batch_size, 3),					config.vocab_size			)
			lowercase__					:      str      =		ids_tensor((self.batch_size, 3),					vocab_size=2			)

			# append to next input_ids and
			lowercase__					:      Tuple      =		torch.cat([input_ids, next_tokens],					dim=-1			)
			lowercase__					:      List[str]      =		torch.cat([input_mask, next_mask],					dim=-1			)

			lowercase__					:      int      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					output_hidden_states=lowerCamelCase_,					)['hidden_states'][0]
			lowercase__					:      int      =		model(
			    lowerCamelCase_,					attention_mask=lowerCamelCase_,					encoder_hidden_states=lowerCamelCase_,					encoder_attention_mask=lowerCamelCase_,					past_key_values=lowerCamelCase_,					output_hidden_states=lowerCamelCase_,					)['hidden_states'][0]

			# select random slice
			lowercase__					:      Tuple      =		ids_tensor((1,),					output_from_past.shape[-1]			).item()
			lowercase__					:      Union[str, Any]      =		output_from_no_past[:, -3:, random_slice_idx].detach()
			lowercase__					:      Tuple      =		output_from_past[:, :, random_slice_idx].detach()

			self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]			)

			# test that outputs are equal for slice
			self.parent.assertTrue(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-3			)			)





	def    snake_case__(						self:						Optional[Any]			):
			lowercase__					:      Dict      =		self.prepare_config_and_inputs()
			(
			    (
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				(
			    lowercase__
			)   ,				
			)					:      List[str]      =		config_and_inputs
			lowercase__					:      Any      =		{'input_ids': input_ids, 'attention_mask': input_mask}
			return config, inputs_dict


@require_torch
class       __lowerCAmelCase					(       _UpperCamelCase   ,    _UpperCamelCase   ,    _UpperCamelCase   ,    unittest.TestCase      ):


	'''simple docstring'''





	_A													=				(
	    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
	)
	_A													=				(OpenLlamaForCausalLM,) if is_torch_available() else ()
	_A													=				(
	    {
	        "feature-extraction": OpenLlamaModel,
	        "text-classification": OpenLlamaForSequenceClassification,
	        "text-generation": OpenLlamaForCausalLM,
	        "zero-shot": OpenLlamaForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A													=				False
	_A													=				False

	def    snake_case__(						self:						Any			):
			lowercase__					:      List[Any]      =		OpenLlamaModelTester(self			)
			lowercase__					:      Dict      =		ConfigTester(self,					config_class=lowerCamelCase_,					hidden_size=37			)

	def    snake_case__(						self:						Any			):
			self.config_tester.run_common_tests()

	def    snake_case__(						self:						Dict			):
			lowercase__					:      Any      =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*lowerCamelCase_			)

	def    snake_case__(						self:						str			):
			lowercase__					:      Any      =		self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
					lowercase__					:      Dict      =		type
					self.model_tester.create_and_check_model(*lowerCamelCase_			)

	def    snake_case__(						self:						Any			):
			lowercase__   ,				lowercase__					:      List[Any]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Tuple      =		3
			lowercase__					:      Union[str, Any]      =		input_dict['input_ids']
			lowercase__					:      Union[str, Any]      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      int      =		ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size			)
			lowercase__					:      Optional[Any]      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Optional[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	def    snake_case__(						self:						Any			):
			lowercase__   ,				lowercase__					:      int      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      str      =		3
			lowercase__					:      Optional[Any]      =		'single_label_classification'
			lowercase__					:      Union[str, Any]      =		input_dict['input_ids']
			lowercase__					:      str      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      Any      =		ids_tensor([self.model_tester.batch_size],					self.model_tester.type_sequence_label_size			)
			lowercase__					:      Dict      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      Optional[int]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	def    snake_case__(						self:						Union[str, Any]			):
			lowercase__   ,				lowercase__					:      Optional[int]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Dict      =		3
			lowercase__					:      List[Any]      =		'multi_label_classification'
			lowercase__					:      Any      =		input_dict['input_ids']
			lowercase__					:      int      =		input_ids.ne(1			).to(lowerCamelCase_			)
			lowercase__					:      Optional[int]      =		ids_tensor(
			    [self.model_tester.batch_size, config.num_labels],					self.model_tester.type_sequence_label_size			).to(torch.float			)
			lowercase__					:      Optional[int]      =		OpenLlamaForSequenceClassification(lowerCamelCase_			)
			model.to(lowerCamelCase_			)
			model.eval()
			lowercase__					:      List[Any]      =		model(lowerCamelCase_,					attention_mask=lowerCamelCase_,					labels=lowerCamelCase_			)
			self.assertEqual(result.logits.shape,					(self.model_tester.batch_size, self.model_tester.num_labels)			)

	@unittest.skip('Open-Llama buffers include complex numbers, which breaks this test'			)
	def    snake_case__(						self:						List[str]			):
			pass





	@parameterized.expand([('linear',), ('dynamic',)]			)
	def    snake_case__(						self:						Any,					lowerCamelCase_:						List[Any]			):
			lowercase__   ,				lowercase__					:      Optional[Any]      =		self.model_tester.prepare_config_and_inputs_for_common()
			lowercase__					:      Optional[int]      =		ids_tensor([1, 10],					config.vocab_size			)
			lowercase__					:      Optional[Any]      =		ids_tensor([1, int(config.max_position_embeddings * 1.5			)],					config.vocab_size			)

			set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
			lowercase__					:      str      =		OpenLlamaModel(lowerCamelCase_			)
			original_model.to(lowerCamelCase_			)
			original_model.eval()
			lowercase__					:      Union[str, Any]      =		original_model(lowerCamelCase_			).last_hidden_state
			lowercase__					:      List[str]      =		original_model(lowerCamelCase_			).last_hidden_state

			set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
			lowercase__					:      List[str]      =		{'type': scaling_type, 'factor': 1_0.0}
			lowercase__					:      Union[str, Any]      =		OpenLlamaModel(lowerCamelCase_			)
			scaled_model.to(lowerCamelCase_			)
			scaled_model.eval()
			lowercase__					:      Dict      =		scaled_model(lowerCamelCase_			).last_hidden_state
			lowercase__					:      Union[str, Any]      =		scaled_model(lowerCamelCase_			).last_hidden_state

			# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
			# maximum sequence length, so the outputs for the short input should match.
			if scaling_type == "dynamic":
					self.assertTrue(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)
			else:
					self.assertFalse(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)

			# The output should be different for long inputs
			self.assertFalse(torch.allclose(lowerCamelCase_,					lowerCamelCase_,					atol=1E-5			)			)


