





'''simple docstring'''


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       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Optional[int]	,      lowercase		:					List[str]=0.01	,      lowercase		:					Dict=1_000						):





        """simple docstring"""




        lowercase_			:int    			= p_stop
        lowercase_			:Dict    			= max_length







     def __iter__( self		:					Dict						):





        """simple docstring"""




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







class       a_   (				unittest.TestCase							):





     def 							lowercase__      ( self		:					Tuple	,      lowercase		:					Union[str, Any]	,      lowercase		:					int	,      lowercase		:					Union[str, Any]=False	,      lowercase		:					Union[str, Any]=True						):





        """simple docstring"""




        lowercase_			:Optional[int]    			= [
            BatchSamplerShard(lowercase	,      2	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)
            for i in range(2						)
        ]
        lowercase_			:Optional[Any]    			= [list(lowercase						) for batch_sampler_shard in batch_sampler_shards]
        if not split_batches:
           self.assertListEqual([len(lowercase						) for shard in batch_sampler_shards]	,      [len(lowercase						) for e in expected]						)
        self.assertListEqual(lowercase	,      lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Dict    			= BatchSampler(range(24						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:List[str]    			= [
            [[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(lowercase	,      lowercase						)

        lowercase_			:str    			= BatchSampler(range(24						)	,      batch_size=3	,      drop_last=lowercase						)
        # Expected shouldn't change
        self.check_batch_sampler_shards(lowercase	,      lowercase						)

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

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

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

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

        # 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[str]    			= BatchSampler(range(20						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:Tuple    			= [
            [[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(lowercase	,      lowercase						)

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

        # Check the shards when the dataset is very small.
        lowercase_			:Optional[int]    			= BatchSampler(range(2						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:Optional[Any]    			= [[[0, 1, 0]], [[1, 0, 1]]]
        self.check_batch_sampler_shards(lowercase	,      lowercase						)

        lowercase_			:Any    			= BatchSampler(range(2						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:int    			= [[], []]
        self.check_batch_sampler_shards(lowercase	,      lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= BatchSampler(range(24						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Union[str, 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(lowercase	,      lowercase	,      split_batches=lowercase						)

        lowercase_			:Optional[int]    			= BatchSampler(range(24						)	,      batch_size=4	,      drop_last=lowercase						)
        # Expected shouldn't change
        self.check_batch_sampler_shards(lowercase	,      lowercase	,      split_batches=lowercase						)

        # 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=lowercase						)
        lowercase_			:List[str]    			= [
            [[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(lowercase	,      lowercase	,      split_batches=lowercase						)

        lowercase_			:Tuple    			= BatchSampler(range(22						)	,      batch_size=4	,      drop_last=lowercase						)
        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(lowercase	,      lowercase	,      split_batches=lowercase						)

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

        lowercase_			:List[Any]    			= BatchSampler(range(21						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Optional[Any]    			= [
            [[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(lowercase	,      lowercase	,      split_batches=lowercase						)

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

        lowercase_			:List[str]    			= BatchSampler(range(2						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Dict    			= [[], []]
        self.check_batch_sampler_shards(lowercase	,      lowercase	,      split_batches=lowercase						)





     def 							lowercase__      ( self		:					Dict						):





        """simple docstring"""




        lowercase_			:Tuple    			= BatchSampler(range(24						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:List[str]    			= [
            [[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(lowercase	,      lowercase	,      even_batches=lowercase						)

        lowercase_			:Optional[Any]    			= BatchSampler(range(24						)	,      batch_size=3	,      drop_last=lowercase						)
        # Expected shouldn't change
        self.check_batch_sampler_shards(lowercase	,      lowercase	,      even_batches=lowercase						)

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

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

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

        lowercase_			:List[Any]    			= BatchSampler(range(22						)	,      batch_size=3	,      drop_last=lowercase						)
        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(lowercase	,      lowercase	,      even_batches=lowercase						)

        # 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_			:int    			= BatchSampler(range(20						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:Optional[Any]    			= [
            [[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(lowercase	,      lowercase	,      even_batches=lowercase						)

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

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

        lowercase_			:Optional[Any]    			= BatchSampler(range(2						)	,      batch_size=3	,      drop_last=lowercase						)
        lowercase_			:List[Any]    			= [[], []]
        self.check_batch_sampler_shards(lowercase	,      lowercase	,      even_batches=lowercase						)





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= BatchSampler(range(24						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Tuple    			= [
            [[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(lowercase	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)

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

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

        lowercase_			:str    			= BatchSampler(range(22						)	,      batch_size=4	,      drop_last=lowercase						)
        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(lowercase	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)

        # 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=lowercase						)
        lowercase_			:str    			= [
            [[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(lowercase	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)

        lowercase_			:Dict    			= BatchSampler(range(21						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Union[str, Any]    			= [
            [[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(lowercase	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)

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

        lowercase_			:Tuple    			= BatchSampler(range(2						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:Dict    			= [[], []]
        self.check_batch_sampler_shards(lowercase	,      lowercase	,      split_batches=lowercase	,      even_batches=lowercase						)





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




        lowercase_			:str    			= [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
        lowercase_			:Tuple    			= [BatchSamplerShard(lowercase	,      2	,      lowercase	,      even_batches=lowercase						) 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 							lowercase__      ( self		:					Union[str, Any]	,      lowercase		:					Union[str, Any]	,      lowercase		:					Dict	,      lowercase		:					Tuple	,      lowercase		:					List[Any]=False	,      lowercase		:					int=2	,      lowercase		:					List[str]=False						):





        """simple docstring"""




        random.seed(lowercase						)
        lowercase_			:Optional[int]    			= list(lowercase						)

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

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

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

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





     def 							lowercase__      ( self		:					int						):





        """simple docstring"""




        lowercase_			:str    			= 42
        lowercase_			:Any    			= RandomIterableDataset()

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

        # Edge case with a very small dataset
        lowercase_			:List[Any]    			= RandomIterableDataset(max_length=2						)

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





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




        lowercase_			:Any    			= BatchSampler(range(16						)	,      batch_size=4	,      drop_last=lowercase						)
        lowercase_			:List[Any]    			= SkipBatchSampler(lowercase	,      2						)
        self.assertListEqual(list(lowercase						)	,      [[8, 9, 10, 11], [12, 13, 14, 15]]						)





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        lowercase_			: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 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:Tuple    			= DataLoader(list(range(16						)						)	,      batch_size=4						)
        lowercase_			:Optional[int]    			= skip_first_batches(lowercase	,      num_batches=2						)
        self.assertListEqual([t.tolist() for t in new_dataloader]	,      [[8, 9, 10, 11], [12, 13, 14, 15]]						)





     def 							lowercase__      ( self		:					Optional[Any]						):





        """simple docstring"""




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

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







     def 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        Accelerator()
        lowercase_			:Union[str, Any]    			= DataLoaderDispatcher(range(16						)	,      batch_size=4						)
        for idx, _ in enumerate(lowercase						):
           self.assertEqual(dataloader.end_of_dataloader	,      idx == 3						)

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









'''simple docstring'''


lowerCAmelCase					:     str												='''
# Installazione di Transformers
! pip install transformers datasets
# Per installare dalla fonte invece dell\'ultima versione rilasciata, commenta il comando sopra e
# rimuovi la modalità commento al comando seguente.
# ! pip install git+https://github.com/huggingface/transformers.git
'''

lowerCAmelCase					:     int												=[{'''type''': '''code''', '''content''': INSTALL_CONTENT}]
lowerCAmelCase					:     List[str]												={
    '''{processor_class}''': '''FakeProcessorClass''',
    '''{model_class}''': '''FakeModelClass''',
    '''{object_class}''': '''FakeObjectClass''',
}


