




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 	_UpperCamelCase					( UpperCAmelCase__						):
      def __init__(	self      :Dict       ,					lowerCamelCase      :Optional[int]=0.01       ,					lowerCamelCase      :Optional[Any]=1000					)   ->	Optional[int]:
       UpperCAmelCase__							       =  p_stop
       UpperCAmelCase__							       =  max_length







      def __iter__(	self      :List[Any]					)   ->	List[Any]:
       UpperCAmelCase__							       =  0
       UpperCAmelCase__							       =  False
       while not stop and count < self.max_length:
        yield count
        count += 1
        UpperCAmelCase__							       =  random.random() < self.p_stop



class 	_UpperCamelCase					( unittest.TestCase						):
      def 		UpperCAmelCase_  (	self      :List[str]       ,					lowerCamelCase      :Dict       ,					lowerCamelCase      :List[str]       ,					lowerCamelCase      :List[str]=False       ,					lowerCamelCase      :Any=True					)   ->	Optional[Any]:
       UpperCAmelCase__							       =  [
           BatchSamplerShard(UpperCAmelCase__       ,					2       ,					UpperCAmelCase__       ,					split_batches=UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					)
           for i in range(2					)
       ]
       UpperCAmelCase__							       =  [list(UpperCAmelCase__					) for batch_sampler_shard in batch_sampler_shards]
       if not split_batches:
        self.assertListEqual([len(UpperCAmelCase__					) for shard in batch_sampler_shards]       ,					[len(UpperCAmelCase__					) for e in expected]					)
       self.assertListEqual(UpperCAmelCase__       ,					UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Dict					)   ->	Any:
       UpperCAmelCase__							       =  BatchSampler(range(24					)       ,					batch_size=3       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [
           [[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(UpperCAmelCase__       ,					UpperCAmelCase__					)

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

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

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

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

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

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

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

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

       UpperCAmelCase__							       =  BatchSampler(range(2					)       ,					batch_size=3       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [[], []]
       self.check_batch_sampler_shards(UpperCAmelCase__       ,					UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Optional[Any]					)   ->	Optional[Any]:
       UpperCAmelCase__							       =  BatchSampler(range(24					)       ,					batch_size=4       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [
           [[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(UpperCAmelCase__       ,					UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)

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

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

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

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

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

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

       UpperCAmelCase__							       =  BatchSampler(range(2					)       ,					batch_size=4       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [[], []]
       self.check_batch_sampler_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	Union[str, Any]:
       UpperCAmelCase__							       =  BatchSampler(range(24					)       ,					batch_size=3       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [
           [[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(UpperCAmelCase__       ,					UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					)

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

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

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

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

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

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

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

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

       UpperCAmelCase__							       =  BatchSampler(range(2					)       ,					batch_size=3       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [[], []]
       self.check_batch_sampler_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Optional[Any]					)   ->	Tuple:
       UpperCAmelCase__							       =  BatchSampler(range(24					)       ,					batch_size=4       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [
           [[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(UpperCAmelCase__       ,					UpperCAmelCase__       ,					split_batches=UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					)

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

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

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

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

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

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

       UpperCAmelCase__							       =  BatchSampler(range(2					)       ,					batch_size=4       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  [[], []]
       self.check_batch_sampler_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					split_batches=UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Union[str, Any]					)   ->	Any:
       UpperCAmelCase__							       =  [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       UpperCAmelCase__							       =  [BatchSamplerShard(UpperCAmelCase__       ,					2       ,					UpperCAmelCase__       ,					even_batches=UpperCAmelCase__					) 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      :int       ,					lowerCamelCase      :List[Any]       ,					lowerCamelCase      :Optional[int]       ,					lowerCamelCase      :Tuple       ,					lowerCamelCase      :Tuple=False       ,					lowerCamelCase      :Optional[Any]=2       ,					lowerCamelCase      :List[str]=False					)   ->	Optional[Any]:
       random.seed(UpperCAmelCase__					)
       UpperCAmelCase__							       =  list(UpperCAmelCase__					)

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

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

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

       if not drop_last:
        while len(UpperCAmelCase__					) < len(UpperCAmelCase__					):
         reference += reference
       self.assertListEqual(UpperCAmelCase__       ,					reference[: len(UpperCAmelCase__					)]					)
      def 		UpperCAmelCase_  (	self      :Any					)   ->	Optional[Any]:
       UpperCAmelCase__							       =  42
       UpperCAmelCase__							       =  RandomIterableDataset()

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

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

       self.check_iterable_dataset_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					batch_size=4       ,					drop_last=UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)
       self.check_iterable_dataset_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					batch_size=4       ,					drop_last=UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)
       self.check_iterable_dataset_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					batch_size=4       ,					drop_last=UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)
       self.check_iterable_dataset_shards(UpperCAmelCase__       ,					UpperCAmelCase__       ,					batch_size=4       ,					drop_last=UpperCAmelCase__       ,					split_batches=UpperCAmelCase__					)
      def 		UpperCAmelCase_  (	self      :Dict					)   ->	Optional[int]:
       UpperCAmelCase__							       =  BatchSampler(range(16					)       ,					batch_size=4       ,					drop_last=UpperCAmelCase__					)
       UpperCAmelCase__							       =  SkipBatchSampler(UpperCAmelCase__       ,					2					)
       self.assertListEqual(list(UpperCAmelCase__					)       ,					[[8, 9, 10, 11], [12, 13, 14, 15]]					)
      def 		UpperCAmelCase_  (	self      :Optional[int]					)   ->	List[str]:
       UpperCAmelCase__							       =  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      :Optional[Any]					)   ->	Optional[Any]:
       UpperCAmelCase__							       =  DataLoader(list(range(16					)					)       ,					batch_size=4					)
       UpperCAmelCase__							       =  skip_first_batches(UpperCAmelCase__       ,					num_batches=2					)
       self.assertListEqual([t.tolist() for t in new_dataloader]       ,					[[8, 9, 10, 11], [12, 13, 14, 15]]					)
      def 		UpperCAmelCase_  (	self      :Optional[int]					)   ->	List[str]:
       UpperCAmelCase__							       =  DataLoaderShard(list(range(16					)					)       ,					batch_size=4					)
       for idx, _ in enumerate(UpperCAmelCase__					):
        self.assertEqual(dataloader.end_of_dataloader       ,					idx == 3					)

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







      def 		UpperCAmelCase_  (	self      :List[str]					)   ->	Dict:
       Accelerator()
       UpperCAmelCase__							       =  DataLoaderDispatcher(range(16					)       ,					batch_size=4					)
       for idx, _ in enumerate(UpperCAmelCase__					):
        self.assertEqual(dataloader.end_of_dataloader       ,					idx == 3					)

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











import io
import itertools
import json
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa
import pyarrow.json as paj

import datasets
from datasets.table import table_cast
from datasets.utils.file_utils import readline


_lowerCamelCase	:		Optional[Any]							       =   datasets.utils.logging.get_logger(__name__)


@dataclass
class 				UpperCamelCase_    ( datasets.BuilderConfig     ):

    '''simple docstring'''






    UpperCAmelCase__							  =      None
    UpperCAmelCase__							  =      "utf-8"
    UpperCAmelCase__							  =      None
    UpperCAmelCase__							  =      None
    UpperCAmelCase__							  =      True  # deprecated
    UpperCAmelCase__							  =      None  # deprecated
    UpperCAmelCase__							  =      10 << 20  # 10MB
    UpperCAmelCase__							  =      None


class 				UpperCamelCase_    ( datasets.ArrowBasedBuilder     ):

    '''simple docstring'''






    UpperCAmelCase__							  =      JsonConfig


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





           '''simple docstring'''

           if self.config.block_size is not None:
                  logger.warning('''The JSON loader parameter `block_size` is deprecated. Please use `chunksize` instead''')
                  A__       =    self.config.block_size
           if self.config.use_threads is not True:
                  logger.warning(
                      '''The JSON loader parameter `use_threads` is deprecated and doesn\'t have any effect anymore.''')
           if self.config.newlines_in_values is not None:
                  raise ValueError('''The JSON loader parameter `newlines_in_values` is no longer supported''')
           return datasets.DatasetInfo(features=self.config.features)


    def        SCREAMING_SNAKE_CASE							(			self      :	List[str]					, UpperCAmelCase__      :	List[Any])						->Dict:





           '''simple docstring'''

           if not self.config.data_files:
                  raise ValueError(f"""At least one data file must be specified, but got data_files={self.config.data_files}""")
           A__       =    dl_manager.download_and_extract(self.config.data_files)
           if isinstance(UpperCAmelCase__					, (str, list, tuple)):
                  A__       =    data_files
                  if isinstance(UpperCAmelCase__					, UpperCAmelCase__):
                         A__       =    [files]
                  A__       =    [dl_manager.iter_files(UpperCAmelCase__) for file in files]
                  return [datasets.SplitGenerator(name=datasets.Split.TRAIN					, gen_kwargs={'''files''': files})]
           A__       =    []
           for split_name, files in data_files.items():
                  if isinstance(UpperCAmelCase__					, UpperCAmelCase__):
                         A__       =    [files]
                  A__       =    [dl_manager.iter_files(UpperCAmelCase__) for file in files]
                  splits.append(datasets.SplitGenerator(name=UpperCAmelCase__					, gen_kwargs={'''files''': files}))
           return splits


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, UpperCAmelCase__      :	pa.Table)						->pa.Table:





           '''simple docstring'''

           if self.config.features is not None:
                  # adding missing columns
                  for column_name in set(self.config.features) - set(pa_table.column_names):
                         A__       =    self.config.features.arrow_schema.field(UpperCAmelCase__).type
                         A__       =    pa_table.append_column(UpperCAmelCase__					, pa.array([None] * len(UpperCAmelCase__)					, type=UpperCAmelCase__))
                  # more expensive cast to support nested structures with keys in a different order
                  # allows str <-> int/float or str to Audio for example
                  A__       =    table_cast(UpperCAmelCase__					, self.config.features.arrow_schema)
           return pa_table





    def        SCREAMING_SNAKE_CASE							(			self      :	Union[str, Any]					, UpperCAmelCase__      :	Tuple)						->str:





           '''simple docstring'''

           for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase__)):
                  # If the file is one json object and if we need to look at the list of items in one specific field
                  if self.config.field is not None:
                         with open(UpperCAmelCase__					, encoding=self.config.encoding					, errors=self.config.encoding_errors) as f:
                                A__       =    json.load(UpperCAmelCase__)

                         # We keep only the field we are interested in
                         A__       =    dataset[self.config.field]

                         # We accept two format: a list of dicts or a dict of lists
                         if isinstance(UpperCAmelCase__					, (list, tuple)):
                                A__       =    set().union(*[row.keys() for row in dataset])
                                A__       =    {col: [row.get(UpperCAmelCase__) for row in dataset] for col in keys}
                         else:
                                A__       =    dataset
                         A__       =    pa.Table.from_pydict(UpperCAmelCase__)
                         yield file_idx, self._cast_table(UpperCAmelCase__)

                  # If the file has one json object per line
                  else:
                         with open(UpperCAmelCase__					, '''rb''') as f:
                                A__       =    0
                                # Use block_size equal to the chunk size divided by 32 to leverage multithreading
                                # Set a default minimum value of 16kB if the chunk size is really small
                                A__       =    max(self.config.chunksize // 32					, 16 << 10)
                                A__       =    (
                                    self.config.encoding_errors if self.config.encoding_errors is not None else '''strict'''
                                )
                                while True:
                                       A__       =    f.read(self.config.chunksize)
                                       if not batch:
                                              break
                                       # Finish current line
                                       try:
                                              batch += f.readline()
                                       except (AttributeError, io.UnsupportedOperation):
                                              batch += readline(UpperCAmelCase__)
                                       # PyArrow only accepts utf-8 encoded bytes
                                       if self.config.encoding != "utf-8":
                                              A__       =    batch.decode(self.config.encoding					, errors=UpperCAmelCase__).encode('''utf-8''')
                                       try:
                                              while True:
                                                     try:
                                                            A__       =    paj.read_json(
                                                                io.BytesIO(UpperCAmelCase__)					, read_options=paj.ReadOptions(block_size=UpperCAmelCase__))
                                                            break
                                                     except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as e:
                                                            if (
                                                                isinstance(UpperCAmelCase__					, pa.ArrowInvalid)
                                                                and "straddling" not in str(UpperCAmelCase__)
                                                                or block_size > len(UpperCAmelCase__)
                                                            ):
                                                                   raise
                                                            else:
                                                                   # Increase the block size in case it was too small.
                                                                   # The block size will be reset for the next file.
                                                                   logger.debug(
                                                                       f"""Batch of {len(UpperCAmelCase__)} bytes couldn't be parsed with block_size={block_size}. Retrying with block_size={block_size * 2}.""")
                                                                   block_size *= 2
                                       except pa.ArrowInvalid as e:
                                              try:
                                                     with open(
                                                         UpperCAmelCase__					, encoding=self.config.encoding					, errors=self.config.encoding_errors) as f:
                                                            A__       =    json.load(UpperCAmelCase__)
                                              except json.JSONDecodeError:
                                                     logger.error(f"""Failed to read file '{file}' with error {type(UpperCAmelCase__)}: {e}""")
                                                     raise e
                                              # If possible, parse the file as a list of json objects and exit the loop
                                              if isinstance(UpperCAmelCase__					, UpperCAmelCase__):  # list is the only sequence type supported in JSON
                                                     try:
                                                            A__       =    set().union(*[row.keys() for row in dataset])
                                                            A__       =    {col: [row.get(UpperCAmelCase__) for row in dataset] for col in keys}
                                                            A__       =    pa.Table.from_pydict(UpperCAmelCase__)
                                                     except (pa.ArrowInvalid, AttributeError) as e:
                                                            logger.error(f"""Failed to read file '{file}' with error {type(UpperCAmelCase__)}: {e}""")
                                                            raise ValueError(f"""Not able to read records in the JSON file at {file}.""") from None
                                                     yield file_idx, self._cast_table(UpperCAmelCase__)
                                                     break
                                              else:
                                                     logger.error(f"""Failed to read file '{file}' with error {type(UpperCAmelCase__)}: {e}""")
                                                     raise ValueError(
                                                         f"""Not able to read records in the JSON file at {file}. """
                                                         f"""You should probably indicate the field of the JSON file containing your records. """
                                                         f"""This JSON file contain the following fields: {str(list(dataset.keys()))}. """
                                                         f"""Select the correct one and provide it as `field='XXX'` to the dataset loading method. """) from None
                        # Uncomment for debugging (will print the Arrow table size and elements)
                        # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                        # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                       yield (file_idx, batch_idx), self._cast_table(UpperCAmelCase__)
                                       batch_idx += 1


