





'''simple docstring'''






import os
import re
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


A_		:  Dict		 =						logging.get_logger(__name__)

A_		:  Union[str, Any]		 =						{
    "vocab_file": "vocab.txt",
    "merges_file": "bpe.codes",
}

A_		:  List[Any]		 =						{
    "vocab_file": {
        "vinai/phobert-base": "https://huggingface.co/vinai/phobert-base/resolve/main/vocab.txt",
        "vinai/phobert-large": "https://huggingface.co/vinai/phobert-large/resolve/main/vocab.txt",
    },
    "merges_file": {
        "vinai/phobert-base": "https://huggingface.co/vinai/phobert-base/resolve/main/bpe.codes",
        "vinai/phobert-large": "https://huggingface.co/vinai/phobert-large/resolve/main/bpe.codes",
    },
}

A_		:  List[str]		 =						{
    "vinai/phobert-base": 256,
    "vinai/phobert-large": 256,
}



def        UpperCamelCase__		(						__magic_name__  :   str      )				->			str:
 '''simple docstring'''






 snake_case__     :					int    =  set()
 snake_case__     :					Tuple    =  word[0]
 for char in word[1:]:
  pairs.add((prev_char, char)      )
  snake_case__     :					Dict    =  char

 snake_case__     :					int    =  set(__magic_name__      )
 return pairs





class 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   lowerCamelCase__    							=      VOCAB_FILES_NAMES
   lowerCamelCase__    							=      PRETRAINED_VOCAB_FILES_MAP
   lowerCamelCase__    							=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES

   def __init__(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE="<s>"	,		__SCREAMING_SNAKE_CASE="</s>"	,		__SCREAMING_SNAKE_CASE="</s>"	,		__SCREAMING_SNAKE_CASE="<s>"	,		__SCREAMING_SNAKE_CASE="<unk>"	,		__SCREAMING_SNAKE_CASE="<pad>"	,		__SCREAMING_SNAKE_CASE="<mask>"	,		**__SCREAMING_SNAKE_CASE	,		):
    super().__init__(
        bos_token=__SCREAMING_SNAKE_CASE	,		eos_token=__SCREAMING_SNAKE_CASE	,		unk_token=__SCREAMING_SNAKE_CASE	,		sep_token=__SCREAMING_SNAKE_CASE	,		cls_token=__SCREAMING_SNAKE_CASE	,		pad_token=__SCREAMING_SNAKE_CASE	,		mask_token=__SCREAMING_SNAKE_CASE	,		**__SCREAMING_SNAKE_CASE	,		)

    snake_case__     :					Dict    =  vocab_file
    snake_case__     :					Optional[Any]    =  merges_file

    snake_case__     :					Dict    =  {}
    snake_case__     :					Any    =  0
    snake_case__     :					int    =  1
    snake_case__     :					int    =  2
    snake_case__     :					List[Any]    =  3

    self.add_from_file(__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[Any]    =  {v: k for k, v in self.encoder.items()}

    with open(__SCREAMING_SNAKE_CASE	,		encoding="""utf-8"""       ) as merges_handle:
     snake_case__     :					Any    =  merges_handle.read().split("""\n"""       )[:-1]
    snake_case__     :					int    =  [tuple(merge.split()[:-1]       ) for merge in merges]
    snake_case__     :					List[str]    =  dict(zip(__SCREAMING_SNAKE_CASE	,		range(len(__SCREAMING_SNAKE_CASE       )       )       )       )
    snake_case__     :					List[str]    =  {}

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = None       ):

    if token_ids_a is None:
     return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
    snake_case__     :					int    =  [self.cls_token_id]
    snake_case__     :					List[str]    =  [self.sep_token_id]
    return cls + token_ids_a + sep + sep + token_ids_a + sep

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = None	,		__SCREAMING_SNAKE_CASE = False       ):

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

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

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = None       ):
    snake_case__     :					Any    =  [self.sep_token_id]
    snake_case__     :					List[str]    =  [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 + sep + token_ids_a + sep       ) * [0]

   @property
   def 		__UpperCamelCase						(							self       ):
    return len(self.encoder       )

   def 		__UpperCamelCase						(							self       ):
    return dict(self.encoder	,		**self.added_tokens_encoder       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    if token in self.cache:
     return self.cache[token]
    snake_case__     :					List[Any]    =  tuple(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[Any]    =  tuple(list(word[:-1]       ) + [word[-1] + """</w>"""]       )
    snake_case__     :					Any    =  get_pairs(__SCREAMING_SNAKE_CASE       )

    if not pairs:
     return token

    while True:
     snake_case__     :					Optional[Any]    =  min(__SCREAMING_SNAKE_CASE	,		key=lambda __SCREAMING_SNAKE_CASE       : self.bpe_ranks.get(__SCREAMING_SNAKE_CASE	,		float("""inf"""       )       )       )
     if bigram not in self.bpe_ranks:
      break
     snake_case__  ,							snake_case__     :					Tuple    =  bigram
     snake_case__     :					Dict    =  []
     snake_case__     :					str    =  0
     while i < len(__SCREAMING_SNAKE_CASE       ):
      try:
       snake_case__     :					Tuple    =  word.index(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
      except ValueError:
       new_word.extend(word[i:]       )
       break
      else:
       new_word.extend(word[i:j]       )
       snake_case__     :					List[str]    =  j

      if word[i] == first and i < len(__SCREAMING_SNAKE_CASE       ) - 1 and word[i + 1] == second:
       new_word.append(first + second       )
       i += 2
      else:
       new_word.append(word[i]       )
       i += 1
     snake_case__     :					Dict    =  tuple(__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Tuple    =  new_word
     if len(__SCREAMING_SNAKE_CASE       ) == 1:
      break
     else:
      snake_case__     :					Union[str, Any]    =  get_pairs(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[Any]    =  """@@ """.join(__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Optional[int]    =  word[:-4]
    snake_case__     :					Dict    =  word
    return word

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[int]    =  []

    snake_case__     :					Any    =  re.findall(R"""\S+\n?"""	,		__SCREAMING_SNAKE_CASE       )

    for token in words:
     split_tokens.extend(list(self.bpe(__SCREAMING_SNAKE_CASE       ).split(""" """       )       )       )
    return split_tokens

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    return self.encoder.get(__SCREAMING_SNAKE_CASE	,		self.encoder.get(self.unk_token       )       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    return self.decoder.get(__SCREAMING_SNAKE_CASE	,		self.unk_token       )

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[Any]    =  """ """.join(__SCREAMING_SNAKE_CASE       ).replace("""@@ """	,		""""""       ).strip()
    return out_string

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE = None       ):
    if not os.path.isdir(__SCREAMING_SNAKE_CASE       ):
     logger.error(f"Vocabulary path ({save_directory}) should be a directory"       )
     return
    snake_case__     :					Optional[int]    =  os.path.join(
        __SCREAMING_SNAKE_CASE	,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]       )
    snake_case__     :					Any    =  os.path.join(
        __SCREAMING_SNAKE_CASE	,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]       )

    if os.path.abspath(self.vocab_file       ) != os.path.abspath(__SCREAMING_SNAKE_CASE       ):
     copyfile(self.vocab_file	,		__SCREAMING_SNAKE_CASE       )

    if os.path.abspath(self.merges_file       ) != os.path.abspath(__SCREAMING_SNAKE_CASE       ):
     copyfile(self.merges_file	,		__SCREAMING_SNAKE_CASE       )

    return out_vocab_file, out_merge_file


   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    if isinstance(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
     try:
      with open(__SCREAMING_SNAKE_CASE	,		"""r"""	,		encoding="""utf-8"""       ) as fd:
       self.add_from_file(__SCREAMING_SNAKE_CASE       )
     except FileNotFoundError as fnfe:
      raise fnfe
     except UnicodeError:
      raise Exception(f"Incorrect encoding detected in {f}, please rebuild the dataset"       )
     return

    snake_case__     :					Tuple    =  f.readlines()
    for lineTmp in lines:
     snake_case__     :					Any    =  lineTmp.strip()
     snake_case__     :					Optional[Any]    =  line.rfind(""" """       )
     if idx == -1:
      raise ValueError("""Incorrect dictionary format, expected '<token> <cnt>'"""       )
     snake_case__     :					Optional[int]    =  line[:idx]
     snake_case__     :					Union[str, Any]    =  len(self.encoder       )






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_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					Optional[int]     ,				UpperCamelCase__     :					Optional[Any]=0.01     ,				UpperCamelCase__     :					Tuple=1000       )     ->      str:







    """simple docstring"""






    __magic_name__								=			p_stop
    __magic_name__								=			max_length





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







    """simple docstring"""






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







class 			UpperCAmelCase_     (					unittest.TestCase	):


  '''simple docstring'''







  def 				_lowercase				(					self     :					Optional[Any]     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					str=False     ,				UpperCamelCase__     :					Optional[Any]=True       )     ->      Dict:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

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

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

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

    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1, 0]], [[1, 0, 1]]]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__       )
  def 				_lowercase				(					self     :					Tuple       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], [[0, 1]]]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					int       )     ->      str:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

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

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

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

    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				even_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Optional[int]       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Dict       )     ->      str:







    """simple docstring"""






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







    """simple docstring"""






    random.seed(UpperCamelCase__       )
    __magic_name__								=			list(UpperCamelCase__       )

    __magic_name__								=			[
        IterableDatasetShard(
            UpperCamelCase__     ,				batch_size=UpperCamelCase__     ,				drop_last=UpperCamelCase__     ,				num_processes=UpperCamelCase__     ,				process_index=UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				)
        for i in range(UpperCamelCase__       )
    ]
    __magic_name__								=			[]
    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__       )       )

    __magic_name__								=			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
    __magic_name__								=			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       )

    __magic_name__								=			[]
    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 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			42
    __magic_name__								=			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
    __magic_name__								=			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 				_lowercase				(					self     :					List[Any]       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(16       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			SkipBatchSampler(UpperCamelCase__     ,				2       )
    self.assertListEqual(list(UpperCamelCase__       )     ,				[[8, 9, 10, 11], [12, 13, 14, 15]]       )
  def 				_lowercase				(					self     :					Tuple       )     ->      str:







    """simple docstring"""






    __magic_name__								=			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 				_lowercase				(					self     :					str       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			DataLoader(list(range(16       )       )     ,				batch_size=4       )
    __magic_name__								=			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 				_lowercase				(					self     :					List[str]       )     ->      Tuple:







    """simple docstring"""






    __magic_name__								=			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 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






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

