code: 
  
  import os
  import sys
  from pathlib import Path
  
  import pytest
  
  from datasets import Dataset, IterableDataset
  from datasets.distributed import split_dataset_by_node
  
  from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch
  def        __SCREAMING_SNAKE_CASE    (      )				->							Any:
  	__A   : int								= Dataset.from_dict({"""i""": range(17					)}					)
  	__A   : Any								= len(a__					)
  	__A   : Union[str, Any]								= 3
  	__A   : int								= [
  	    split_dataset_by_node(a__ ,rank=a__ ,world_size=a__					) for rank in range(a__					)
  	]
  	assert sum(len(a__					) for ds in datasets_per_rank					) == full_size
  	assert len({tuple(x.values()					) for ds in datasets_per_rank for x in ds}					) == full_size
  def        __SCREAMING_SNAKE_CASE    (      )				->							Dict:
  	def gen():
  		return ({"i": i} for i in range(17					))
  
  	__A   : Tuple								= 3
  	__A   : Optional[int]								= IterableDataset.from_generator(a__					)
  	__A   : Optional[Any]								= len(list(a__					)					)
  	__A   : Optional[int]								= [
  	    split_dataset_by_node(a__ ,rank=a__ ,world_size=a__					) for rank in range(a__					)
  	]
  	assert sum(len(list(a__					)					) for ds in datasets_per_rank					) == full_size
  	assert len({tuple(x.values()					) for ds in datasets_per_rank for x in ds}					) == full_size
  @pytest.mark.parametrize("""shards_per_node""" ,[1, 2, 3]					)
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Optional[Any]					)				->							Tuple:
  	def gen(a__						:      Union[str, Any]					):
  		for shard in shards:
  			yield from ({"i": i, "shard": shard} for i in range(17					))
  
  	__A   : int								= 3
  	__A   : str								= shards_per_node * world_size
  	__A   : str								= {"""shards""": [f"""shard_{shard_idx}.txt""" for shard_idx in range(a__					)]}
  	__A   : Tuple								= IterableDataset.from_generator(a__ ,gen_kwargs=a__					)
  	__A   : Optional[Any]								= len(list(a__					)					)
  	assert full_ds.n_shards == world_size * shards_per_node
  	__A   : Tuple								= [
  	    split_dataset_by_node(a__ ,rank=a__ ,world_size=a__					) for rank in range(a__					)
  	]
  	assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
  	assert sum(len(list(a__					)					) for ds in datasets_per_rank					) == full_size
  	assert len({tuple(x.values()					) for ds in datasets_per_rank for x in ds}					) == full_size
  def        __SCREAMING_SNAKE_CASE    (      )				->							Any:
  	def gen():
  		return ({"i": i} for i in range(17					))
  
  	__A   : Optional[Any]								= 2
  	__A   : Union[str, Any]								= IterableDataset.from_generator(a__					)
  	__A   : Optional[int]								= len(list(a__					)					)
  
  	__A   : Tuple								= split_dataset_by_node(a__ ,rank=0 ,world_size=a__					).shuffle(seed=42					)
  	assert len(list(a__					)					) == 1 + full_size // world_size
  	with pytest.raises(a__					):
  		split_dataset_by_node(a__ ,rank=0 ,world_size=a__					).shuffle()
  
  	__A   : List[Any]								= split_dataset_by_node(full_ds.shuffle(seed=42					) ,rank=0 ,world_size=a__					)
  	assert len(list(a__					)					) == 1 + full_size // world_size
  	with pytest.raises(a__					):
  		split_dataset_by_node(full_ds.shuffle() ,rank=0 ,world_size=a__					)
  @pytest.mark.parametrize("""streaming""" ,[False, True]					)
  @require_torch
  @pytest.mark.skipif(os.name == """nt""" ,reason="""execute_subprocess_async doesn't support windows"""					)
  @pytest.mark.integration
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Any					)				->							Tuple:
  	__A   : Optional[int]								= 2
  	__A   : Optional[Any]								= get_torch_dist_unique_port()
  	__A   : str								= Path(a__					).resolve().parent / """distributed_scripts""" / """launch_torch_distributed.py"""
  	__A   : Union[str, Any]								= f"""
          -m torch.distributed.launch
          --nproc_per_node={nproc_per_node}
          --master_port={master_port}
          {test_script}
      """.split()
  	__A   : Tuple								= f"""
          --streaming={streaming}
      """.split()
  	__A   : List[Any]								= [sys.executable] + distributed_args + args
  	execute_subprocess_async(a__ ,env=os.environ.copy()					)
  
  
  
  
  
  
  
  @pytest.mark.parametrize(
      """nproc_per_node, num_workers""" ,[
          (2, 2),  # each node has 2 shards and each worker has 1 shards
          (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
      ] ,)
  @require_torch
  @pytest.mark.skipif(os.name == """nt""" ,reason="""execute_subprocess_async doesn't support windows"""					)
  @pytest.mark.integration
  def        __SCREAMING_SNAKE_CASE    (      a__						:      Any ,a__						:      Union[str, Any]					)				->							Tuple:
  	__A   : int								= True
  	__A   : int								= get_torch_dist_unique_port()
  	__A   : Optional[int]								= Path(a__					).resolve().parent / """distributed_scripts""" / """launch_torch_distributed.py"""
  	__A   : Any								= f"""
          -m torch.distributed.launch
          --nproc_per_node={nproc_per_node}
          --master_port={master_port}
          {test_script}
      """.split()
  	__A   : Any								= f"""
          --streaming={streaming}
          --num_workers={num_workers}
      """.split()
  	__A   : Any								= [sys.executable] + distributed_args + args
  	execute_subprocess_async(a__ ,env=os.environ.copy()					)
  
  
  
  
  
  
style_context: 
  
  import os
  import tempfile
  import unittest
  from typing import List
  
  from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
  from transformers.testing_utils import require_tokenizers
  from transformers.tokenization_utils import PreTrainedTokenizer
  from transformers.tokenization_utils_base import PreTrainedTokenizerBase
  
  
  
  
  
  
  @require_tokenizers
  class     lowerCamelCase_			(  unittest.TestCase  ):
  	_lowercase		:      List[str]    				=							EsmTokenizer
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    List[Any]       ):
  		super().setUp()
  		__A   : Dict								= tempfile.mkdtemp()
  		# fmt: off
  		__A   : List[str]								= ["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
  		# fmt: on
  		__A   : int								= os.path.join(self.tmpdirname				,			VOCAB_FILES_NAMES["""vocab_file"""]       )
  		with open(self.vocab_file				,			"""w"""				,			encoding="""utf-8"""       ) as vocab_writer:
  			vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]       )       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Dict				,			**__A	:    int       ):
  		return [self.get_tokenizer(**__A       )]
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Optional[int]				,			**__A	:    List[Any]       ):
  		return self.tokenizer_class.from_pretrained(self.tmpdirname				,			**__A       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    List[str]       ):
  		__A   : Union[str, Any]								= self.tokenizer_class(self.vocab_file       )
  
  		__A   : int								= tokenizer.tokenize("""LAGVS"""       )
  		self.assertListEqual(__A				,			["""L""", """A""", """G""", """V""", """S"""]       )
  		self.assertListEqual(tokenizer.convert_tokens_to_ids(__A       )				,			[4, 5, 6, 7, 8]       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    str       ):
  		__A   : Union[str, Any]								= self.tokenizer_class(self.vocab_file       )
  
  		__A   : int								= """LAGVS"""
  		self.assertListEqual(tokenizer.encode(__A       )				,			[0, 4, 5, 6, 7, 8, 2]       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    Dict       ):
  		__A   : str								= self.tokenizer_class(self.vocab_file       )
  
  		__A   : Tuple								= ["""LAGVS""", """WCB"""]
  		__A   : str								= tokenizer(__A				,			padding=__A       )["""input_ids"""]
  
  		self.assertListEqual(__A				,			[[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2]]       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    List[Any]       ):
  		__A   : Dict								= self.tokenizer_class(self.vocab_file       )
  
  		__A   : Any								= ["""LAGVS""", """WCB"""]
  		__A   : str								= tokenizer(__A				,			padding=__A       )["""input_ids"""]
  
  		self.assertListEqual(__A				,			[[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2, 1, 1]]       )
  
  
  
  
  	def 					lowerCAmelCase_					(    self	:    List[str]       ):
  		__A   : Optional[int]								= self.get_tokenizers(fast=__A       )
  		for tokenizer in tokenizers:
  			with self.subTest(F"""{tokenizer.__class__.__name__}"""       ):
  				__A   : Any								= """<unk>"""
  				__A   : str								= """<mask>"""
  
  				__A   : Optional[Any]								= tokenizer.tokenize(__A       )
  				__A   : str								= tokenizer.tokenize(__A       )
  
  				self.assertEqual(len(__A       )				,			1       )
  				self.assertEqual(len(__A       )				,			1       )
  				self.assertEqual(token_a[0]				,			__A       )
  				self.assertEqual(token_a[0]				,			__A       )
  
  
  
  
  
  