

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 	__magic_name__    (			lowerCAmelCase_      ):





			def __init__(		self					,  __snake_case=0.01					,  __snake_case=1000					)   -> Optional[int]:




						'''simple docstring'''
						__a        =p_stop
						__a        =max_length





			def __iter__(		self					)   -> Any:




						'''simple docstring'''
						__a        =0
						__a        =False
						while not stop and count < self.max_length:
									yield count
									count += 1
									__a        =random.random() < self.p_stop






class 	__magic_name__    (			unittest.TestCase      ):





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case=False					,  __snake_case=True					)   -> Tuple:




						'''simple docstring'''
						__a        =[
						    BatchSamplerShard(__snake_case					,  2					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)
						    for i in range(2					)
						]
						__a        =[list(__snake_case					) for batch_sampler_shard in batch_sampler_shards]
						if not split_batches:
									self.assertListEqual([len(__snake_case					) for shard in batch_sampler_shards]					,  [len(__snake_case					) for e in expected]					)
						self.assertListEqual(__snake_case					,  __snake_case					)





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						# Check the shards when the dataset is a round multiple of total batch size.
						__a        =BatchSampler(range(24					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					)

						__a        =BatchSampler(range(24					)					,  batch_size=3					,  drop_last=__snake_case					)
						# Expected shouldn't change
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					)

						# Check the shards when the dataset is a round multiple of batch size but not total batch size.
						__a        =BatchSampler(range(21					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					)

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

						# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
						# num_processes batch.
						__a        =BatchSampler(range(22					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					)

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

						# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
						# num_processes batch.
						__a        =BatchSampler(range(20					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					)

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

						# Check the shards when the dataset is very small.
						__a        =BatchSampler(range(2					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[[[0, 1, 0]], [[1, 0, 1]]]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					)

						__a        =BatchSampler(range(2					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[[], []]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					)





			def 		__magic_name__    (		self					)   -> str:




						'''simple docstring'''
						# Check the shards when the dataset is a round multiple of batch size.
						__a        =BatchSampler(range(24					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						__a        =BatchSampler(range(24					)					,  batch_size=4					,  drop_last=__snake_case					)
						# Expected shouldn't change
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						# Check the shards when the dataset is not a round multiple of batch size.
						__a        =BatchSampler(range(22					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						__a        =BatchSampler(range(22					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						# Check the shards when the dataset is not a round multiple of batch size or num_processes.
						__a        =BatchSampler(range(21					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						__a        =BatchSampler(range(21					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						# Check the shards when the dataset is very small.
						__a        =BatchSampler(range(2					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[[[0, 1]], [[0, 1]]]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  split_batches=__snake_case					)

						__a        =BatchSampler(range(2					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[[], []]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  split_batches=__snake_case					)





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						# Check the shards when the dataset is a round multiple of total batch size.
						__a        =BatchSampler(range(24					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

						__a        =BatchSampler(range(24					)					,  batch_size=3					,  drop_last=__snake_case					)
						# Expected shouldn't change
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

						# Check the shards when the dataset is a round multiple of batch size but not total batch size.
						__a        =BatchSampler(range(21					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

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

						# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
						# num_processes batch.
						__a        =BatchSampler(range(22					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

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

						# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
						# num_processes batch.
						__a        =BatchSampler(range(20					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

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

						# Check the shards when the dataset is very small.
						__a        =BatchSampler(range(2					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[[[0, 1]], []]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  even_batches=__snake_case					)

						__a        =BatchSampler(range(2					)					,  batch_size=3					,  drop_last=__snake_case					)
						__a        =[[], []]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  even_batches=__snake_case					)





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						# Check the shards when the dataset is a round multiple of batch size.
						__a        =BatchSampler(range(24					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

						__a        =BatchSampler(range(24					)					,  batch_size=4					,  drop_last=__snake_case					)
						# Expected shouldn't change
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

						# Check the shards when the dataset is not a round multiple of batch size.
						__a        =BatchSampler(range(22					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

						__a        =BatchSampler(range(22					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

						# Check the shards when the dataset is not a round multiple of batch size or num_processes.
						__a        =BatchSampler(range(21					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

						__a        =BatchSampler(range(21					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[
						    [[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(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)

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

						__a        =BatchSampler(range(2					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =[[], []]
						self.check_batch_sampler_shards(__snake_case					,  __snake_case					,  split_batches=__snake_case					,  even_batches=__snake_case					)





			def 		__magic_name__    (		self					)   -> List[str]:




						'''simple docstring'''
						__a        =[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
						__a        =[BatchSamplerShard(__snake_case					,  2					,  __snake_case					,  even_batches=__snake_case					) 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 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case=False					,  __snake_case=2					,  __snake_case=False					)   -> str:




						'''simple docstring'''
						random.seed(__snake_case					)
						__a        =list(__snake_case					)

						__a        =[
						    IterableDatasetShard(
						        __snake_case					,  batch_size=__snake_case					,  drop_last=__snake_case					,  num_processes=__snake_case					,  process_index=__snake_case					,  split_batches=__snake_case					,  )
						    for i in range(__snake_case					)
						]
						__a        =[]
						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(__snake_case					)
									iterable_dataset_lists.append(list(__snake_case					)					)

						__a        =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
						__a        =iterable_dataset_lists[0]
						for l in iterable_dataset_lists[1:]:
									self.assertEqual(len(__snake_case					)					,  len(__snake_case					)					)
									self.assertTrue(len(__snake_case					) % shard_batch_size == 0					)

						__a        =[]
						for idx in range(0					,  len(__snake_case					)					,  __snake_case					):
									for l in iterable_dataset_lists:
												observed += l[idx : idx + shard_batch_size]

						if not drop_last:
									while len(__snake_case					) < len(__snake_case					):
												reference += reference
						self.assertListEqual(__snake_case					,  reference[: len(__snake_case					)]					)





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						__a        =42
						__a        =RandomIterableDataset()

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

						# Edge case with a very small dataset
						__a        =RandomIterableDataset(max_length=2					)

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





			def 		__magic_name__    (		self					)   -> str:




						'''simple docstring'''
						__a        =BatchSampler(range(16					)					,  batch_size=4					,  drop_last=__snake_case					)
						__a        =SkipBatchSampler(__snake_case					,  2					)
						self.assertListEqual(list(__snake_case					)					,  [[8, 9, 10, 11], [12, 13, 14, 15]]					)





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




						'''simple docstring'''
						__a        =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 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						__a        =DataLoader(list(range(16					)					)					,  batch_size=4					)
						__a        =skip_first_batches(__snake_case					,  num_batches=2					)
						self.assertListEqual([t.tolist() for t in new_dataloader]					,  [[8, 9, 10, 11], [12, 13, 14, 15]]					)





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						__a        =DataLoaderShard(list(range(16					)					)					,  batch_size=4					)
						for idx, _ in enumerate(__snake_case					):
									self.assertEqual(dataloader.end_of_dataloader					,  idx == 3					)

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





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						Accelerator()
						__a        =DataLoaderDispatcher(range(16					)					,  batch_size=4					)
						for idx, _ in enumerate(__snake_case					):
									self.assertEqual(dataloader.end_of_dataloader					,  idx == 3					)

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









from ... import PretrainedConfig


_lowerCAmelCase					:			Any  				=			{
    "sijunhe/nezha-cn-base": "https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json",
}






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
			SCREAMING_SNAKE_CASE							   =						'nezha'





			def __init__(		self					,  __snake_case=2_1128					,  __snake_case=768					,  __snake_case=12					,  __snake_case=12					,  __snake_case=3072					,  __snake_case="gelu"					,  __snake_case=0.1					,  __snake_case=0.1					,  __snake_case=512					,  __snake_case=64					,  __snake_case=2					,  __snake_case=0.02					,  __snake_case=1e-12					,  __snake_case=0.1					,  __snake_case=0					,  __snake_case=2					,  __snake_case=3					,  __snake_case=True					,  **__snake_case					,  )   -> str:




						'''simple docstring'''
						super().__init__(pad_token_id=__snake_case					,  bos_token_id=__snake_case					,  eos_token_id=__snake_case					,  **__snake_case					)

						__a        =vocab_size
						__a        =hidden_size
						__a        =num_hidden_layers
						__a        =num_attention_heads
						__a        =hidden_act
						__a        =intermediate_size
						__a        =hidden_dropout_prob
						__a        =attention_probs_dropout_prob
						__a        =max_position_embeddings
						__a        =max_relative_position
						__a        =type_vocab_size
						__a        =initializer_range
						__a        =layer_norm_eps
						__a        =classifier_dropout
						__a        =use_cache






