import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowerCAmelCase__					:       Optional[int]		=       datasets.utils.logging.get_logger(__name__)





@dataclass
class 			__snake_case  (      datasets.BuilderConfig     ):
					__lowerCamelCase	 =       None





class 			__snake_case  (      datasets.ArrowBasedBuilder     ):
					__lowerCamelCase	 =       ArrowConfig




					def       __a      (       self		)					->  Any:



									'''simple docstring'''







									return datasets.DatasetInfo(features=self.config.features		)




					def       __a      (       self   ,	__UpperCamelCase		)					->  List[Any]:



									'''simple docstring'''







									if not self.config.data_files:
													raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""		)
									snake_case__ :     int 							=				dl_manager.download_and_extract(self.config.data_files		)
									if isinstance(_UpperCAmelCase   ,	(str, list, tuple)		):
													snake_case__ :     int 							=				data_files
													if isinstance(_UpperCAmelCase   ,	_UpperCAmelCase		):
																	snake_case__ :     List[Any] 							=				[files]
													# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
													snake_case__ :     int 							=				[dl_manager.iter_files(_UpperCAmelCase		) for file in files]
													return [datasets.SplitGenerator(name=datasets.Split.TRAIN   ,	gen_kwargs={'files': files}		)]
									snake_case__ :     int 							=				[]
									for split_name, files in data_files.items():
													if isinstance(_UpperCAmelCase   ,	_UpperCAmelCase		):
																	snake_case__ :     str 							=				[files]
													# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
													snake_case__ :     List[Any] 							=				[dl_manager.iter_files(_UpperCAmelCase		) for file in files]
													# Infer features is they are stoed in the arrow schema
													if self.info.features is None:
																	for file in itertools.chain.from_iterable(_UpperCAmelCase		):
																					with open(_UpperCAmelCase   ,	'rb'		) as f:
																									snake_case__ :     List[str] 							=				datasets.Features.from_arrow_schema(pa.ipc.open_stream(_UpperCAmelCase		).schema		)
																					break
													splits.append(datasets.SplitGenerator(name=_UpperCAmelCase   ,	gen_kwargs={'files': files}		)		)
									return splits




					def       __a      (       self   ,	__UpperCamelCase		)					->  pa.Table:



									'''simple docstring'''







									if self.info.features is not None:
													# more expensive cast to support nested features with keys in a different order
													# allows str <-> int/float or str to Audio for example
													snake_case__ :     int 							=				table_cast(_UpperCAmelCase   ,	self.info.features.arrow_schema		)
									return pa_table




					def       __a      (       self   ,	__UpperCamelCase		)					->  List[str]:



									'''simple docstring'''







									for file_idx, file in enumerate(itertools.chain.from_iterable(_UpperCAmelCase		)		):
													with open(_UpperCAmelCase   ,	'rb'		) as f:
																	try:
																					for batch_idx, record_batch in enumerate(pa.ipc.open_stream(_UpperCAmelCase		)		):
																									snake_case__ :     Dict 							=				pa.Table.from_batches([record_batch]		)
																									# Uncomment for debugging (will print the Arrow table size and elements)
																									# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																									# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																									yield F"""{file_idx}_{batch_idx}""", self._cast_table(_UpperCAmelCase		)
																	except ValueError as e:
																					logger.error(F"""Failed to read file '{file}' with error {type(_UpperCAmelCase		)}: {e}"""		)
																					raise



import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_multi_gpu,
    require_single_gpu,
    test_sync,
)
from accelerate.utils import patch_environment





class 			__snake_case  (      unittest.TestCase     ):




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



									'''simple docstring'''







									snake_case__ :     List[Any] 							=				inspect.getfile(accelerate.test_utils		)
									snake_case__ :     Tuple 							=				os.path.sep.join(mod_file.split(os.path.sep		)[:-1] + ['scripts', 'test_sync.py']		)




					@require_cpu
					def       __a      (       self		)					->  int:



									'''simple docstring'''







									debug_launcher(test_sync.main   ,	num_processes=1		)




					@require_cpu
					def       __a      (       self		)					->  Dict:



									'''simple docstring'''







									debug_launcher(test_sync.main		)




					@require_single_gpu
					def       __a      (       self		)					->  int:



									'''simple docstring'''







									test_sync.main()




					@require_multi_gpu
					def       __a      (       self		)					->  str:



									'''simple docstring'''







									print(F"""Found {torch.cuda.device_count()} devices."""		)
									snake_case__ :     Any 							=				['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
									with patch_environment(omp_num_threads=1		):
													execute_subprocess_async(__UpperCamelCase   ,	env=os.environ.copy()		)


