


'''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     lowercase__  (						_snake_case					):





							'''simple docstring'''


							def __init__(    self    ,					__snake_case=0.01    ,					__snake_case=1000  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      p_stop
											_SCREAMING_SNAKE_CASE :       List[str]     		=      max_length






							def __iter__(    self  ):
											_SCREAMING_SNAKE_CASE :       int     		=      0
											_SCREAMING_SNAKE_CASE :       str     		=      False
											while not stop and count < self.max_length:
															yield count
															count += 1
															_SCREAMING_SNAKE_CASE :       Tuple     		=      random.random() < self.p_stop





class     lowercase__  (						unittest.TestCase					):





							'''simple docstring'''


							def UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case=False    ,					__snake_case=True  ):
											_SCREAMING_SNAKE_CASE :       int     		=      [
											    BatchSamplerShard(__snake_case    ,					2    ,					__snake_case    ,					split_batches=__snake_case    ,					even_batches=__snake_case  )
											    for i in range(2  )
											]
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [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 UpperCAmelCase_	(    self  ):
											# Check the shards when the dataset is a round multiple of total batch size.
											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(24  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      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.
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      BatchSampler(range(21  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      BatchSampler(range(21  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      [
											    [[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.
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      BatchSampler(range(22  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__snake_case    ,					__snake_case  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(22  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      [
											    [[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.
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      BatchSampler(range(20  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      BatchSampler(range(20  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      [
											    [[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.
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      BatchSampler(range(2  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [[[0, 1, 0]], [[1, 0, 1]]]
											self.check_batch_sampler_shards(__snake_case    ,					__snake_case  )

											_SCREAMING_SNAKE_CASE :       Dict     		=      BatchSampler(range(2  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      [[], []]
											self.check_batch_sampler_shards(__snake_case    ,					__snake_case  )






							def UpperCAmelCase_	(    self  ):
											# Check the shards when the dataset is a round multiple of batch size.
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      BatchSampler(range(24  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__snake_case    ,					__snake_case    ,					split_batches=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      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.
											_SCREAMING_SNAKE_CASE :       int     		=      BatchSampler(range(22  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__snake_case    ,					__snake_case    ,					split_batches=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      BatchSampler(range(22  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__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.
											_SCREAMING_SNAKE_CASE :       str     		=      BatchSampler(range(21  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       str     		=      [
											    [[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  )

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

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

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






							def UpperCAmelCase_	(    self  ):
											# Check the shards when the dataset is a round multiple of total batch size.
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      BatchSampler(range(24  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       Dict     		=      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.
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      BatchSampler(range(21  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      BatchSampler(range(21  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Any     		=      [
											    [[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.
											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(22  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Any     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      BatchSampler(range(22  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__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.
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      BatchSampler(range(20  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[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(__snake_case    ,					__snake_case    ,					even_batches=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      BatchSampler(range(20  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       int     		=      [
											    [[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.
											_SCREAMING_SNAKE_CASE :       int     		=      BatchSampler(range(2  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      [[[0, 1]], []]
											self.check_batch_sampler_shards(__snake_case    ,					__snake_case    ,					even_batches=__snake_case  )

											_SCREAMING_SNAKE_CASE :       str     		=      BatchSampler(range(2  )    ,					batch_size=3    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      [[], []]
											self.check_batch_sampler_shards(__snake_case    ,					__snake_case    ,					even_batches=__snake_case  )






							def UpperCAmelCase_	(    self  ):
											# Check the shards when the dataset is a round multiple of batch size.
											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(24  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       int     		=      [
											    [[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  )

											_SCREAMING_SNAKE_CASE :       str     		=      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.
											_SCREAMING_SNAKE_CASE :       str     		=      BatchSampler(range(22  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[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(__snake_case    ,					__snake_case    ,					split_batches=__snake_case    ,					even_batches=__snake_case  )

											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(22  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       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(__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.
											_SCREAMING_SNAKE_CASE :       str     		=      BatchSampler(range(21  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [
											    [[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  )

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

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

											_SCREAMING_SNAKE_CASE :       List[str]     		=      BatchSampler(range(2  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      [[], []]
											self.check_batch_sampler_shards(__snake_case    ,					__snake_case    ,					split_batches=__snake_case    ,					even_batches=__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      [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 UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case=False    ,					__snake_case=2    ,					__snake_case=False  ):
											random.seed(__snake_case  )
											_SCREAMING_SNAKE_CASE :       int     		=      list(__snake_case  )

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      [
											    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  )
											]
											_SCREAMING_SNAKE_CASE :       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(__snake_case  )
															iterable_dataset_lists.append(list(__snake_case  )  )

											_SCREAMING_SNAKE_CASE :       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
											_SCREAMING_SNAKE_CASE :       List[str]     		=      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  )

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      []
											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 UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      42
											_SCREAMING_SNAKE_CASE :       Tuple     		=      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
											_SCREAMING_SNAKE_CASE :       Tuple     		=      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 UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      BatchSampler(range(16  )    ,					batch_size=4    ,					drop_last=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      SkipBatchSampler(__snake_case    ,					2  )
											self.assertListEqual(list(__snake_case  )    ,					[[8, 9, 10, 11], [12, 13, 14, 15]]  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[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  ):
											_SCREAMING_SNAKE_CASE :       str     		=      DataLoader(list(range(16  )  )    ,					batch_size=4  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      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 UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      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 UpperCAmelCase_	(    self  ):
											Accelerator()
											_SCREAMING_SNAKE_CASE :       str     		=      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  )








'''simple docstring'''


from sklearn.metrics import fa_score

import datasets


UpperCAmelCase_		:   Optional[int]							   =		'\nThe F1 score is the harmonic mean of the precision and recall. It can be computed with the equation:\nF1 = 2 * (precision * recall) / (precision + recall)\n'


UpperCAmelCase_		:   Tuple							   =		'\nArgs:\n    predictions (`list` of `int`): Predicted labels.\n    references (`list` of `int`): Ground truth labels.\n    labels (`list` of `int`): The set of labels to include when `average` is not set to `\'binary\'`, and the order of the labels if `average` is `None`. Labels present in the data can be excluded, for example to calculate a multiclass average ignoring a majority negative class. Labels not present in the data will result in 0 components in a macro average. For multilabel targets, labels are column indices. By default, all labels in `predictions` and `references` are used in sorted order. Defaults to None.\n    pos_label (`int`): The class to be considered the positive class, in the case where `average` is set to `binary`. Defaults to 1.\n    average (`string`): This parameter is required for multiclass/multilabel targets. If set to `None`, the scores for each class are returned. Otherwise, this determines the type of averaging performed on the data. Defaults to `\'binary\'`.\n\n        - \'binary\': Only report results for the class specified by `pos_label`. This is applicable only if the classes found in `predictions` and `references` are binary.\n        - \'micro\': Calculate metrics globally by counting the total true positives, false negatives and false positives.\n        - \'macro\': Calculate metrics for each label, and find their unweighted mean. This does not take label imbalance into account.\n        - \'weighted\': Calculate metrics for each label, and find their average weighted by support (the number of true instances for each label). This alters `\'macro\'` to account for label imbalance. This option can result in an F-score that is not between precision and recall.\n        - \'samples\': Calculate metrics for each instance, and find their average (only meaningful for multilabel classification).\n    sample_weight (`list` of `float`): Sample weights Defaults to None.\n\nReturns:\n    f1 (`float` or `array` of `float`): F1 score or list of f1 scores, depending on the value passed to `average`. Minimum possible value is 0. Maximum possible value is 1. Higher f1 scores are better.\n\nExamples:\n\n    Example 1-A simple binary example\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0])\n        >>> print(results)\n        {\'f1\': 0.5}\n\n    Example 2-The same simple binary example as in Example 1, but with `pos_label` set to `0`.\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], pos_label=0)\n        >>> print(round(results[\'f1\'], 2))\n        0.67\n\n    Example 3-The same simple binary example as in Example 1, but with `sample_weight` included.\n        >>> f1_metric = datasets.load_metric("f1")\n        >>> results = f1_metric.compute(references=[0, 1, 0, 1, 0], predictions=[0, 0, 1, 1, 0], sample_weight=[0.9, 0.5, 3.9, 1.2, 0.3])\n        >>> print(round(results[\'f1\'], 2))\n        0.35\n\n    Example 4-A multiclass example, with different values for the `average` input.\n        >>> predictions = [0, 2, 1, 0, 0, 1]\n        >>> references = [0, 1, 2, 0, 1, 2]\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="macro")\n        >>> print(round(results[\'f1\'], 2))\n        0.27\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="micro")\n        >>> print(round(results[\'f1\'], 2))\n        0.33\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average="weighted")\n        >>> print(round(results[\'f1\'], 2))\n        0.27\n        >>> results = f1_metric.compute(predictions=predictions, references=references, average=None)\n        >>> print(results)\n        {\'f1\': array([0.8, 0. , 0. ])}\n'


UpperCAmelCase_		:   Dict							   =		'\n@article{scikit-learn,\n    title={Scikit-learn: Machine Learning in {P}ython},\n    author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n           and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n           and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n           Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n    journal={Journal of Machine Learning Research},\n    volume={12},\n    pages={2825--2830},\n    year={2011}\n}\n'





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,     _KWARGS_DESCRIPTION					)
class     lowercase__  (						datasets.Metric					):





							'''simple docstring'''


							def UpperCAmelCase_	(    self  ):
											return datasets.MetricInfo(
											    description=_DESCRIPTION    ,					citation=_CITATION    ,					inputs_description=_KWARGS_DESCRIPTION    ,					features=datasets.Features(
											        {
											            """predictions""": datasets.Sequence(datasets.Value("""int32"""  )  ),
											            """references""": datasets.Sequence(datasets.Value("""int32"""  )  ),
											        }
											        if self.config_name == """multilabel"""
											        else {
											            """predictions""": datasets.Value("""int32"""  ),
											            """references""": datasets.Value("""int32"""  ),
											        }  )    ,					reference_urls=["""https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html"""]    ,					)






							def UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case=None    ,					__snake_case=1    ,					__snake_case="binary"    ,					__snake_case=None  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      fa_score(
											    __snake_case    ,					__snake_case    ,					labels=__snake_case    ,					pos_label=__snake_case    ,					average=__snake_case    ,					sample_weight=__snake_case  )
											return {"f1": float(__snake_case  ) if score.size == 1 else score}




