
'''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 			SCREAMING_SNAKE_CASE__      (     snake_case_):







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







											'''simple docstring'''






											UpperCamelCase   			=     p_stop
											UpperCamelCase   			=     max_length







				def __iter__(       self     )->		Any:







											'''simple docstring'''






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


class 			SCREAMING_SNAKE_CASE__      (     unittest.TestCase):







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_=False						,				A_=True     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     [
											    BatchSamplerShard(A_						,				2						,				A_						,				split_batches=A_						,				even_batches=A_     )
											    for i in range(2     )
											]
											UpperCamelCase   			=     [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 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BatchSampler(range(24     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_     )

											UpperCamelCase   			=     BatchSampler(range(24     )						,				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.
											UpperCamelCase   			=     BatchSampler(range(21     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_     )

											UpperCamelCase   			=     BatchSampler(range(21     )						,				batch_size=3						,				drop_last=A_     )
											UpperCamelCase   			=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											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.
											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_     )

											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=3						,				drop_last=A_     )
											UpperCamelCase   			=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											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.
											UpperCamelCase   			=     BatchSampler(range(20     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_     )

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

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

											UpperCamelCase   			=     BatchSampler(range(2     )						,				batch_size=3						,				drop_last=A_     )
											UpperCamelCase   			=     [[], []]
											self.check_batch_sampler_shards(A_						,				A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     BatchSampler(range(24     )						,				batch_size=4						,				drop_last=A_     )
											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(A_						,				A_						,				split_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(24     )						,				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.
											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=4						,				drop_last=A_     )
											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(A_						,				A_						,				split_batches=A_     )

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

											# 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=A_     )
											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(A_						,				A_						,				split_batches=A_     )

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

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

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







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







											'''simple docstring'''






											UpperCamelCase   			=     BatchSampler(range(24     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_						,				even_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(24     )						,				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.
											UpperCamelCase   			=     BatchSampler(range(21     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_						,				even_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(21     )						,				batch_size=3						,				drop_last=A_     )
											UpperCamelCase   			=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											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.
											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_						,				even_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=3						,				drop_last=A_     )
											UpperCamelCase   			=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											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.
											UpperCamelCase   			=     BatchSampler(range(20     )						,				batch_size=3						,				drop_last=A_     )
											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(A_						,				A_						,				even_batches=A_     )

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

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

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







				def 						UpperCAmelCase_    (       self     )->		str:







											'''simple docstring'''






											UpperCamelCase   			=     BatchSampler(range(24     )						,				batch_size=4						,				drop_last=A_     )
											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(A_						,				A_						,				split_batches=A_						,				even_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(24     )						,				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.
											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=4						,				drop_last=A_     )
											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(A_						,				A_						,				split_batches=A_						,				even_batches=A_     )

											UpperCamelCase   			=     BatchSampler(range(22     )						,				batch_size=4						,				drop_last=A_     )
											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(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.
											UpperCamelCase   			=     BatchSampler(range(21     )						,				batch_size=4						,				drop_last=A_     )
											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(A_						,				A_						,				split_batches=A_						,				even_batches=A_     )

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

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

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







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											UpperCamelCase   			=     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
											UpperCamelCase   			=     [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], [12, 13]]     )
											self.assertListEqual(list(batch_sampler_shards[1]     )						,				[[3, 4], [9, 10, 11]]     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_=False						,				A_=2						,				A_=False     )->		Optional[int]:







											'''simple docstring'''






											random.seed(A_     )
											UpperCamelCase   			=     list(A_     )

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

											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(A_     )						,				len(A_     )     )
																		self.assertTrue(len(A_     ) % shard_batch_size == 0     )

											UpperCamelCase   			=     []
											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 						UpperCAmelCase_    (       self     )->		List[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     42
											UpperCamelCase   			=     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
											UpperCamelCase   			=     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 						UpperCAmelCase_    (       self     )->		Optional[int]:







											'''simple docstring'''






											UpperCamelCase   			=     BatchSampler(range(16     )						,				batch_size=4						,				drop_last=A_     )
											UpperCamelCase   			=     SkipBatchSampler(A_						,				2     )
											self.assertListEqual(list(A_     )						,				[[8, 9, 10, 11], [12, 13, 14, 15]]     )







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







											'''simple docstring'''






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







											'''simple docstring'''






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







				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     DataLoaderShard(list(range(16     )     )						,				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 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											Accelerator()
											UpperCamelCase   			=     DataLoaderDispatcher(range(16     )						,				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     )






'''simple docstring'''
from typing import TYPE_CHECKING

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


lowerCAmelCase					:    Optional[Any]           = {
    'configuration_falcon': ['FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP', 'FalconConfig'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCAmelCase					:    str           = [
						    'FALCON_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'FalconForCausalLM',
						    'FalconModel',
						    'FalconPreTrainedModel',
						    'FalconForSequenceClassification',
						    'FalconForTokenClassification',
						    'FalconForQuestionAnswering',
						]


if TYPE_CHECKING:
						from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_falcon import (
												    FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
												    FalconForCausalLM,
												    FalconForQuestionAnswering,
												    FalconForSequenceClassification,
												    FalconForTokenClassification,
												    FalconModel,
												    FalconPreTrainedModel,
												)


else:
						import sys

						lowerCAmelCase					:    Tuple           = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




