




"""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 					__magic_name__							:






				'''simple docstring'''





				def __init__( self	,				_a	,				_a=13	,				_a=7	,				_a=True	,				_a=True	,				_a=True	,				_a=True	,				_a=99	,				_a=64	,				_a=5	,				_a=4	,				_a=37	,				_a="gelu"	,				_a=0.1	,				_a=0.1	,				_a=512	,				_a=16	,				_a=2	,				_a=0.02	,				_a=3	,				_a=4	,				_a=None	,				):





							"""simple docstring"""
							lowerCamelCase       =      parent
							lowerCamelCase       =      batch_size
							lowerCamelCase       =      seq_length
							lowerCamelCase       =      is_training
							lowerCamelCase       =      use_input_mask
							lowerCamelCase       =      use_token_type_ids
							lowerCamelCase       =      use_labels
							lowerCamelCase       =      vocab_size
							lowerCamelCase       =      hidden_size
							lowerCamelCase       =      num_hidden_layers
							lowerCamelCase       =      num_attention_heads
							lowerCamelCase       =      intermediate_size
							lowerCamelCase       =      hidden_act
							lowerCamelCase       =      hidden_dropout_prob
							lowerCamelCase       =      attention_probs_dropout_prob
							lowerCamelCase       =      max_position_embeddings
							lowerCamelCase       =      type_vocab_size
							lowerCamelCase       =      type_sequence_label_size
							lowerCamelCase       =      initializer_range
							lowerCamelCase       =      num_labels
							lowerCamelCase       =      num_choices
							lowerCamelCase       =      scope
							lowerCamelCase       =      vocab_size - 1





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      ids_tensor([self.batch_size, self.seq_length]	,				self.vocab_size				)

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

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

							lowerCamelCase       =      self.get_config()

							return config, input_ids, input_mask, token_labels





				def      _lowerCAmelCase       ( self				):





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





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self.prepare_config_and_inputs()

							lowerCamelCase       =      True

							return config, input_ids, input_mask, token_labels





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.num_labels
							lowerCamelCase       =      GPTNeoXForQuestionAnswering(_a				)
							model.to(_a				)
							model.eval()
							lowerCamelCase       =      model(_a	,				attention_mask=_a				)
							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	,				_a	,				_a	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a	,				_a				):





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





				def      _lowerCAmelCase       ( self	,				_a	,				_a	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      True
							lowerCamelCase       =      GPTNeoXForCausalLM(config=_a				)
							model.to(_a				)
							model.eval()

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

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

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

							lowerCamelCase       =      model(_a	,				attention_mask=_a	,				output_hidden_states=_a				)
							lowerCamelCase       =      output_from_no_past["""hidden_states"""][0]
							lowerCamelCase       =      model(
							    _a	,				attention_mask=_a	,				past_key_values=_a	,				output_hidden_states=_a	,				)["""hidden_states"""][0]

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





							"""simple docstring"""
							lowerCamelCase       =      self.prepare_config_and_inputs()
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      config_and_inputs
							lowerCamelCase       =      {"""input_ids""": input_ids, """attention_mask""": input_mask}
							return config, inputs_dict






@require_torch
class 					__magic_name__							( UpperCAmelCase__		,      UpperCAmelCase__		,      UpperCAmelCase__		,      unittest.TestCase       ):






				'''simple docstring'''
				__UpperCamelCase    						=		(
				    (
				        GPTNeoXModel,
				        GPTNeoXForCausalLM,
				        GPTNeoXForQuestionAnswering,
				        GPTNeoXForSequenceClassification,
				        GPTNeoXForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__UpperCamelCase    						=		(GPTNeoXForCausalLM,) if is_torch_available() else ()
				__UpperCamelCase    						=		(
				    {
				        "feature-extraction": GPTNeoXModel,
				        "question-answering": GPTNeoXForQuestionAnswering,
				        "text-classification": GPTNeoXForSequenceClassification,
				        "text-generation": GPTNeoXForCausalLM,
				        "token-classification": GPTNeoXForTokenClassification,
				        "zero-shot": GPTNeoXForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		False





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      GPTNeoXModelTester(self				)
							lowerCamelCase       =      ConfigTester(self	,				config_class=_a	,				hidden_size=64	,				num_attention_heads=8				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							self.config_tester.run_common_tests()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# This regression test was failing with PyTorch < 1.3
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs_for_decoder()

							lowerCamelCase       =      None

							self.model_tester.create_and_check_model_as_decoder(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_decoder_model_past_large_inputs(_a	,				_a	,				_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_causal_lm(*_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*_a				)





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*_a				)





				@unittest.skip(reason="""Feed forward chunking is not implemented"""				)
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							pass





				@parameterized.expand([("""linear""",), ("""dynamic""",)]				)
				def      _lowerCAmelCase       ( self	,				_a				):





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

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






@require_torch
class 					__magic_name__							( unittest.TestCase       ):






				'''simple docstring'''





				@slow
				def      _lowerCAmelCase       ( self				):





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

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

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

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

										self.assertEqual(_a	,				_a				)









"""simple docstring"""




from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


lowerCAmelCase							:    List[Any]  							=							logging.get_logger(__name__)


lowerCAmelCase							:    int  							=							TypeVar("""DatasetType""", Dataset, IterableDataset)



def 					a__				(       snake_case__ ,				snake_case__ = None ,				snake_case__ = None ,				snake_case__ = None ,				snake_case__ = None ,				snake_case__ = "first_exhausted" ,				)		->      DatasetType:
			from .arrow_dataset import Dataset
			from .iterable_dataset import IterableDataset

			if not datasets:
						raise ValueError("""Unable to interleave an empty list of datasets."""       )
			for i, dataset in enumerate(snake_case__       ):
						if not isinstance(snake_case__ ,				(Dataset, IterableDataset)       ):
									if isinstance(snake_case__ ,				(DatasetDict, IterableDatasetDict)       ):
												if not dataset:
															raise ValueError(
															    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} '
															    """is an empty dataset dictionary."""       )
												raise ValueError(
												    F'Dataset at position {i} has at least one split: {list(snake_case__       )}\n'
												    F'Please pick one to interleave with the other datasets, for example: dataset[\'{next(iter(snake_case__       )       )}\']'       )
									raise ValueError(
									    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(snake_case__       ).__name__}.'       )
						if i == 0:
									lowerCamelCase   ,     lowerCamelCase       =      (
									    (Dataset, IterableDataset) if isinstance(snake_case__ ,				snake_case__       ) else (IterableDataset, Dataset)
									)
						elif not isinstance(snake_case__ ,				snake_case__       ):
									raise ValueError(
									    F'Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects.'       )
			if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
						raise ValueError(F'{stopping_strategy} is not supported. Please enter a valid stopping_strategy.'       )
			if dataset_type is Dataset:
						return _interleave_map_style_datasets(
						    snake_case__ ,				snake_case__ ,				snake_case__ ,				info=snake_case__ ,				split=snake_case__ ,				stopping_strategy=snake_case__       )
			else:
						return _interleave_iterable_datasets(
						    snake_case__ ,				snake_case__ ,				snake_case__ ,				info=snake_case__ ,				split=snake_case__ ,				stopping_strategy=snake_case__       )



def 					a__				(       snake_case__ ,				snake_case__ = None ,				snake_case__ = None ,				snake_case__ = 0 ,				)		->      DatasetType:

			if not dsets:
						raise ValueError("""Unable to concatenate an empty list of datasets."""       )
			for i, dataset in enumerate(snake_case__       ):
						if not isinstance(snake_case__ ,				(Dataset, IterableDataset)       ):
									if isinstance(snake_case__ ,				(DatasetDict, IterableDatasetDict)       ):
												if not dataset:
															raise ValueError(
															    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} '
															    """is an empty dataset dictionary."""       )
												raise ValueError(
												    F'Dataset at position {i} has at least one split: {list(snake_case__       )}\n'
												    F'Please pick one to interleave with the other datasets, for example: dataset[\'{next(iter(snake_case__       )       )}\']'       )
									raise ValueError(
									    F'Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(snake_case__       ).__name__}.'       )
						if i == 0:
									lowerCamelCase   ,     lowerCamelCase       =      (
									    (Dataset, IterableDataset) if isinstance(snake_case__ ,				snake_case__       ) else (IterableDataset, Dataset)
									)
						elif not isinstance(snake_case__ ,				snake_case__       ):
									raise ValueError(
									    F'Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects.'       )
			if dataset_type is Dataset:
						return _concatenate_map_style_datasets(snake_case__ ,				info=snake_case__ ,				split=snake_case__ ,				axis=snake_case__       )
			else:
						return _concatenate_iterable_datasets(snake_case__ ,				info=snake_case__ ,				split=snake_case__ ,				axis=snake_case__       )



