





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__    (	_UpperCAmelCase  ):


      def __init__(self	:				Optional[int]    ,      a__	:				Optional[int]=0.0_1    ,      a__	:				Optional[Any]=1000				):





         """simple docstring"""



         __snake_case								=       p_stop
         __snake_case								=       max_length


      def __iter__(self	:				str				):





         """simple docstring"""



         __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        a			(self	:				Optional[int]    ,      a__	:				Any    ,      a__	:				Tuple    ,      a__	:				Any=False    ,      a__	:				int=True				):





         """simple docstring"""



         __snake_case								=       [
             BatchSamplerShard(a__    ,      2    ,      a__    ,      split_batches=a__    ,      even_batches=a__				)
             for i in range(2				)
         ]
         __snake_case								=       [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        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       BatchSampler(range(24				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__				)

         __snake_case								=       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.
         __snake_case								=       BatchSampler(range(21				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__				)

         __snake_case								=       BatchSampler(range(21				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__				)

         # 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=a__				)
         __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(a__    ,      a__				)

         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(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.
         __snake_case								=       BatchSampler(range(20				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__				)

         __snake_case								=       BatchSampler(range(20				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__				)

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

         __snake_case								=       BatchSampler(range(2				)    ,      batch_size=3    ,      drop_last=a__				)
         __snake_case								=       [[], []]
         self.check_batch_sampler_shards(a__    ,      a__				)


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       BatchSampler(range(24				)    ,      batch_size=4    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      split_batches=a__				)

         __snake_case								=       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.
         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=4    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      split_batches=a__				)

         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=4    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      split_batches=a__				)

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

         __snake_case								=       BatchSampler(range(21				)    ,      batch_size=4    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      split_batches=a__				)

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

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


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





         """simple docstring"""



         __snake_case								=       BatchSampler(range(24				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      even_batches=a__				)

         __snake_case								=       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.
         __snake_case								=       BatchSampler(range(21				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      even_batches=a__				)

         __snake_case								=       BatchSampler(range(21				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(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.
         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      even_batches=a__				)

         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(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.
         __snake_case								=       BatchSampler(range(20				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      even_batches=a__				)

         __snake_case								=       BatchSampler(range(20				)    ,      batch_size=3    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      even_batches=a__				)

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

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


      def        a			(self	:				Tuple				):





         """simple docstring"""



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

         __snake_case								=       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.
         __snake_case								=       BatchSampler(range(22				)    ,      batch_size=4    ,      drop_last=a__				)
         __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(a__    ,      a__    ,      split_batches=a__    ,      even_batches=a__				)

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

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

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

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


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
         __snake_case								=       [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        a			(self	:				Tuple    ,      a__	:				str    ,      a__	:				Any    ,      a__	:				Union[str, Any]    ,      a__	:				List[str]=False    ,      a__	:				Tuple=2    ,      a__	:				Union[str, Any]=False				):





         """simple docstring"""



         random.seed(a__				)
         __snake_case								=       list(a__				)

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

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

         __snake_case								=       []
         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        a			(self	:				Optional[int]				):





         """simple docstring"""



         __snake_case								=       42
         __snake_case								=       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
         __snake_case								=       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        a			(self	:				List[str]				):





         """simple docstring"""



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


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __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        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case								=       DataLoader(list(range(16				)				)    ,      batch_size=4				)
         __snake_case								=       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        a			(self	:				str				):





         """simple docstring"""



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





         """simple docstring"""



         Accelerator()
         __snake_case								=       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				)








from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


snake_case_           =      logging.get_logger(__name__)

snake_case_           =      {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

snake_case_           =      {
    'vocab_file': {
        'allegro/herbert-base-cased': 'https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json'
    },
    'merges_file': {
        'allegro/herbert-base-cased': 'https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt'
    },
}

snake_case_           =      {'allegro/herbert-base-cased': 514}
snake_case_           =      {}

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase  ):
      A_						:    Dict							=   VOCAB_FILES_NAMES
      A_						:    Optional[Any]							=   PRETRAINED_VOCAB_FILES_MAP
      A_						:    Optional[int]							=   PRETRAINED_INIT_CONFIGURATION
      A_						:    int							=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      A_						:    Any							=   HerbertTokenizer


      def __init__(self	:				Dict    ,      a__	:				Tuple=None    ,      a__	:				Optional[int]=None    ,      a__	:				List[str]=None    ,      a__	:				Optional[int]="<s>"    ,      a__	:				Optional[Any]="<unk>"    ,      a__	:				Any="<pad>"    ,      a__	:				List[Any]="<mask>"    ,      a__	:				Any="</s>"    ,      **a__	:				Tuple    ,      ):





         """simple docstring"""



         super().__init__(
             a__    ,      a__    ,      tokenizer_file=a__    ,      cls_token=a__    ,      unk_token=a__    ,      pad_token=a__    ,      mask_token=a__    ,      sep_token=a__    ,      **a__    ,      )


      def        a			(self	:				List[str]    ,      a__	:				List[int]    ,      a__	:				Optional[List[int]] = None				):





         """simple docstring"""



         __snake_case								=       [self.cls_token_id]
         __snake_case								=       [self.sep_token_id]
         if token_ids_a is None:
            return cls + token_ids_a + sep

         return cls + token_ids_a + sep + token_ids_a + sep


      def        a			(self	:				List[str]    ,      a__	:				List[int]    ,      a__	:				Optional[List[int]] = None    ,      a__	:				bool = False				):





         """simple docstring"""



         if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_a=a__    ,      token_ids_a=a__    ,      already_has_special_tokens=a__				)

         if token_ids_a is None:
            return [1] + ([0] * len(a__				)) + [1]
         return [1] + ([0] * len(a__				)) + [1] + ([0] * len(a__				)) + [1]


      def        a			(self	:				Optional[int]    ,      a__	:				List[int]    ,      a__	:				Optional[List[int]] = None				):





         """simple docstring"""



         __snake_case								=       [self.sep_token_id]
         __snake_case								=       [self.cls_token_id]

         if token_ids_a is None:
            return len(cls + token_ids_a + sep				) * [0]
         return len(cls + token_ids_a + sep				) * [0] + len(token_ids_a + sep				) * [1]


      def        a			(self	:				int    ,      a__	:				str    ,      a__	:				Optional[str] = None				):





         """simple docstring"""



         __snake_case								=       self._tokenizer.model.save(a__    ,      name=a__				)
         return tuple(a__				)

