


import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)






class _lowerCamelCase(						_a   ):
       def __init__(					self,      lowerCamelCase=0.0_1,      lowerCamelCase=10_00)       ->   Dict:




           """simple docstring"""

           _lowercase       :							List[str]      		=		p_stop
           _lowercase       :							Tuple      		=		max_length




       def __iter__(					self)       ->   str:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		0
           _lowercase       :							Union[str, Any]      		=		False
           while not stop and count < self.max_length:
               yield count
               count += 1
               _lowercase       :							int      		=		random.random() < self.p_stop






class _lowerCamelCase(						unittest.TestCase   ):
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase=False,      lowerCamelCase=True)       ->   Optional[int]:




           """simple docstring"""

           _lowercase       :							Dict      		=		[
               BatchSamplerShard(lowerCamelCase,      2,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)
               for i in range(2)
           ]
           _lowercase       :							str      		=		[list(lowerCamelCase) for batch_sampler_shard in batch_sampler_shards]
           if not split_batches:
               self.assertListEqual([len(lowerCamelCase) for shard in batch_sampler_shards],      [len(lowerCamelCase) for e in expected])
           self.assertListEqual(lowerCamelCase,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							List[str]      		=		BatchSampler(range(24),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(24),      batch_size=3,      drop_last=lowerCamelCase)
           # Expected shouldn't change
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           # Check the shards when the dataset is a round multiple of batch size but not total batch size.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(21),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							List[Any]      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							Dict      		=		BatchSampler(range(21),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
           # num_processes batch.
           _lowercase       :							Any      		=		BatchSampler(range(22),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							List[Any]      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							List[str]      		=		BatchSampler(range(22),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
           # num_processes batch.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(20),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Dict      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(20),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Tuple      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           # Check the shards when the dataset is very small.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(2),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Dict      		=		[[[0, 1, 0]], [[1, 0, 1]]]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)

           _lowercase       :							List[str]      		=		BatchSampler(range(2),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Dict      		=		[[], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							str      		=		BatchSampler(range(24),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[Any]      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           _lowercase       :							List[str]      		=		BatchSampler(range(24),      batch_size=4,      drop_last=lowerCamelCase)
           # Expected shouldn't change
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(22),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							List[Any]      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           _lowercase       :							Optional[Any]      		=		BatchSampler(range(22),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Dict      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size or num_processes.
           _lowercase       :							Any      		=		BatchSampler(range(21),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           _lowercase       :							Tuple      		=		BatchSampler(range(21),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           # Check the shards when the dataset is very small.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(2),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[[[0, 1]], [[0, 1]]]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)

           _lowercase       :							Dict      		=		BatchSampler(range(2),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Dict      		=		[[], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							List[str]      		=		BatchSampler(range(24),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Any      		=		BatchSampler(range(24),      batch_size=3,      drop_last=lowerCamelCase)
           # Expected shouldn't change
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is a round multiple of batch size but not total batch size.
           _lowercase       :							Any      		=		BatchSampler(range(21),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							List[str]      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Optional[Any]      		=		BatchSampler(range(21),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Tuple      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
           # num_processes batch.
           _lowercase       :							str      		=		BatchSampler(range(22),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							str      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							str      		=		BatchSampler(range(22),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
           # num_processes batch.
           _lowercase       :							List[Any]      		=		BatchSampler(range(20),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							str      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Dict      		=		BatchSampler(range(20),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[
               [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
               [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is very small.
           _lowercase       :							Union[str, Any]      		=		BatchSampler(range(2),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[[[0, 1]], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Dict      		=		BatchSampler(range(2),      batch_size=3,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[[], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      even_batches=lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           _lowercase       :							str      		=		BatchSampler(range(24),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Any      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Optional[Any]      		=		BatchSampler(range(24),      batch_size=4,      drop_last=lowerCamelCase)
           # Expected shouldn't change
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size.
           _lowercase       :							int      		=		BatchSampler(range(22),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							str      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							List[str]      		=		BatchSampler(range(22),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							str      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is not a round multiple of batch size or num_processes.
           _lowercase       :							List[str]      		=		BatchSampler(range(21),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							Any      		=		BatchSampler(range(21),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[
               [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
               [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
           ]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           # Check the shards when the dataset is very small.
           _lowercase       :							Any      		=		BatchSampler(range(2),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							int      		=		[[[0, 1]], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)

           _lowercase       :							str      		=		BatchSampler(range(2),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		[[], []]
           self.check_batch_sampler_shards(lowerCamelCase,      lowerCamelCase,      split_batches=lowerCamelCase,      even_batches=lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
           _lowercase       :							List[str]      		=		[BatchSamplerShard(lowerCamelCase,      2,      lowerCamelCase,      even_batches=lowerCamelCase) for i in range(2)]

           self.assertEqual(len(batch_sampler_shards[0]),      3)
           self.assertEqual(len(batch_sampler_shards[1]),      2)

           self.assertListEqual(list(batch_sampler_shards[0]),      [[0, 1, 2], [5, 6, 7, 8], [12, 13]])
           self.assertListEqual(list(batch_sampler_shards[1]),      [[3, 4], [9, 10, 11]])
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase=False,      lowerCamelCase=2,      lowerCamelCase=False)       ->   str:




           """simple docstring"""

           random.seed(lowerCamelCase)
           _lowercase       :							List[str]      		=		list(lowerCamelCase)

           _lowercase       :							Optional[int]      		=		[
               IterableDatasetShard(
                   lowerCamelCase,      batch_size=lowerCamelCase,      drop_last=lowerCamelCase,      num_processes=lowerCamelCase,      process_index=lowerCamelCase,      split_batches=lowerCamelCase,      )
               for i in range(lowerCamelCase)
           ]
           _lowercase       :							Dict      		=		[]
           for iterable_dataset_shard in iterable_dataset_shards:
               # Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
               random.seed(lowerCamelCase)
               iterable_dataset_lists.append(list(lowerCamelCase))

           _lowercase       :							Tuple      		=		batch_size // num_processes if split_batches else batch_size
           # All iterable dataset shard should have the same length, a round multiple of shard_batch_size
           _lowercase       :							str      		=		iterable_dataset_lists[0]
           for l in iterable_dataset_lists[1:]:
               self.assertEqual(len(lowerCamelCase),      len(lowerCamelCase))
               self.assertTrue(len(lowerCamelCase) % shard_batch_size == 0)

           _lowercase       :							Tuple      		=		[]
           for idx in range(0,      len(lowerCamelCase),      lowerCamelCase):
               for l in iterable_dataset_lists:
                   observed += l[idx : idx + shard_batch_size]

           if not drop_last:
               while len(lowerCamelCase) < len(lowerCamelCase):
                   reference += reference
           self.assertListEqual(lowerCamelCase,      reference[: len(lowerCamelCase)])
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		42
           _lowercase       :							int      		=		RandomIterableDataset()

           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)

           # Edge case with a very small dataset
           _lowercase       :							Union[str, Any]      		=		RandomIterableDataset(max_length=2)

           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
           self.check_iterable_dataset_shards(lowerCamelCase,      lowerCamelCase,      batch_size=4,      drop_last=lowerCamelCase,      split_batches=lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		BatchSampler(range(16),      batch_size=4,      drop_last=lowerCamelCase)
           _lowercase       :							str      		=		SkipBatchSampler(lowerCamelCase,      2)
           self.assertListEqual(list(lowerCamelCase),      [[8, 9, 10, 11], [12, 13, 14, 15]])
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		SkipDataLoader(list(range(16)),      batch_size=4,      skip_batches=2)
           self.assertListEqual([t.tolist() for t in dataloader],      [[8, 9, 10, 11], [12, 13, 14, 15]])
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		DataLoader(list(range(16)),      batch_size=4)
           _lowercase       :							Optional[int]      		=		skip_first_batches(lowerCamelCase,      num_batches=2)
           self.assertListEqual([t.tolist() for t in new_dataloader],      [[8, 9, 10, 11], [12, 13, 14, 15]])
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		DataLoaderShard(list(range(16)),      batch_size=4)
           for idx, _ in enumerate(lowerCamelCase):
               self.assertEqual(dataloader.end_of_dataloader,      idx == 3)

           # Test it also works on the second iteration
           for idx, _ in enumerate(lowerCamelCase):
               self.assertEqual(dataloader.end_of_dataloader,      idx == 3)




       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           Accelerator()
           _lowercase       :							Optional[Any]      		=		DataLoaderDispatcher(range(16),      batch_size=4)
           for idx, _ in enumerate(lowerCamelCase):
               self.assertEqual(dataloader.end_of_dataloader,      idx == 3)

           # Test it also works on the second iteration
           for idx, _ in enumerate(lowerCamelCase):
               self.assertEqual(dataloader.end_of_dataloader,      idx == 3)







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





	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.0_2		,		A=3		,		A=4		,		A=None		,		)			->			Optional[int]:
								UpperCAmelCase     :		List[Any]							     =		parent
								UpperCAmelCase     :		Optional[int]							     =		batch_size
								UpperCAmelCase     :		Union[str, Any]							     =		seq_length
								UpperCAmelCase     :		Optional[Any]							     =		is_training
								UpperCAmelCase     :		Dict							     =		use_input_mask
								UpperCAmelCase     :		str							     =		use_token_type_ids
								UpperCAmelCase     :		List[Any]							     =		use_labels
								UpperCAmelCase     :		List[Any]							     =		vocab_size
								UpperCAmelCase     :		Dict							     =		hidden_size
								UpperCAmelCase     :		Dict							     =		num_hidden_layers
								UpperCAmelCase     :		Optional[int]							     =		num_attention_heads
								UpperCAmelCase     :		int							     =		intermediate_size
								UpperCAmelCase     :		List[str]							     =		hidden_act
								UpperCAmelCase     :		List[str]							     =		hidden_dropout_prob
								UpperCAmelCase     :		int							     =		attention_probs_dropout_prob
								UpperCAmelCase     :		str							     =		max_position_embeddings
								UpperCAmelCase     :		Optional[Any]							     =		type_vocab_size
								UpperCAmelCase     :		List[str]							     =		type_sequence_label_size
								UpperCAmelCase     :		int							     =		initializer_range
								UpperCAmelCase     :		str							     =		num_labels
								UpperCAmelCase     :		Optional[int]							     =		num_choices
								UpperCAmelCase     :		Dict							     =		scope
								UpperCAmelCase     :		Union[str, Any]							     =		vocab_size - 1





	def  _lowercase(  self    )			->			Union[str, Any]:
								UpperCAmelCase     :		List[str]							     =		ids_tensor([self.batch_size, self.seq_length]		,		self.vocab_size    )

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

								UpperCAmelCase     :		List[str]							     =		None
								if self.use_labels:
															UpperCAmelCase     :		Optional[int]							     =		ids_tensor([self.batch_size, self.seq_length]		,		self.num_labels    )

								UpperCAmelCase     :		Optional[int]							     =		self.get_config()

								return config, input_ids, input_mask, token_labels





	def  _lowercase(  self    )			->			Optional[Any]:
								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  _lowercase(  self    )			->			Optional[Any]:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		self.prepare_config_and_inputs()

								UpperCAmelCase     :		Any							     =		True

								return config, input_ids, input_mask, token_labels





	def  _lowercase(  self		,		A		,		A		,		A    )			->			int:
								UpperCAmelCase     :		str							     =		GPTNeoXModel(config=A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[str]							     =		model(A		,		attention_mask=A    )
								UpperCAmelCase     :		List[str]							     =		model(A    )
								self.parent.assertEqual(result.last_hidden_state.shape		,		(self.batch_size, self.seq_length, self.hidden_size)    )





	def  _lowercase(  self		,		A		,		A		,		A    )			->			Optional[int]:
								UpperCAmelCase     :		str							     =		True
								UpperCAmelCase     :		Optional[Any]							     =		GPTNeoXModel(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[Any]							     =		model(A		,		attention_mask=A    )
								self.parent.assertEqual(result.last_hidden_state.shape		,		(self.batch_size, self.seq_length, self.hidden_size)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			List[str]:
								UpperCAmelCase     :		Tuple							     =		GPTNeoXForCausalLM(config=A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		str							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.seq_length, self.vocab_size)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			Tuple:
								UpperCAmelCase     :		List[str]							     =		self.num_labels
								UpperCAmelCase     :		Any							     =		GPTNeoXForQuestionAnswering(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		str							     =		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  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			int:
								UpperCAmelCase     :		Tuple							     =		self.num_labels
								UpperCAmelCase     :		List[str]							     =		GPTNeoXForSequenceClassification(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		List[Any]							     =		ids_tensor([self.batch_size]		,		self.type_sequence_label_size    )
								UpperCAmelCase     :		Tuple							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.num_labels)    )





	def  _lowercase(  self		,		A		,		A		,		A		,		A    )			->			str:
								UpperCAmelCase     :		List[Any]							     =		self.num_labels
								UpperCAmelCase     :		Tuple							     =		GPTNeoXForTokenClassification(A    )
								model.to(A    )
								model.eval()
								UpperCAmelCase     :		int							     =		model(A		,		attention_mask=A		,		labels=A    )
								self.parent.assertEqual(result.logits.shape		,		(self.batch_size, self.seq_length, self.num_labels)    )





	def  _lowercase(  self		,		A		,		A		,		A    )			->			Union[str, Any]:
								UpperCAmelCase     :		Optional[int]							     =		True
								UpperCAmelCase     :		str							     =		GPTNeoXForCausalLM(config=A    )
								model.to(A    )
								model.eval()

								# first forward pass
								UpperCAmelCase     :		List[str]							     =		model(A		,		attention_mask=A		,		use_cache=A    )
								UpperCAmelCase     :		List[Any]							     =		outputs.past_key_values

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

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

								UpperCAmelCase     :		Dict							     =		model(A		,		attention_mask=A		,		output_hidden_states=A    )
								UpperCAmelCase     :		Any							     =		output_from_no_past["""hidden_states"""][0]
								UpperCAmelCase     :		List[str]							     =		model(
								    A		,		attention_mask=A		,		past_key_values=A		,		output_hidden_states=A		,		)["""hidden_states"""][0]

								# select random slice
								UpperCAmelCase     :		Tuple							     =		ids_tensor((1,)		,		output_from_past.shape[-1]    ).item()
								UpperCAmelCase     :		List[Any]							     =		output_from_no_past[:, -3:, random_slice_idx].detach()
								UpperCAmelCase     :		List[str]							     =		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  _lowercase(  self    )			->			int:
								UpperCAmelCase     :		Tuple							     =		self.prepare_config_and_inputs()
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Optional[Any]							     =		config_and_inputs
								UpperCAmelCase     :		Union[str, Any]							     =		{"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_torch
class 			UpperCamelCase_ (							__magic_name__      ,     __magic_name__      ,     __magic_name__      ,     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  _lowercase(  self    )			->			Union[str, Any]:
								UpperCAmelCase     :		str							     =		GPTNeoXModelTester(self    )
								UpperCAmelCase     :		Optional[Any]							     =		ConfigTester(self		,		config_class=A		,		hidden_size=64		,		num_attention_heads=8    )





	def  _lowercase(  self    )			->			Optional[Any]:
								self.config_tester.run_common_tests()





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(A		,		A		,		A    )





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Union[str, Any]							     =		self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(A		,		A		,		A    )





	def  _lowercase(  self    )			->			Optional[Any]:
								# This regression test was failing with PyTorch < 1.3
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		Optional[Any]							     =		self.model_tester.prepare_config_and_inputs_for_decoder()

								UpperCAmelCase     :		Optional[Any]							     =		None

								self.model_tester.create_and_check_model_as_decoder(A		,		A		,		A    )





	def  _lowercase(  self    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     ,     UpperCAmelCase     :		List[str]							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_decoder_model_past_large_inputs(A		,		A		,		A    )





	def  _lowercase(  self    )			->			int:
								UpperCAmelCase     :		Any							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_causal_lm(*A    )





	def  _lowercase(  self    )			->			Optional[int]:
								UpperCAmelCase     :		Tuple							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*A    )





	def  _lowercase(  self    )			->			Any:
								UpperCAmelCase     :		Tuple							     =		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*A    )





	def  _lowercase(  self    )			->			Optional[Any]:
								UpperCAmelCase     :		Optional[Any]							     =		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  _lowercase(  self    )			->			Optional[int]:
								pass






	@parameterized.expand([("""linear""",), ("""dynamic""",)]    )
	def  _lowercase(  self		,		A    )			->			str:
								UpperCAmelCase     ,     UpperCAmelCase     :		List[Any]							     =		self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase     :		int							     =		ids_tensor([1, 10]		,		config.vocab_size    )
								UpperCAmelCase     :		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
								UpperCAmelCase     :		Dict							     =		GPTNeoXModel(A    )
								original_model.to(A    )
								original_model.eval()
								UpperCAmelCase     :		List[str]							     =		original_model(A    ).last_hidden_state
								UpperCAmelCase     :		Any							     =		original_model(A    ).last_hidden_state

								set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
								UpperCAmelCase     :		Any							     =		{"""type""": scaling_type, """factor""": 1_0.0}
								UpperCAmelCase     :		str							     =		GPTNeoXModel(A    )
								scaled_model.to(A    )
								scaled_model.eval()
								UpperCAmelCase     :		Optional[Any]							     =		scaled_model(A    ).last_hidden_state
								UpperCAmelCase     :		Optional[Any]							     =		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 			UpperCamelCase_ (							unittest.TestCase			):





	@slow
	def  _lowercase(  self    )			->			List[Any]:
								UpperCAmelCase     :		str							     =		AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )
								for checkpointing in [True, False]:
															UpperCAmelCase     :		int							     =		GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""    )

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

															UpperCAmelCase     :		List[Any]							     =		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
															UpperCAmelCase     :		List[str]							     =		"""My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"""

															UpperCAmelCase     :		Union[str, Any]							     =		model.generate(**A		,		do_sample=A		,		max_new_tokens=20    )
															UpperCAmelCase     :		Tuple							     =		tokenizer.batch_decode(A    )[0]

															self.assertEqual(A		,		A    )

