code: 
  
  
  import os
  import tempfile
  from pathlib import Path
  from unittest import TestCase
  
  import pyarrow as pa
  import pytest
  
  from datasets.arrow_dataset import Dataset
  from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
  from datasets.info import DatasetInfo
  from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo
  
  from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
  
  
  
  
  class 					_UpperCamelCase (      lowerCamelCase__	):
  
  
  
  
  
  
  
  							'''simple docstring'''
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				Optional[Any]		,				SCREAMING_SNAKE_CASE_  :				int		,				SCREAMING_SNAKE_CASE_  :				Any=False		):
  													_a ,	_a ,	_a     =			(
  													    filename_skip_take['filename'],
  													    filename_skip_take['skip'] if 'skip' in filename_skip_take else None,
  													    filename_skip_take['take'] if 'take' in filename_skip_take else None,
  													)
  													open(os.path.join(SCREAMING_SNAKE_CASE_		)		,				'wb'		).close()
  													_a     =			pa.Table.from_pydict({'filename': [Path(SCREAMING_SNAKE_CASE_		).name] * 1_0_0}		)
  													if take == -1:
  																			_a     =			len(SCREAMING_SNAKE_CASE_		) - skip
  													if skip is not None and take is not None:
  																			_a     =			pa_table.slice(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  													return pa_table
  
  
  
  
  class 					_UpperCamelCase (      lowerCamelCase__	):
  
  
  
  
  
  
  
  							'''simple docstring'''
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				str		):
  													_a     =			'my_name'
  													_a     =			SplitInfo(name='train'		,				num_examples=1_0_0		)
  													_a     =			SplitInfo(name='test'		,				num_examples=1_0_0		)
  													_a     =			[train_info, test_info]
  													_a     =			SplitDict()
  													split_dict.add(SCREAMING_SNAKE_CASE_		)
  													split_dict.add(SCREAMING_SNAKE_CASE_		)
  													_a     =			DatasetInfo(splits=SCREAMING_SNAKE_CASE_		)
  
  													with tempfile.TemporaryDirectory() as tmp_dir:
  																			_a     =			ReaderTest(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  
  																			_a     =			'test[:33%]'
  																			_a     =			Dataset(**reader.read(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)		)
  																			self.assertEqual(dset['filename'][0]		,				f"""{name}-test"""		)
  																			self.assertEqual(dset.num_rows		,				3_3		)
  																			self.assertEqual(dset.num_columns		,				1		)
  
  																			_a     =			['train', 'test[:33%]']
  																			_a     =			[Split.TRAIN, ReadInstruction.from_spec('test[:33%]'		)]
  																			for instructions in [instructionsa, instructionsa]:
  																									_a     =			[reader.read(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		) for instr in instructions]
  																									_a ,	_a     =			(Dataset(**SCREAMING_SNAKE_CASE_		) for dataset_kwargs in datasets_kwargs)
  																									self.assertEqual(train_dset['filename'][0]		,				f"""{name}-train"""		)
  																									self.assertEqual(train_dset.num_rows		,				1_0_0		)
  																									self.assertEqual(train_dset.num_columns		,				1		)
  																									self.assertIsInstance(train_dset.split		,				SCREAMING_SNAKE_CASE_		)
  																									self.assertEqual(str(train_dset.split		)		,				'train'		)
  																									self.assertEqual(test_dset['filename'][0]		,				f"""{name}-test"""		)
  																									self.assertEqual(test_dset.num_rows		,				3_3		)
  																									self.assertEqual(test_dset.num_columns		,				1		)
  																									self.assertIsInstance(test_dset.split		,				SCREAMING_SNAKE_CASE_		)
  																									self.assertEqual(str(test_dset.split		)		,				'test[:33%]'		)
  																									del train_dset, test_dset
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				Optional[int]		):
  													_a     =			'my_name'
  													_a     =			SplitInfo(name='train'		,				num_examples=1_0_0_0		,				shard_lengths=[1_0_0] * 1_0		)
  													_a     =			[train_info]
  													_a     =			SplitDict()
  													split_dict.add(SCREAMING_SNAKE_CASE_		)
  													_a     =			DatasetInfo(splits=SCREAMING_SNAKE_CASE_		)
  
  													with tempfile.TemporaryDirectory() as tmp_dir:
  																			_a     =			ReaderTest(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  
  																			_a     =			'train[:33%]'
  																			_a     =			Dataset(**reader.read(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)		)
  																			self.assertEqual(dset['filename'][0]		,				f"""{name}-train-00000-of-00010"""		)
  																			self.assertEqual(dset['filename'][-1]		,				f"""{name}-train-00003-of-00010"""		)
  																			self.assertEqual(dset.num_rows		,				3_3_0		)
  																			self.assertEqual(dset.num_columns		,				1		)
  
  
  
  
  
  
  
  							def 		_UpperCAmelCase	(		self  :				int		):
  													_a     =			SplitInfo(name='train'		,				num_examples=1_0_0		)
  													_a     =			SplitInfo(name='test'		,				num_examples=1_0_0		)
  													_a     =			SplitDict()
  													split_dict.add(SCREAMING_SNAKE_CASE_		)
  													split_dict.add(SCREAMING_SNAKE_CASE_		)
  													_a     =			DatasetInfo(splits=SCREAMING_SNAKE_CASE_		)
  
  													with tempfile.TemporaryDirectory() as tmp_dir:
  																			_a     =			ReaderTest(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		)
  
  																			_a     =			[
  																			    {'filename': os.path.join(SCREAMING_SNAKE_CASE_		,				'train'		)},
  																			    {'filename': os.path.join(SCREAMING_SNAKE_CASE_		,				'test'		), 'skip': 1_0, 'take': 1_0},
  																			]
  																			_a     =			Dataset(**reader.read_files(SCREAMING_SNAKE_CASE_		,				original_instructions='train+test[10:20]'		)		)
  																			self.assertEqual(dset.num_rows		,				1_1_0		)
  																			self.assertEqual(dset.num_columns		,				1		)
  																			del dset
  
  
  
  
  @pytest.mark.parametrize('in_memory'							,     [False, True]				)
  def 							SCREAMING_SNAKE_CASE					(  _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase				)    ->  Any:
  						_a     =			arrow_file
  						with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
  												_a     =			ArrowReader.read_table(_UpperCAmelCase							,     in_memory=_UpperCAmelCase				)
  						assert table.shape == dataset.data.shape
  						assert set(table.column_names				) == set(dataset.data.column_names				)
  						assert dict(table.to_pydict()				) == dict(dataset.data.to_pydict()				)  # to_pydict returns OrderedDict
  
  
  @pytest.mark.parametrize('in_memory'							,     [False, True]				)
  def 							SCREAMING_SNAKE_CASE					(  _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase				)    ->  List[str]:
  						_a     =			arrow_file
  						_a     =			ArrowReader(''							,     _UpperCAmelCase				)
  						with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
  												_a     =			reader.read_files([{'filename': filename}]							,     in_memory=_UpperCAmelCase				)
  						assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
  						_a     =			dataset_kwargs['arrow_table']
  						assert table.shape == dataset.data.shape
  						assert set(table.column_names				) == set(dataset.data.column_names				)
  						assert dict(table.to_pydict()				) == dict(dataset.data.to_pydict()				)  # to_pydict returns OrderedDict
  
  
  def 							SCREAMING_SNAKE_CASE					(  )    ->  Dict:
  						assert ReadInstruction('train'							,     to=10							,     unit='abs'				).to_spec() == "train[:10]"
  						assert ReadInstruction('train'							,     from_=-80							,     to=10							,     unit='%'				).to_spec() == "train[-80%:10%]"
  
  						_a     =			'train+test'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == spec_train_test
  
  						_a     =			'train[2:10]'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == spec_train_abs
  
  						_a     =			'train[15%:-20%]'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == spec_train_pct
  
  						_a     =			'train[:10%](closest)'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == "train[:10%]"
  
  						_a     =			'train[:10%](pct1_dropremainder)'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == spec_train_pct_rounding
  
  						_a     =			'train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)'
  						assert ReadInstruction.from_spec(_UpperCAmelCase				).to_spec() == spec_train_test_pct_rounding
  
  
  def 							SCREAMING_SNAKE_CASE					(  )    ->  List[str]:
  						_a     =			'dummy'
  						_a     =			[SplitInfo(name='train'							,     num_examples=100				)]
  						_a     =			'train[:33%]'
  						_a     =			'arrow'
  						_a     =			'prefix'
  
  						_a     =			make_file_instructions(_UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase				)
  						assert isinstance(_UpperCAmelCase							,     _UpperCAmelCase				)
  						assert file_instructions.num_examples == 33
  						assert file_instructions.file_instructions == [
  						    {"filename": os.path.join(_UpperCAmelCase							,     f"""{name}-train.arrow"""				), "skip": 0, "take": 33}
  						]
  
  						_a     =			[SplitInfo(name='train'							,     num_examples=100							,     shard_lengths=[10] * 10				)]
  						_a     =			make_file_instructions(_UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase				)
  						assert isinstance(_UpperCAmelCase							,     _UpperCAmelCase				)
  						assert file_instructions.num_examples == 33
  						assert file_instructions.file_instructions == [
  						    {"filename": os.path.join(_UpperCAmelCase							,     f"""{name}-train-00000-of-00010.arrow"""				), "skip": 0, "take": -1},
  						    {"filename": os.path.join(_UpperCAmelCase							,     f"""{name}-train-00001-of-00010.arrow"""				), "skip": 0, "take": -1},
  						    {"filename": os.path.join(_UpperCAmelCase							,     f"""{name}-train-00002-of-00010.arrow"""				), "skip": 0, "take": -1},
  						    {"filename": os.path.join(_UpperCAmelCase							,     f"""{name}-train-00003-of-00010.arrow"""				), "skip": 0, "take": 3},
  						]
  
  
  
  
  
  
  @pytest.mark.parametrize('name, expected_exception'							,     [(None, TypeError), ('', ValueError)]				)
  def 							SCREAMING_SNAKE_CASE					(  _UpperCAmelCase							,     _UpperCAmelCase				)    ->  str:
  						_a     =			[SplitInfo(name='train'							,     num_examples=100				)]
  						_a     =			'train'
  						_a     =			'arrow'
  						_a     =			'prefix_path'
  						with pytest.raises(_UpperCAmelCase				):
  												_a     =			make_file_instructions(_UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase							,     _UpperCAmelCase				)
  
  
  
style_context: 
  
  
  from .chunk_utils import chunk_layer
  from .data_transforms import make_atomaa_masks
  from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
  from .loss import compute_predicted_aligned_error, compute_tm
  from .protein import Protein as OFProtein
  from .protein import to_pdb
  from .rigid_utils import Rigid, Rotation
  from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims
  
  
  