




'''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      A       :







							def __init__(    self			:    List[Any]				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    Union[str, Any]=13				,     lowerCAmelCase_			:    Union[str, Any]=7				,     lowerCAmelCase_			:    str=True				,     lowerCAmelCase_			:    Any=True				,     lowerCAmelCase_			:    int=True				,     lowerCAmelCase_			:    Union[str, Any]=True				,     lowerCAmelCase_			:    Any=99				,     lowerCAmelCase_			:    Dict=64				,     lowerCAmelCase_			:    Union[str, Any]=5				,     lowerCAmelCase_			:    List[Any]=4				,     lowerCAmelCase_			:    Tuple=37				,     lowerCAmelCase_			:    Optional[int]="gelu"				,     lowerCAmelCase_			:    List[Any]=0.1				,     lowerCAmelCase_			:    Optional[Any]=0.1				,     lowerCAmelCase_			:    int=5_12				,     lowerCAmelCase_			:    List[Any]=16				,     lowerCAmelCase_			:    Any=2				,     lowerCAmelCase_			:    Optional[int]=0.0_2				,     lowerCAmelCase_			:    List[str]=3				,     lowerCAmelCase_			:    Tuple=4				,     lowerCAmelCase_			:    Optional[int]=None				,     )     ->				Optional[int]:


														"""simple docstring"""


														_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      	=							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
														_a      	=							vocab_size - 1







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


														"""simple docstring"""


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

														_a      	=							self.get_config()

														return config, input_ids, input_mask, token_labels







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


														"""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=_lowerCAmelCase				,     initializer_range=self.initializer_range				,     pad_token_id=self.pad_token_id				,     )







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


														"""simple docstring"""


														_a      	=							self.prepare_config_and_inputs()

														_a      	=							True

														return config, input_ids, input_mask, token_labels







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


														"""simple docstring"""


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







							def   __lowerCAmelCase     (    self			:    Dict				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    int   )     ->				str:


														"""simple docstring"""


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







							def   __lowerCAmelCase     (    self			:    Tuple				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    Tuple   )     ->				Optional[Any]:


														"""simple docstring"""


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







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


														"""simple docstring"""


														_a      	=							self.num_labels
														_a      	=							GPTNeoXForQuestionAnswering(_lowerCAmelCase   )
														model.to(_lowerCAmelCase   )
														model.eval()
														_a      	=							model(_lowerCAmelCase				,     attention_mask=_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   __lowerCAmelCase     (    self			:    Tuple				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    Optional[int]   )     ->				str:


														"""simple docstring"""


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







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


														"""simple docstring"""


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







							def   __lowerCAmelCase     (    self			:    Optional[Any]				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Optional[int]   )     ->				Optional[Any]:


														"""simple docstring"""


														_a      	=							True
														_a      	=							GPTNeoXForCausalLM(config=_lowerCAmelCase   )
														model.to(_lowerCAmelCase   )
														model.eval()

														# first forward pass
														_a      	=							model(_lowerCAmelCase				,     attention_mask=_lowerCAmelCase				,     use_cache=_lowerCAmelCase   )
														_a      	=							outputs.past_key_values

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

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

														_a      	=							model(_lowerCAmelCase				,     attention_mask=_lowerCAmelCase				,     output_hidden_states=_lowerCAmelCase   )
														_a      	=							output_from_no_past['''hidden_states'''][0]
														_a      	=							model(
														    _lowerCAmelCase				,     attention_mask=_lowerCAmelCase				,     past_key_values=_lowerCAmelCase				,     output_hidden_states=_lowerCAmelCase				,     )['''hidden_states'''][0]

														# select random slice
														_a      	=							ids_tensor((1,)				,     output_from_past.shape[-1]   ).item()
														_a      	=							output_from_no_past[:, -3:, random_slice_idx].detach()
														_a      	=							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   __lowerCAmelCase     (    self			:    Any   )     ->				Optional[Any]:


														"""simple docstring"""


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



@require_torch
class      A       (						__lowercase	,__lowercase	,__lowercase	,unittest.TestCase       ):
							lowercase_							     =							(
							    (
							        GPTNeoXModel,
							        GPTNeoXForCausalLM,
							        GPTNeoXForQuestionAnswering,
							        GPTNeoXForSequenceClassification,
							        GPTNeoXForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase_							     =							(GPTNeoXForCausalLM,) if is_torch_available() else ()
							lowercase_							     =							(
							    {
							        '''feature-extraction''': GPTNeoXModel,
							        '''question-answering''': GPTNeoXForQuestionAnswering,
							        '''text-classification''': GPTNeoXForSequenceClassification,
							        '''text-generation''': GPTNeoXForCausalLM,
							        '''token-classification''': GPTNeoXForTokenClassification,
							        '''zero-shot''': GPTNeoXForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowercase_							     =							False
							lowercase_							     =							False
							lowercase_							     =							False
							lowercase_							     =							False







							def   __lowerCAmelCase     (    self			:    Tuple   )     ->				int:


														"""simple docstring"""


														_a      	=							GPTNeoXModelTester(self   )
														_a      	=							ConfigTester(self				,     config_class=_lowerCAmelCase				,     hidden_size=64				,     num_attention_heads=8   )







							def   __lowerCAmelCase     (    self			:    Dict   )     ->				str:


														"""simple docstring"""


														self.config_tester.run_common_tests()







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(_lowerCAmelCase				,     _lowerCAmelCase				,     _lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs_for_decoder()
														self.model_tester.create_and_check_model_as_decoder(_lowerCAmelCase				,     _lowerCAmelCase				,     _lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs_for_decoder()

														_a      	=							None

														self.model_tester.create_and_check_model_as_decoder(_lowerCAmelCase				,     _lowerCAmelCase				,     _lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_decoder_model_past_large_inputs(_lowerCAmelCase				,     _lowerCAmelCase				,     _lowerCAmelCase   )







							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_causal_lm(*_lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase   )







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


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase   )







							@unittest.skip(reason='''Feed forward chunking is not implemented'''   )
							def   __lowerCAmelCase     (    self			:    List[str]   )     ->				List[str]:


														"""simple docstring"""


														pass







							@parameterized.expand([('''linear''',), ('''dynamic''',)]   )
							def   __lowerCAmelCase     (    self			:    List[Any]				,     lowerCAmelCase_			:    str   )     ->				Union[str, Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs_for_common()
														_a      	=							ids_tensor([1, 10]				,     config.vocab_size   )
														_a      	=							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
														_a      	=							GPTNeoXModel(_lowerCAmelCase   )
														original_model.to(_lowerCAmelCase   )
														original_model.eval()
														_a      	=							original_model(_lowerCAmelCase   ).last_hidden_state
														_a      	=							original_model(_lowerCAmelCase   ).last_hidden_state

														set_seed(42   )  # Fixed seed at init time so the two models get the same random weights
														_a      	=							{'''type''': scaling_type, '''factor''': 1_0.0}
														_a      	=							GPTNeoXModel(_lowerCAmelCase   )
														scaled_model.to(_lowerCAmelCase   )
														scaled_model.eval()
														_a      	=							scaled_model(_lowerCAmelCase   ).last_hidden_state
														_a      	=							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   )   )



@require_torch
class      A       (						unittest.TestCase       ):







							@slow
							def   __lowerCAmelCase     (    self			:    int   )     ->				Tuple:


														"""simple docstring"""


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

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

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

																					_a      	=							model.generate(**_lowerCAmelCase				,     do_sample=_lowerCAmelCase				,     max_new_tokens=20   )
																					_a      	=							tokenizer.batch_decode(_lowerCAmelCase   )[0]

																					self.assertEqual(_lowerCAmelCase				,     _lowerCAmelCase   )








from typing import Callable, Optional

from .. import Features
from ..packaged_modules.generator.generator import Generator
from .abc import AbstractDatasetInputStream





class 				a			(      __lowercase			):






				def __init__(  self			, _lowerCAmelCase			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = False			, _lowerCAmelCase = False			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										super().__init__(
										    features=_lowerCAmelCase			, cache_dir=_lowerCAmelCase			, keep_in_memory=_lowerCAmelCase			, streaming=_lowerCAmelCase			, num_proc=_lowerCAmelCase			, **_lowerCAmelCase			, )
										__SCREAMING_SNAKE_CASE:				Any           =   Generator(
										    cache_dir=_lowerCAmelCase			, features=_lowerCAmelCase			, generator=_lowerCAmelCase			, gen_kwargs=_lowerCAmelCase			, **_lowerCAmelCase			, )






				def 			snake_case_    (  self					):

										"""simple docstring"""




										if self.streaming:
																__SCREAMING_SNAKE_CASE:				List[str]           =   self.builder.as_streaming_dataset(split='''train'''					)
										# Build regular (map-style) dataset
										else:
																__SCREAMING_SNAKE_CASE:				str           =   None
																__SCREAMING_SNAKE_CASE:				List[Any]           =   None
																__SCREAMING_SNAKE_CASE:				Tuple           =   None
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   None

																self.builder.download_and_prepare(
																    download_config=_lowerCAmelCase			, download_mode=_lowerCAmelCase			, verification_mode=_lowerCAmelCase			, base_path=_lowerCAmelCase			, num_proc=self.num_proc			, )
																__SCREAMING_SNAKE_CASE:				List[str]           =   self.builder.as_dataset(
																    split='''train'''			, verification_mode=_lowerCAmelCase			, in_memory=self.keep_in_memory					)
										return dataset



