





'''simple docstring'''


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					:				List[Any]       ,     a__					:				List[str]       ,     a__					:				int=13       ,     a__					:				Optional[Any]=7       ,     a__					:				Dict=True       ,     a__					:				Optional[Any]=True       ,     a__					:				int=False       ,     a__					:				Union[str, Any]=True       ,     a__					:				List[Any]=99       ,     a__					:				Optional[int]=32       ,     a__					:				List[str]=5       ,     a__					:				Optional[Any]=4       ,     a__					:				Optional[Any]=37       ,     a__					:				str="gelu"       ,     a__					:				List[str]=0.1       ,     a__					:				Optional[Any]=0.1       ,     a__					:				Tuple=512       ,     a__					:				List[Any]=16       ,     a__					:				int=2       ,     a__					:				Dict=0.02       ,     a__					:				Tuple=3       ,     a__					:				List[str]=4       ,     a__					:				Optional[Any]=None       ,     ):
              UpperCAmelCase  =						parent
              UpperCAmelCase  =						batch_size
              UpperCAmelCase  =						seq_length
              UpperCAmelCase  =						is_training
              UpperCAmelCase  =						use_input_mask
              UpperCAmelCase  =						use_token_type_ids
              UpperCAmelCase  =						use_labels
              UpperCAmelCase  =						vocab_size
              UpperCAmelCase  =						hidden_size
              UpperCAmelCase  =						num_hidden_layers
              UpperCAmelCase  =						num_attention_heads
              UpperCAmelCase  =						intermediate_size
              UpperCAmelCase  =						hidden_act
              UpperCAmelCase  =						hidden_dropout_prob
              UpperCAmelCase  =						attention_probs_dropout_prob
              UpperCAmelCase  =						max_position_embeddings
              UpperCAmelCase  =						type_vocab_size
              UpperCAmelCase  =						type_sequence_label_size
              UpperCAmelCase  =						initializer_range
              UpperCAmelCase  =						num_labels
              UpperCAmelCase  =						num_choices
              UpperCAmelCase  =						scope

       def 			__snake_case      ( self					:				Dict						):
              UpperCAmelCase  =						ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size						)

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

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

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

              UpperCAmelCase  =						self.get_config()

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

       def 			__snake_case      ( self					:				Dict						):
              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=a__       ,     initializer_range=self.initializer_range       ,     use_stable_embedding=a__       ,     )

       def 			__snake_case      ( self					:				List[Any]       ,     a__					:				Optional[int]       ,     a__					:				str       ,     a__					:				int       ,     a__					:				Optional[Any]       ,     a__					:				Optional[int]       ,     a__					:				Tuple       ,     a__					:				List[str]						):
              UpperCAmelCase  =						OpenLlamaModel(config=a__						)
              model.to(a__						)
              model.eval()
              UpperCAmelCase  =						model(a__       ,     attention_mask=a__						)
              UpperCAmelCase  =						model(a__						)
              self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)						)

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

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

       def 			__snake_case      ( self					:				str       ,     a__					:				Any       ,     a__					:				int       ,     a__					:				Union[str, Any]       ,     a__					:				Tuple       ,     a__					:				Optional[Any]       ,     a__					:				Any       ,     a__					:				str       ,     a__					:				Union[str, Any]       ,     a__					:				str       ,     ):
              UpperCAmelCase  =						True
              UpperCAmelCase  =						True
              UpperCAmelCase  =						OpenLlamaForCausalLM(config=a__						)
              model.to(a__						)
              model.eval()

              # first forward pass
              UpperCAmelCase  =						model(
                  a__       ,     attention_mask=a__       ,     encoder_hidden_states=a__       ,     encoder_attention_mask=a__       ,     use_cache=a__       ,     )
              UpperCAmelCase  =						outputs.past_key_values

              # create hypothetical multiple next token and extent to next_input_ids
              UpperCAmelCase  =						ids_tensor((self.batch_size, 3)       ,     config.vocab_size						)
              UpperCAmelCase  =						ids_tensor((self.batch_size, 3)       ,     vocab_size=2						)

              # append to next input_ids and
              UpperCAmelCase  =						torch.cat([input_ids, next_tokens]       ,     dim=-1						)
              UpperCAmelCase  =						torch.cat([input_mask, next_mask]       ,     dim=-1						)

              UpperCAmelCase  =						model(
                  a__       ,     attention_mask=a__       ,     encoder_hidden_states=a__       ,     encoder_attention_mask=a__       ,     output_hidden_states=a__       ,     )['''hidden_states'''][0]
              UpperCAmelCase  =						model(
                  a__       ,     attention_mask=a__       ,     encoder_hidden_states=a__       ,     encoder_attention_mask=a__       ,     past_key_values=a__       ,     output_hidden_states=a__       ,     )['''hidden_states'''][0]

              # select random slice
              UpperCAmelCase  =						ids_tensor((1,)       ,     output_from_past.shape[-1]						).item()
              UpperCAmelCase  =						output_from_no_past[:, -3:, random_slice_idx].detach()
              UpperCAmelCase  =						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(a__       ,     a__       ,     atol=1e-3						)						)



       def 			__snake_case      ( self					:				Union[str, Any]						):
              UpperCAmelCase  =						self.prepare_config_and_inputs()
              (
                  (
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					(
                  UpperCAmelCase
              ),					
              )  =						config_and_inputs
              UpperCAmelCase  =						{'''input_ids''': input_ids, '''attention_mask''': input_mask}
              return config, inputs_dict






@require_torch
class     lowerCAmelCase__ (				UpperCAmelCase_      , UpperCAmelCase_      , UpperCAmelCase_      , unittest.TestCase    ):

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

       def 			__snake_case      ( self					:				Optional[Any]						):
              UpperCAmelCase  =						OpenLlamaModelTester(self						)
              UpperCAmelCase  =						ConfigTester(self       ,     config_class=a__       ,     hidden_size=37						)

       def 			__snake_case      ( self					:				List[str]						):
              self.config_tester.run_common_tests()

       def 			__snake_case      ( self					:				Optional[Any]						):
              UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*a__						)

       def 			__snake_case      ( self					:				Dict						):
              UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs()
              for type in ["absolute", "relative_key", "relative_key_query"]:
                     UpperCAmelCase  =						type
                     self.model_tester.create_and_check_model(*a__						)

       def 			__snake_case      ( self					:				Dict						):
              UpperCAmelCase,					UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs_for_common()
              UpperCAmelCase  =						3
              UpperCAmelCase  =						input_dict['''input_ids''']
              UpperCAmelCase  =						input_ids.ne(1						).to(a__						)
              UpperCAmelCase  =						ids_tensor([self.model_tester.batch_size]       ,     self.model_tester.type_sequence_label_size						)
              UpperCAmelCase  =						OpenLlamaForSequenceClassification(a__						)
              model.to(a__						)
              model.eval()
              UpperCAmelCase  =						model(a__       ,     attention_mask=a__       ,     labels=a__						)
              self.assertEqual(result.logits.shape       ,     (self.model_tester.batch_size, self.model_tester.num_labels)						)

       def 			__snake_case      ( self					:				Any						):
              UpperCAmelCase,					UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs_for_common()
              UpperCAmelCase  =						3
              UpperCAmelCase  =						'''single_label_classification'''
              UpperCAmelCase  =						input_dict['''input_ids''']
              UpperCAmelCase  =						input_ids.ne(1						).to(a__						)
              UpperCAmelCase  =						ids_tensor([self.model_tester.batch_size]       ,     self.model_tester.type_sequence_label_size						)
              UpperCAmelCase  =						OpenLlamaForSequenceClassification(a__						)
              model.to(a__						)
              model.eval()
              UpperCAmelCase  =						model(a__       ,     attention_mask=a__       ,     labels=a__						)
              self.assertEqual(result.logits.shape       ,     (self.model_tester.batch_size, self.model_tester.num_labels)						)

       def 			__snake_case      ( self					:				Tuple						):
              UpperCAmelCase,					UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs_for_common()
              UpperCAmelCase  =						3
              UpperCAmelCase  =						'''multi_label_classification'''
              UpperCAmelCase  =						input_dict['''input_ids''']
              UpperCAmelCase  =						input_ids.ne(1						).to(a__						)
              UpperCAmelCase  =						ids_tensor(
                  [self.model_tester.batch_size, config.num_labels]       ,     self.model_tester.type_sequence_label_size						).to(torch.float						)
              UpperCAmelCase  =						OpenLlamaForSequenceClassification(a__						)
              model.to(a__						)
              model.eval()
              UpperCAmelCase  =						model(a__       ,     attention_mask=a__       ,     labels=a__						)
              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					:				Optional[int]						):
              pass



       @parameterized.expand([('''linear''',), ('''dynamic''',)]						)
       def 			__snake_case      ( self					:				Any       ,     a__					:				List[str]						):
              UpperCAmelCase,					UpperCAmelCase  =						self.model_tester.prepare_config_and_inputs_for_common()
              UpperCAmelCase  =						ids_tensor([1, 10]       ,     config.vocab_size						)
              UpperCAmelCase  =						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
              UpperCAmelCase  =						OpenLlamaModel(a__						)
              original_model.to(a__						)
              original_model.eval()
              UpperCAmelCase  =						original_model(a__						).last_hidden_state
              UpperCAmelCase  =						original_model(a__						).last_hidden_state

              set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
              UpperCAmelCase  =						{'''type''': scaling_type, '''factor''': 10.0}
              UpperCAmelCase  =						OpenLlamaModel(a__						)
              scaled_model.to(a__						)
              scaled_model.eval()
              UpperCAmelCase  =						scaled_model(a__						).last_hidden_state
              UpperCAmelCase  =						scaled_model(a__						).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(a__       ,     a__       ,     atol=1e-5						)						)
              else:
                     self.assertFalse(torch.allclose(a__       ,     a__       ,     atol=1e-5						)						)

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









'''simple docstring'''




from typing import TYPE_CHECKING

from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


__UpperCAmelCase            =       {"""configuration_gpt_neox""": ["""GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPTNeoXConfig"""]}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__UpperCAmelCase            =       ["""GPTNeoXTokenizerFast"""]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__UpperCAmelCase            =       [
							    """GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST""",
							    """GPTNeoXForCausalLM""",
							    """GPTNeoXForQuestionAnswering""",
							    """GPTNeoXForSequenceClassification""",
							    """GPTNeoXForTokenClassification""",
							    """GPTNeoXLayer""",
							    """GPTNeoXModel""",
							    """GPTNeoXPreTrainedModel""",
							]


if TYPE_CHECKING:
							from .configuration_gpt_neox import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXConfig

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_gpt_neox_fast import GPTNeoXTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_gpt_neox import (
														    GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST,
														    GPTNeoXForCausalLM,
														    GPTNeoXForQuestionAnswering,
														    GPTNeoXForSequenceClassification,
														    GPTNeoXForTokenClassification,
														    GPTNeoXLayer,
														    GPTNeoXModel,
														    GPTNeoXPreTrainedModel,
														)


else:
							import sys

							__UpperCAmelCase            =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



