



'''simple docstring'''
import itertools
from dataclasses import dataclass
from io import StringIO
from typing import Optional

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


_SCREAMING_SNAKE_CASE =	datasets.utils.logging.get_logger(__name__)
@dataclass
class      __UpperCAmelCase							(    datasets.BuilderConfig							):



					'''simple docstring'''





					_UpperCamelCase  	=		None
					_UpperCamelCase  	=		"""utf-8"""
					_UpperCamelCase  	=		None
					_UpperCamelCase  	=		10 << 20  # 10MB
					_UpperCamelCase  	=		False
					_UpperCamelCase  	=		"""line"""





class      __UpperCAmelCase							(    datasets.ArrowBasedBuilder							):



					'''simple docstring'''





					_UpperCamelCase  	=		TextConfig



					def 						__snake_case	(  self    :	str)						->       List[str]:
								return datasets.DatasetInfo(features=self.config.features)



					def 						__snake_case	(  self    :	int			,					_lowercase    :	Union[str, Any])						->       Dict:





								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}')
								A_       		=   dl_manager.download_and_extract(self.config.data_files)
								if isinstance(__a			,					(str, list, tuple)):
											A_       		=   data_files
											if isinstance(__a			,					__a):
														A_       		=   [files]
											A_       		=   [dl_manager.iter_files(__a) for file in files]
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,					gen_kwargs={'files': files})]
								A_       		=   []
								for split_name, files in data_files.items():
											if isinstance(__a			,					__a):
														A_       		=   [files]
											A_       		=   [dl_manager.iter_files(__a) for file in files]
											splits.append(datasets.SplitGenerator(name=__a			,					gen_kwargs={'files': files}))
								return splits



					def 						__snake_case	(  self    :	int			,					_lowercase    :	Optional[Any])						->       List[str]:





								if self.config.features is not None:
											A_       		=   self.config.features.arrow_schema
											if all(not require_storage_cast(__a) for feature in self.config.features.values()):
														# cheaper cast
														A_       		=   pa_table.cast(__a)
											else:
														# more expensive cast; allows str <-> int/float or str to Audio for example
														A_       		=   table_cast(__a			,					__a)
											return pa_table
								else:
											return pa_table.cast(pa.schema({'text': pa.string()}))



					def 						__snake_case	(  self    :	str			,					_lowercase    :	str)						->       Optional[int]:
								A_       		=   list(self.config.features) if self.config.features is not None else ['text']
								for file_idx, file in enumerate(itertools.chain.from_iterable(__a)):
											# open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
											with open(__a			,					encoding=self.config.encoding			,					errors=self.config.errors) as f:
														if self.config.sample_by == "line":
																	A_       		=   0
																	while True:
																				A_       		=   f.read(self.config.chunksize)
																				if not batch:
																							break
																				batch += f.readline()  # finish current line
																				# StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
																				A_       		=   StringIO(__a).readlines()
																				if not self.config.keep_linebreaks:
																							A_       		=   [line.rstrip('\n') for line in batch]
																				A_       		=   pa.Table.from_arrays([pa.array(__a)]			,					names=__a)
																				# 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 (file_idx, batch_idx), self._cast_table(__a)
																				batch_idx += 1
														elif self.config.sample_by == "paragraph":
																	A_       		=   0
																	A_       		=   ''
																	while True:
																				A_       		=   f.read(self.config.chunksize)
																				if not new_batch:
																							break
																				batch += new_batch
																				batch += f.readline()  # finish current line
																				A_       		=   batch.split('\n\n')
																				A_       		=   pa.Table.from_arrays(
																				    [pa.array([example for example in batch[:-1] if example])]			,					names=__a)
																				# 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 (file_idx, batch_idx), self._cast_table(__a)
																				batch_idx += 1
																				A_       		=   batch[-1]
																	if batch:
																				A_       		=   pa.Table.from_arrays([pa.array([batch])]			,					names=__a)
																				yield (file_idx, batch_idx), self._cast_table(__a)
														elif self.config.sample_by == "document":
																	A_       		=   f.read()
																	A_       		=   pa.Table.from_arrays([pa.array([text])]			,					names=__a)
																	yield file_idx, self._cast_table(__a)










'''simple docstring'''
import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_SCREAMING_SNAKE_CASE =	namedtuple("ExampleNamedTuple", "a b c")
class      __UpperCAmelCase							(    unittest.TestCase							):



					'''simple docstring'''





					def 						__snake_case	(  self    :	List[Any])						->       Tuple:
								A_       		=   torch.randn(5			,					2)
								A_       		=   torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

								A_       		=   send_to_device(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta.cpu()			,					_lowercase))

								A_       		=   send_to_device((tensor, [tensor, tensor], 1)			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta[0].cpu()			,					_lowercase))
								self.assertIsInstance(resulta[1]			,					_lowercase)
								self.assertTrue(torch.equal(resulta[1][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta[1][1].cpu()			,					_lowercase))
								self.assertEqual(resulta[2]			,					1)

								A_       		=   send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta['a'].cpu()			,					_lowercase))
								self.assertIsInstance(resulta['b']			,					_lowercase)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta['b'][1].cpu()			,					_lowercase))
								self.assertEqual(resulta['c']			,					1)

								A_       		=   send_to_device(ExampleNamedTuple(a=_lowercase			,					b=[tensor, tensor]			,					c=1)			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta.a.cpu()			,					_lowercase))
								self.assertIsInstance(resulta.b			,					_lowercase)
								self.assertTrue(torch.equal(resulta.b[0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta.b[1].cpu()			,					_lowercase))
								self.assertEqual(resulta.c			,					1)

								A_       		=   send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1})			,					_lowercase)
								self.assertIsInstance(_lowercase			,					_lowercase)
								self.assertTrue(torch.equal(resulta['a'].cpu()			,					_lowercase))
								self.assertIsInstance(resulta['b']			,					_lowercase)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()			,					_lowercase))
								self.assertTrue(torch.equal(resulta['b'][1].cpu()			,					_lowercase))
								self.assertEqual(resulta['c']			,					1)



					def 						__snake_case	(  self    :	Dict)						->       str:
								with self.assertRaises(_lowercase) as cm:
											A_       		=   recursively_apply(torch.tensor			,					(torch.tensor(1), 1)			,					error_on_other_type=_lowercase)
								self.assertEqual(
								    str(cm.exception)			,					'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'			,					)



					def 						__snake_case	(  self    :	Union[str, Any])						->       List[str]:
								A_       		=   torch.tensor([1, 2, 3, 4, 5])
								self.assertEqual(listify(_lowercase)			,					[1, 2, 3, 4, 5])

								A_       		=   torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
								self.assertEqual(listify(_lowercase)			,					[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])

								A_       		=   torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]])
								self.assertEqual(
								    listify(_lowercase)			,					[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]])



					def 						__snake_case	(  self    :	List[Any])						->       Union[str, Any]:
								with patch_environment(aa=1			,					BB=2):
											self.assertEqual(os.environ.get('AA')			,					'1')
											self.assertEqual(os.environ.get('BB')			,					'2')

								self.assertNotIn('AA'			,					os.environ)
								self.assertNotIn('BB'			,					os.environ)



					def 						__snake_case	(  self    :	Any)						->       str:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   extract_model_from_parallel(_lowercase			,					keep_fpaa_wrapper=_lowercase)
								A_       		=   pickle.dumps(_lowercase)



					@require_cuda
					def 						__snake_case	(  self    :	Union[str, Any])						->       Union[str, Any]:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   extract_model_from_parallel(_lowercase			,					keep_fpaa_wrapper=_lowercase)
								A_       		=   pickle.dumps(_lowercase)



					@require_cuda
					@require_torch_min_version(version='2.0')
					def 						__snake_case	(  self    :	Optional[Any])						->       Dict:
								A_       		=   RegressionModel()
								A_       		=   model.forward
								A_       		=   torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
								A_       		=   convert_outputs_to_fpaa(model.forward)
								A_       		=   torch.compile(model.forward			,					backend='inductor')
								A_       		=   torch.randn(4			,					10).cuda()
								A_       		=   model(_lowercase)



					def 						__snake_case	(  self    :	Union[str, Any])						->       Optional[Any]:
								A_       		=   RegressionModel()
								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								A_       		=   torch.nn.parallel.DataParallel(_lowercase)
								A_       		=   extract_model_from_parallel(_lowercase)

								self.assertEqual(_lowercase			,					_lowercase)



					@require_torch_min_version(version='2.0')
					def 						__snake_case	(  self    :	str)						->       int:
								A_       		=   RegressionModel()
								A_       		=   torch.compile(_lowercase)

								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								A_       		=   torch.nn.parallel.DataParallel(_lowercase)
								A_       		=   torch.compile(_lowercase)
								A_       		=   extract_model_from_parallel(_lowercase)

								self.assertEqual(compiled_model._orig_mod			,					compiled_model_unwrapped._orig_mod)



					def 						__snake_case	(  self    :	List[str])						->       Optional[int]:
								self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3])])			,					torch.device('cpu'))
								self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3])})			,					torch.device('cpu'))
								self.assertIsNone(find_device([1, 'a']))




