






'''simple docstring'''


import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, 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 (
					    GPTNeoXForCausalLM,
					    GPTNeoXForQuestionAnswering,
					    GPTNeoXForSequenceClassification,
					    GPTNeoXForTokenClassification,
					    GPTNeoXModel,
					)




class    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





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

									'''simple docstring'''







									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
									UpperCamelCase__         = vocab_size - 1






			def 				A      (   self							:						Tuple			)							->    Optional[int]:

									'''simple docstring'''







									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_labels:
															UpperCamelCase__         = ids_tensor([self.batch_size, self.seq_length] ,      self.num_labels			)

									UpperCamelCase__         = self.get_config()

									return config, input_ids, input_mask, token_labels






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

									'''simple docstring'''







									return GPTNeoXConfig(
									    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=lowercase__ ,      initializer_range=self.initializer_range ,      pad_token_id=self.pad_token_id ,      )






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

									'''simple docstring'''







									UpperCamelCase__         = self.prepare_config_and_inputs()

									UpperCamelCase__         = True

									return config, input_ids, input_mask, token_labels






			def 				A      (   self							:						List[Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						str ,      lowercase							:						Optional[int]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = GPTNeoXModel(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__			)
									UpperCamelCase__         = model(lowercase__			)
									self.parent.assertEqual(result.last_hidden_state.shape ,      (self.batch_size, self.seq_length, self.hidden_size)			)






			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						Optional[int] ,      lowercase							:						Tuple ,      lowercase							:						Union[str, Any]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = True
									UpperCamelCase__         = GPTNeoXModel(lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__			)
									self.parent.assertEqual(result.last_hidden_state.shape ,      (self.batch_size, self.seq_length, self.hidden_size)			)






			def 				A      (   self							:						Tuple ,      lowercase							:						List[Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						Tuple ,      lowercase							:						Optional[Any]			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = GPTNeoXForCausalLM(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__ ,      labels=lowercase__			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.vocab_size)			)






			def 				A      (   self							:						int ,      lowercase							:						Any ,      lowercase							:						Tuple ,      lowercase							:						List[str] ,      lowercase							:						List[str]			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = GPTNeoXForQuestionAnswering(lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__			)
									self.parent.assertEqual(result.start_logits.shape ,      (self.batch_size, self.seq_length)			)
									self.parent.assertEqual(result.end_logits.shape ,      (self.batch_size, self.seq_length)			)






			def 				A      (   self							:						Tuple ,      lowercase							:						Any ,      lowercase							:						Tuple ,      lowercase							:						List[Any] ,      lowercase							:						Any			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = GPTNeoXForSequenceClassification(lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = ids_tensor([self.batch_size] ,      self.type_sequence_label_size			)
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__ ,      labels=lowercase__			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_labels)			)






			def 				A      (   self							:						List[Any] ,      lowercase							:						List[str] ,      lowercase							:						Any ,      lowercase							:						Optional[Any] ,      lowercase							:						Optional[int]			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = GPTNeoXForTokenClassification(lowercase__			)
									model.to(lowercase__			)
									model.eval()
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__ ,      labels=lowercase__			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.num_labels)			)






			def 				A      (   self							:						Optional[Any] ,      lowercase							:						Tuple ,      lowercase							:						Dict ,      lowercase							:						List[Any]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = True
									UpperCamelCase__         = GPTNeoXForCausalLM(config=lowercase__			)
									model.to(lowercase__			)
									model.eval()

									# first forward pass
									UpperCamelCase__         = model(lowercase__ ,      attention_mask=lowercase__ ,      use_cache=lowercase__			)
									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(lowercase__ ,      attention_mask=lowercase__ ,      output_hidden_states=lowercase__			)
									UpperCamelCase__         = output_from_no_past['''hidden_states'''][0]
									UpperCamelCase__         = model(
									    lowercase__ ,      attention_mask=lowercase__ ,      past_key_values=lowercase__ ,      output_hidden_states=lowercase__ ,      )['''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(lowercase__ ,      lowercase__ ,      atol=1e-3			)			)







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

									'''simple docstring'''







									UpperCamelCase__         = self.prepare_config_and_inputs()
									UpperCamelCase__         = config_and_inputs
									UpperCamelCase__         = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
									return config, inputs_dict




@require_torch
class    _SCREAMING_SNAKE_CASE							(		a__       ,a__       ,a__       ,unittest.TestCase	):


			'''simple docstring'''





			__a							:							List[Any]					     =					(
			    (
			        GPTNeoXModel,
			        GPTNeoXForCausalLM,
			        GPTNeoXForQuestionAnswering,
			        GPTNeoXForSequenceClassification,
			        GPTNeoXForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__a							:							Optional[int]					     =					(GPTNeoXForCausalLM,) if is_torch_available() else ()
			__a							:							Union[str, Any]					     =					(
			    {
			        """feature-extraction""": GPTNeoXModel,
			        """question-answering""": GPTNeoXForQuestionAnswering,
			        """text-classification""": GPTNeoXForSequenceClassification,
			        """text-generation""": GPTNeoXForCausalLM,
			        """token-classification""": GPTNeoXForTokenClassification,
			        """zero-shot""": GPTNeoXForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__a							:							List[Any]					     =					False
			__a							:							Optional[Any]					     =					False
			__a							:							Optional[int]					     =					False
			__a							:							Union[str, Any]					     =					False






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

									'''simple docstring'''







									UpperCamelCase__         = GPTNeoXModelTester(self			)
									UpperCamelCase__         = ConfigTester(self ,      config_class=lowercase__ ,      hidden_size=6_4 ,      num_attention_heads=8			)






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

									'''simple docstring'''







									self.config_tester.run_common_tests()






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

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(lowercase__ ,      lowercase__ ,      lowercase__			)






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

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs_for_decoder()
									self.model_tester.create_and_check_model_as_decoder(lowercase__ ,      lowercase__ ,      lowercase__			)






			def 				A      (   self							:						Tuple			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs_for_decoder()

									UpperCamelCase__         = None

									self.model_tester.create_and_check_model_as_decoder(lowercase__ ,      lowercase__ ,      lowercase__			)






			def 				A      (   self							:						Tuple			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_decoder_model_past_large_inputs(lowercase__ ,      lowercase__ ,      lowercase__			)






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

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_causal_lm(*lowercase__			)






			def 				A      (   self							:						List[Any]			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase__			)






			def 				A      (   self							:						Tuple			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase__			)






			def 				A      (   self							:						int			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*lowercase__			)






			@unittest.skip(reason="""Feed forward chunking is not implemented"""			)
			def 				A      (   self							:						Optional[int]			)							->    str:

									'''simple docstring'''







									pass







			@parameterized.expand([("""linear""",), ("""dynamic""",)]			)
			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						Dict			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs_for_common()
									UpperCamelCase__         = ids_tensor([1, 1_0] ,      config.vocab_size			)
									UpperCamelCase__         = ids_tensor([1, int(config.max_position_embeddings * 1.5			)] ,      config.vocab_size			)

									set_seed(4_2			)  # Fixed seed at init time so the two models get the same random weights
									UpperCamelCase__         = GPTNeoXModel(lowercase__			)
									original_model.to(lowercase__			)
									original_model.eval()
									UpperCamelCase__         = original_model(lowercase__			).last_hidden_state
									UpperCamelCase__         = original_model(lowercase__			).last_hidden_state

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

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




@require_torch
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			@slow
			def 				A      (   self							:						Dict			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""			)
									for checkpointing in [True, False]:
															UpperCamelCase__         = GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""			)

															if checkpointing:
																					model.gradient_checkpointing_enable()
															else:
																					model.gradient_checkpointing_disable()
															model.to(lowercase__			)

															UpperCamelCase__         = tokenizer("""My favorite food is""" ,      return_tensors="""pt"""			).to(lowercase__			)
															# The hub repo. is updated on 2023-04-04, resulting in poor outputs.
															# See: https://github.com/huggingface/transformers/pull/24193
															UpperCamelCase__         = '''My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'''

															UpperCamelCase__         = model.generate(**lowercase__ ,      do_sample=lowercase__ ,      max_new_tokens=2_0			)
															UpperCamelCase__         = tokenizer.batch_decode(lowercase__			)[0]

															self.assertEqual(lowercase__ ,      lowercase__			)








'''simple docstring'''


from typing import TYPE_CHECKING

from ..utils import _LazyModule


lowerCamelCase_				: List[Any]    	=							{
    '''config''': [
        '''EXTERNAL_DATA_FORMAT_SIZE_LIMIT''',
        '''OnnxConfig''',
        '''OnnxConfigWithPast''',
        '''OnnxSeq2SeqConfigWithPast''',
        '''PatchingSpec''',
    ],
    '''convert''': ['''export''', '''validate_model_outputs'''],
    '''features''': ['''FeaturesManager'''],
    '''utils''': ['''ParameterFormat''', '''compute_serialized_parameters_size'''],
}


if TYPE_CHECKING:
					from .config import (
					    EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
					    OnnxConfig,
					    OnnxConfigWithPast,
					    OnnxSeqaSeqConfigWithPast,
					    PatchingSpec,
					)
					from .convert import export, validate_model_outputs
					from .features import FeaturesManager
					from .utils import ParameterFormat, compute_serialized_parameters_size

else:
					import sys

					lowerCamelCase_				: List[Any]    	=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
