






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		(	a_	):
							def __init__(     self:      int       ,	UpperCamelCase:      List[str]=0.01       ,	UpperCamelCase:      Dict=10_00						)		->		List[Any]:
											snake_case__									=     p_stop
											snake_case__									=     max_length


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


class   __SCREAMING_SNAKE_CASE		(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      List[str]       ,	UpperCamelCase:      List[str]       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      Union[str, Any]=False       ,	UpperCamelCase:      Any=True						)		->		List[Any]:
											snake_case__									=     [
											    BatchSamplerShard(UpperCamelCase       ,	2       ,	UpperCamelCase       ,	split_batches=UpperCamelCase       ,	even_batches=UpperCamelCase						)
											    for i in range(2						)
											]
											snake_case__									=     [list(UpperCamelCase						) for batch_sampler_shard in batch_sampler_shards]
											if not split_batches:
															self.assertListEqual([len(UpperCamelCase						) for shard in batch_sampler_shards]       ,	[len(UpperCamelCase						) for e in expected]						)
											self.assertListEqual(UpperCamelCase       ,	UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Optional[int]						)		->		List[Any]:
											# Check the shards when the dataset is a round multiple of total batch size.
											snake_case__									=     BatchSampler(range(24						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase						)

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

											# Check the shards when the dataset is a round multiple of batch size but not total batch size.
											snake_case__									=     BatchSampler(range(21						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
											# num_processes batch.
											snake_case__									=     BatchSampler(range(22						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
											# num_processes batch.
											snake_case__									=     BatchSampler(range(20						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase						)

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

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

											snake_case__									=     BatchSampler(range(2						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [[], []]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Tuple						)		->		Optional[int]:
											# Check the shards when the dataset is a round multiple of batch size.
											snake_case__									=     BatchSampler(range(24						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size.
											snake_case__									=     BatchSampler(range(22						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size or num_processes.
											snake_case__									=     BatchSampler(range(21						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase						)

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

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

											snake_case__									=     BatchSampler(range(2						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [[], []]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Optional[Any]						)		->		Any:
											# Check the shards when the dataset is a round multiple of total batch size.
											snake_case__									=     BatchSampler(range(24						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is a round multiple of batch size but not total batch size.
											snake_case__									=     BatchSampler(range(21						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
											# num_processes batch.
											snake_case__									=     BatchSampler(range(22						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
											# num_processes batch.
											snake_case__									=     BatchSampler(range(20						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
											    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

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

											snake_case__									=     BatchSampler(range(2						)       ,	batch_size=3       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [[], []]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Tuple						)		->		Any:
											# Check the shards when the dataset is a round multiple of batch size.
											snake_case__									=     BatchSampler(range(24						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size.
											snake_case__									=     BatchSampler(range(22						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

											# Check the shards when the dataset is not a round multiple of batch size or num_processes.
											snake_case__									=     BatchSampler(range(21						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [
											    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
											    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
											]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase       ,	even_batches=UpperCamelCase						)

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

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

											snake_case__									=     BatchSampler(range(2						)       ,	batch_size=4       ,	drop_last=UpperCamelCase						)
											snake_case__									=     [[], []]
											self.check_batch_sampler_shards(UpperCamelCase       ,	UpperCamelCase       ,	split_batches=UpperCamelCase       ,	even_batches=UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		str:
											snake_case__									=     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
											snake_case__									=     [BatchSamplerShard(UpperCamelCase       ,	2       ,	UpperCamelCase       ,	even_batches=UpperCamelCase						) for i in range(2						)]

											self.assertEqual(len(batch_sampler_shards[0]						)       ,	3						)
											self.assertEqual(len(batch_sampler_shards[1]						)       ,	2						)

											self.assertListEqual(list(batch_sampler_shards[0]						)       ,	[[0, 1, 2], [5, 6, 7, 8], [12, 13]]						)
											self.assertListEqual(list(batch_sampler_shards[1]						)       ,	[[3, 4], [9, 10, 11]]						)
							def lowerCAmelCase_ (     self:      Tuple       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      str=False       ,	UpperCamelCase:      Optional[Any]=2       ,	UpperCamelCase:      List[Any]=False						)		->		Tuple:
											random.seed(UpperCamelCase						)
											snake_case__									=     list(UpperCamelCase						)

											snake_case__									=     [
											    IterableDatasetShard(
											        UpperCamelCase       ,	batch_size=UpperCamelCase       ,	drop_last=UpperCamelCase       ,	num_processes=UpperCamelCase       ,	process_index=UpperCamelCase       ,	split_batches=UpperCamelCase       ,	)
											    for i in range(UpperCamelCase						)
											]
											snake_case__									=     []
											for iterable_dataset_shard in iterable_dataset_shards:
															# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
															random.seed(UpperCamelCase						)
															iterable_dataset_lists.append(list(UpperCamelCase						)						)

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

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

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

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

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

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

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


							def lowerCAmelCase_ (     self:      Optional[int]						)		->		Tuple:
											Accelerator()
											snake_case__									=     DataLoaderDispatcher(range(16						)       ,	batch_size=4						)
											for idx, _ in enumerate(UpperCamelCase						):
															self.assertEqual(dataloader.end_of_dataloader       ,	idx == 3						)

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












import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    LiltForQuestionAnswering,
							    LiltForSequenceClassification,
							    LiltForTokenClassification,
							    LiltModel,
							)
							from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST


class   __SCREAMING_SNAKE_CASE:
							def __init__(     self:      Any       ,	UpperCamelCase:      Tuple       ,	UpperCamelCase:      str=13       ,	UpperCamelCase:      int=7       ,	UpperCamelCase:      Optional[int]=True       ,	UpperCamelCase:      Tuple=True       ,	UpperCamelCase:      List[Any]=True       ,	UpperCamelCase:      Dict=True       ,	UpperCamelCase:      Any=99       ,	UpperCamelCase:      int=24       ,	UpperCamelCase:      List[Any]=2       ,	UpperCamelCase:      Any=6       ,	UpperCamelCase:      Union[str, Any]=37       ,	UpperCamelCase:      int="gelu"       ,	UpperCamelCase:      Optional[Any]=0.1       ,	UpperCamelCase:      Optional[Any]=0.1       ,	UpperCamelCase:      Union[str, Any]=5_12       ,	UpperCamelCase:      List[Any]=16       ,	UpperCamelCase:      str=2       ,	UpperCamelCase:      Any=0.02       ,	UpperCamelCase:      int=3       ,	UpperCamelCase:      str=None       ,	UpperCamelCase:      str=10_00       ,	)		->		Any:
											snake_case__									=     parent
											snake_case__									=     batch_size
											snake_case__									=     seq_length
											snake_case__									=     is_training
											snake_case__									=     use_input_mask
											snake_case__									=     use_token_type_ids
											snake_case__									=     use_labels
											snake_case__									=     vocab_size
											snake_case__									=     hidden_size
											snake_case__									=     num_hidden_layers
											snake_case__									=     num_attention_heads
											snake_case__									=     intermediate_size
											snake_case__									=     hidden_act
											snake_case__									=     hidden_dropout_prob
											snake_case__									=     attention_probs_dropout_prob
											snake_case__									=     max_position_embeddings
											snake_case__									=     type_vocab_size
											snake_case__									=     type_sequence_label_size
											snake_case__									=     initializer_range
											snake_case__									=     num_labels
											snake_case__									=     scope
											snake_case__									=     range_bbox
							def lowerCAmelCase_ (     self:      Optional[Any]						)		->		Dict:
											snake_case__									=     ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size						)

											snake_case__									=     ids_tensor([self.batch_size, self.seq_length, 4]       ,	self.range_bbox						)
											# Ensure that bbox is legal
											for i in range(bbox.shape[0]						):
															for j in range(bbox.shape[1]						):
																			if bbox[i, j, 3] < bbox[i, j, 1]:
																							snake_case__									=     bbox[i, j, 3]
																							snake_case__									=     bbox[i, j, 1]
																							snake_case__									=     t
																			if bbox[i, j, 2] < bbox[i, j, 0]:
																							snake_case__									=     bbox[i, j, 2]
																							snake_case__									=     bbox[i, j, 0]
																							snake_case__									=     t

											snake_case__									=     None
											if self.use_input_mask:
															snake_case__									=     ids_tensor([self.batch_size, self.seq_length]       ,	vocab_size=2						)

											snake_case__									=     None
											if self.use_token_type_ids:
															snake_case__									=     ids_tensor([self.batch_size, self.seq_length]       ,	self.type_vocab_size						)

											snake_case__									=     None
											snake_case__									=     None
											if self.use_labels:
															snake_case__									=     ids_tensor([self.batch_size]       ,	self.type_sequence_label_size						)
															snake_case__									=     ids_tensor([self.batch_size, self.seq_length]       ,	self.num_labels						)

											snake_case__									=     self.get_config()

											return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels
							def lowerCAmelCase_ (     self:      int						)		->		Optional[Any]:
											return LiltConfig(
											    vocab_size=self.vocab_size       ,	hidden_size=self.hidden_size       ,	num_hidden_layers=self.num_hidden_layers       ,	num_attention_heads=self.num_attention_heads       ,	intermediate_size=self.intermediate_size       ,	hidden_act=self.hidden_act       ,	hidden_dropout_prob=self.hidden_dropout_prob       ,	attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,	max_position_embeddings=self.max_position_embeddings       ,	type_vocab_size=self.type_vocab_size       ,	initializer_range=self.initializer_range       ,	)
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      str       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      int       ,	UpperCamelCase:      Union[str, Any]       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      Tuple       ,	)		->		Dict:
											snake_case__									=     LiltModel(config=UpperCamelCase						)
											model.to(UpperCamelCase						)
											model.eval()
											snake_case__									=     model(UpperCamelCase       ,	bbox=UpperCamelCase       ,	attention_mask=UpperCamelCase       ,	token_type_ids=UpperCamelCase						)
											snake_case__									=     model(UpperCamelCase       ,	bbox=UpperCamelCase       ,	token_type_ids=UpperCamelCase						)
											snake_case__									=     model(UpperCamelCase       ,	bbox=UpperCamelCase						)

											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)						)
											self.parent.assertEqual(result.pooler_output.shape       ,	(self.batch_size, self.hidden_size)						)
							def lowerCAmelCase_ (     self:      Optional[Any]       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      Optional[Any]       ,	UpperCamelCase:      List[str]       ,	UpperCamelCase:      str       ,	UpperCamelCase:      Union[str, Any]       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      Optional[Any]       ,	)		->		Tuple:
											snake_case__									=     self.num_labels
											snake_case__									=     LiltForTokenClassification(config=UpperCamelCase						)
											model.to(UpperCamelCase						)
											model.eval()
											snake_case__									=     model(
											    UpperCamelCase       ,	bbox=UpperCamelCase       ,	attention_mask=UpperCamelCase       ,	token_type_ids=UpperCamelCase       ,	labels=UpperCamelCase						)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.num_labels)						)
							def lowerCAmelCase_ (     self:      Optional[int]       ,	UpperCamelCase:      List[str]       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      Tuple       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      List[Any]       ,	UpperCamelCase:      Union[str, Any]       ,	UpperCamelCase:      List[Any]       ,	)		->		Any:
											snake_case__									=     LiltForQuestionAnswering(config=UpperCamelCase						)
											model.to(UpperCamelCase						)
											model.eval()
											snake_case__									=     model(
											    UpperCamelCase       ,	bbox=UpperCamelCase       ,	attention_mask=UpperCamelCase       ,	token_type_ids=UpperCamelCase       ,	start_positions=UpperCamelCase       ,	end_positions=UpperCamelCase       ,	)
											self.parent.assertEqual(result.start_logits.shape       ,	(self.batch_size, self.seq_length)						)
											self.parent.assertEqual(result.end_logits.shape       ,	(self.batch_size, self.seq_length)						)


							def lowerCAmelCase_ (     self:      List[str]						)		->		Optional[int]:
											snake_case__									=     self.prepare_config_and_inputs()
											(
											    (
											    snake_case__
											)							,						(
											    snake_case__
											)							,						(
											    snake_case__
											)							,						(
											    snake_case__
											)							,						(
											    snake_case__
											)							,						(
											    snake_case__
											)							,						(
											    snake_case__
											)							,						
											)									=     config_and_inputs
											snake_case__									=     {
											    'input_ids': input_ids,
											    'bbox': bbox,
											    'token_type_ids': token_type_ids,
											    'attention_mask': input_mask,
											}
											return config, inputs_dict


@require_torch
class   __SCREAMING_SNAKE_CASE(	a_    ,  a_    ,  a_    ,  unittest.TestCase	):
							_UpperCAmelCase     =       (
							    (
							        LiltModel,
							        LiltForSequenceClassification,
							        LiltForTokenClassification,
							        LiltForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							_UpperCAmelCase     =       (
							    {
							        "feature-extraction": LiltModel,
							        "question-answering": LiltForQuestionAnswering,
							        "text-classification": LiltForSequenceClassification,
							        "token-classification": LiltForTokenClassification,
							        "zero-shot": LiltForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							_UpperCAmelCase     =       False
							_UpperCAmelCase     =       False
							def lowerCAmelCase_ (     self:      Dict       ,	UpperCamelCase:      Dict       ,	UpperCamelCase:      List[str]       ,	UpperCamelCase:      Any       ,	UpperCamelCase:      Union[str, Any]       ,	UpperCamelCase:      int						)		->		int:
											return True
							def lowerCAmelCase_ (     self:      Dict						)		->		List[Any]:
											snake_case__									=     LiltModelTester(self						)
											snake_case__									=     ConfigTester(self       ,	config_class=UpperCamelCase       ,	hidden_size=37						)
							def lowerCAmelCase_ (     self:      Optional[Any]						)		->		Tuple:
											self.config_tester.run_common_tests()
							def lowerCAmelCase_ (     self:      Union[str, Any]						)		->		str:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Tuple						)		->		Dict:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											for type in ["absolute", "relative_key", "relative_key_query"]:
															snake_case__									=     type
															self.model_tester.create_and_check_model(*UpperCamelCase						)
							def lowerCAmelCase_ (     self:      int						)		->		Dict:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*UpperCamelCase						)
							def lowerCAmelCase_ (     self:      Dict						)		->		int:
											snake_case__									=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*UpperCamelCase						)


							@slow
							def lowerCAmelCase_ (     self:      int						)		->		Any:
											for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															snake_case__									=     LiltModel.from_pretrained(UpperCamelCase						)
															self.assertIsNotNone(UpperCamelCase						)


@require_torch
@slow
class   __SCREAMING_SNAKE_CASE(	unittest.TestCase	):
							def lowerCAmelCase_ (     self:      List[Any]						)		->		Dict:
											snake_case__									=     LiltModel.from_pretrained('SCUT-DLVCLab/lilt-roberta-en-base'						).to(UpperCamelCase						)

											snake_case__									=     torch.tensor([[1, 2]]       ,	device=UpperCamelCase						)
											snake_case__									=     torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]       ,	device=UpperCamelCase						)

											# forward pass
											with torch.no_grad():
															snake_case__									=     model(input_ids=UpperCamelCase       ,	bbox=UpperCamelCase						)

											snake_case__									=     torch.Size([1, 2, 7_68]						)
											snake_case__									=     torch.tensor(
											    [[-0.0_653, 0.0_950, -0.0_061], [-0.0_545, 0.0_926, -0.0_324]]       ,	device=UpperCamelCase       ,	)

											self.assertTrue(outputs.last_hidden_state.shape       ,	UpperCamelCase						)
											self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]       ,	UpperCamelCase       ,	atol=1e-3						)						)




