




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


UpperCamelCase__       							=      datasets.utils.logging.get_logger(__name__)





@dataclass
class   A							(       datasets.BuilderConfig  ):
   __UpperCAmelCase				:				List[Any]          =    None
   __UpperCAmelCase				:				Union[str, Any]          =    'utf-8'
   __UpperCAmelCase				:				Optional[int]          =    None
   __UpperCAmelCase				:				Optional[Any]          =    10 << 20  # 10MB
   __UpperCAmelCase				:				Optional[int]          =    False
   __UpperCAmelCase				:				Union[str, Any]          =    'line'





class   A							(       datasets.ArrowBasedBuilder  ):
   __UpperCAmelCase				:				List[str]          =    TextConfig


   def     lowercase_   (self  :		int	)		->       List[str]:
         """simple docstring"""



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


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


   def     lowercase_   (self  :		Optional[Any] ,						__UpperCAmelCase  :		Optional[Any]	)		->       pa.Table:
         """simple docstring"""



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





   def     lowercase_   (self  :		Tuple ,						__UpperCAmelCase  :		Tuple	)		->       List[str]:
         """simple docstring"""



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




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,
)


UpperCamelCase__       							=      namedtuple('ExampleNamedTuple', 'a b c')





class   A							(       unittest.TestCase  ):


   def     lowercase_   (self  :		str	)		->       List[Any]:
         """simple docstring"""



         UpperCAmelCase__			=						torch.randn(5 ,						2	)
         UpperCAmelCase__			=						torch.device("cuda"	) if torch.cuda.is_available() else torch.device("cpu"	)

         UpperCAmelCase__			=						send_to_device(__UpperCAmelCase ,						__UpperCAmelCase	)
         self.assertTrue(torch.equal(resulta.cpu() ,						__UpperCAmelCase	)	)

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

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

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

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


   def     lowercase_   (self  :		int	)		->       int:
         """simple docstring"""



         with self.assertRaises(__UpperCAmelCase	) as cm:
               UpperCAmelCase__			=						recursively_apply(torch.tensor ,						(torch.tensor(1	), 1) ,						error_on_other_type=__UpperCAmelCase	)
         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     lowercase_   (self  :		List[str]	)		->       Any:
         """simple docstring"""



         UpperCAmelCase__			=						torch.tensor([1, 2, 3, 4, 5]	)
         self.assertEqual(listify(__UpperCAmelCase	) ,						[1, 2, 3, 4, 5]	)

         UpperCAmelCase__			=						torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)
         self.assertEqual(listify(__UpperCAmelCase	) ,						[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)

         UpperCAmelCase__			=						torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)
         self.assertEqual(
             listify(__UpperCAmelCase	) ,						[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)


   def     lowercase_   (self  :		Union[str, Any]	)		->       Dict:
         """simple docstring"""



         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     lowercase_   (self  :		int	)		->       str:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase ,						keep_fpaa_wrapper=__UpperCAmelCase	)
         UpperCAmelCase__			=						pickle.dumps(__UpperCAmelCase	)


   @require_cuda
   def     lowercase_   (self  :		List[str]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						extract_model_from_parallel(__UpperCAmelCase ,						keep_fpaa_wrapper=__UpperCAmelCase	)
         UpperCAmelCase__			=						pickle.dumps(__UpperCAmelCase	)


   @require_cuda
   @require_torch_min_version(version="2.0"	)
   def     lowercase_   (self  :		Optional[int]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						model.forward
         UpperCAmelCase__			=						torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
         UpperCAmelCase__			=						convert_outputs_to_fpaa(model.forward	)
         UpperCAmelCase__			=						torch.compile(model.forward ,						backend="inductor"	)
         UpperCAmelCase__			=						torch.randn(4 ,						1_0	).cuda()
         UpperCAmelCase__			=						model(__UpperCAmelCase	)


   def     lowercase_   (self  :		List[Any]	)		->       Union[str, Any]:
         """simple docstring"""



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

         self.assertEqual(__UpperCAmelCase ,						__UpperCAmelCase	)


   @require_torch_min_version(version="2.0"	)
   def     lowercase_   (self  :		Union[str, Any]	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						RegressionModel()
         UpperCAmelCase__			=						torch.compile(__UpperCAmelCase	)

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

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





   def     lowercase_   (self  :		int	)		->       int:
         """simple docstring"""



         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"]	)	)



