

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_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a=0.01       ,							_a=1_0_0_0       )    ->    List[str]:
												_a	:    Any         =		p_stop
												_a	:    Dict         =		max_length




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



class 				UpperCAmelCase_		(			unittest.TestCase   ):
							"""simple docstring"""







							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a=False       ,							_a=True       )    ->    int:
												_a	:    List[Any]         =		[
												    BatchSamplerShard(_a       ,							2       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )
												    for i in range(2       )
												]
												_a	:    Optional[int]         =		[list(_a       ) for batch_sampler_shard in batch_sampler_shards]
												if not split_batches:
																	self.assertListEqual([len(_a       ) for shard in batch_sampler_shards]       ,							[len(_a       ) for e in expected]       )
												self.assertListEqual(_a       ,							_a       )


							def 				__lowercase				(			self       )    ->    List[Any]:
												# Check the shards when the dataset is a round multiple of total batch size.
												_a	:    List[str]         =		BatchSampler(range(2_4       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												_a	:    str         =		BatchSampler(range(2_4       )       ,							batch_size=3       ,							drop_last=_a       )
												# Expected shouldn't change
												self.check_batch_sampler_shards(_a       ,							_a       )

												# Check the shards when the dataset is a round multiple of batch size but not total batch size.
												_a	:    Any         =		BatchSampler(range(2_1       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    str         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [0, 1, 2]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												_a	:    Any         =		BatchSampler(range(2_1       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Union[str, Any]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
												# num_processes batch.
												_a	:    Dict         =		BatchSampler(range(2_2       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 0, 1]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												_a	:    Tuple         =		BatchSampler(range(2_2       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    str         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
												# num_processes batch.
												_a	:    Optional[Any]         =		BatchSampler(range(2_0       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [1, 2, 3]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

												_a	:    Union[str, Any]         =		BatchSampler(range(2_0       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Optional[Any]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       )

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

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


							def 				__lowercase				(			self       )    ->    Optional[Any]:
												# Check the shards when the dataset is a round multiple of batch size.
												_a	:    Any         =		BatchSampler(range(2_4       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Optional[Any]         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

												_a	:    Tuple         =		BatchSampler(range(2_4       )       ,							batch_size=4       ,							drop_last=_a       )
												# Expected shouldn't change
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size.
												_a	:    Tuple         =		BatchSampler(range(2_2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    int         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [0, 1]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

												_a	:    Any         =		BatchSampler(range(2_2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size or num_processes.
												_a	:    Tuple         =		BatchSampler(range(2_1       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Dict         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 0]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [1, 2]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

												_a	:    int         =		BatchSampler(range(2_1       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )

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

												_a	:    Tuple         =		BatchSampler(range(2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    List[Any]         =		[[], []]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       )


							def 				__lowercase				(			self       )    ->    Optional[Any]:
												# Check the shards when the dataset is a round multiple of total batch size.
												_a	:    Dict         =		BatchSampler(range(2_4       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Optional[Any]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												_a	:    Any         =		BatchSampler(range(2_4       )       ,							batch_size=3       ,							drop_last=_a       )
												# Expected shouldn't change
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												# Check the shards when the dataset is a round multiple of batch size but not total batch size.
												_a	:    Dict         =		BatchSampler(range(2_1       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Tuple         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												_a	:    str         =		BatchSampler(range(2_1       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Optional[int]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
												# num_processes batch.
												_a	:    Dict         =		BatchSampler(range(2_2       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    List[str]         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												_a	:    str         =		BatchSampler(range(2_2       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Dict         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
												# num_processes batch.
												_a	:    Optional[int]         =		BatchSampler(range(2_0       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Tuple         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

												_a	:    List[Any]         =		BatchSampler(range(2_0       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    str         =		[
												    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
												    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )

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

												_a	:    Optional[Any]         =		BatchSampler(range(2       )       ,							batch_size=3       ,							drop_last=_a       )
												_a	:    Tuple         =		[[], []]
												self.check_batch_sampler_shards(_a       ,							_a       ,							even_batches=_a       )


							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												# Check the shards when the dataset is a round multiple of batch size.
												_a	:    Tuple         =		BatchSampler(range(2_4       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    str         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

												_a	:    Union[str, Any]         =		BatchSampler(range(2_4       )       ,							batch_size=4       ,							drop_last=_a       )
												# Expected shouldn't change
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size.
												_a	:    Optional[Any]         =		BatchSampler(range(2_2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Optional[int]         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

												_a	:    int         =		BatchSampler(range(2_2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Optional[Any]         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

												# Check the shards when the dataset is not a round multiple of batch size or num_processes.
												_a	:    Dict         =		BatchSampler(range(2_1       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Dict         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

												_a	:    Optional[int]         =		BatchSampler(range(2_1       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Any         =		[
												    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
												    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
												]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )

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

												_a	:    Dict         =		BatchSampler(range(2       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Optional[int]         =		[[], []]
												self.check_batch_sampler_shards(_a       ,							_a       ,							split_batches=_a       ,							even_batches=_a       )


							def 				__lowercase				(			self       )    ->    Optional[Any]:
												_a	:    int         =		[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
												_a	:    str         =		[BatchSamplerShard(_a       ,							2       ,							_a       ,							even_batches=_a       ) 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], [1_2, 1_3]]       )
												self.assertListEqual(list(batch_sampler_shards[1]       )       ,							[[3, 4], [9, 1_0, 1_1]]       )


							def 				__lowercase				(			self       ,							_a       ,							_a       ,							_a       ,							_a=False       ,							_a=2       ,							_a=False       )    ->    List[Any]:
												random.seed(_a       )
												_a	:    str         =		list(_a       )

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

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

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

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


							def 				__lowercase				(			self       )    ->    Dict:
												_a	:    List[Any]         =		4_2
												_a	:    Union[str, Any]         =		RandomIterableDataset()

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

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

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


							def 				__lowercase				(			self       )    ->    int:
												_a	:    List[Any]         =		BatchSampler(range(1_6       )       ,							batch_size=4       ,							drop_last=_a       )
												_a	:    Optional[int]         =		SkipBatchSampler(_a       ,							2       )
												self.assertListEqual(list(_a       )       ,							[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]       )


							def 				__lowercase				(			self       )    ->    Dict:
												_a	:    Dict         =		SkipDataLoader(list(range(1_6       )       )       ,							batch_size=4       ,							skip_batches=2       )
												self.assertListEqual([t.tolist() for t in dataloader]       ,							[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]       )


							def 				__lowercase				(			self       )    ->    Optional[int]:
												_a	:    Tuple         =		DataLoader(list(range(1_6       )       )       ,							batch_size=4       )
												_a	:    List[str]         =		skip_first_batches(_a       ,							num_batches=2       )
												self.assertListEqual([t.tolist() for t in new_dataloader]       ,							[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]       )


							def 				__lowercase				(			self       )    ->    List[Any]:
												_a	:    int         =		DataLoaderShard(list(range(1_6       )       )       ,							batch_size=4       )
												for idx, _ in enumerate(_a       ):
																	self.assertEqual(dataloader.end_of_dataloader       ,							idx == 3       )

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




							def 				__lowercase				(			self       )    ->    Optional[Any]:
												Accelerator()
												_a	:    Optional[int]         =		DataLoaderDispatcher(range(1_6       )       ,							batch_size=4       )
												for idx, _ in enumerate(_a       ):
																	self.assertEqual(dataloader.end_of_dataloader       ,							idx == 3       )

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






from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


a__				=		{'''configuration_mmbt''': ['''MMBTConfig''']}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			a__				=		['''MMBTForClassification''', '''MMBTModel''', '''ModalEmbeddings''']


if TYPE_CHECKING:
			from .configuration_mmbt import MMBTConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings

else:
			import sys

			a__				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



