


'''simple docstring'''





import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


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




@dataclass
class 							UpperCAmelCase ( datasets.BuilderConfig	):
 UpperCAmelCase				=				None




class 							UpperCAmelCase ( datasets.ArrowBasedBuilder	):
 UpperCAmelCase				=				ArrowConfig


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    return datasets.DatasetInfo(features=self.config.features   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any] ,					__lowerCamelCase       :  Union[str, Any]   ):



    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__  :int            =			dl_manager.download_and_extract(self.config.data_files   )
    if isinstance(__lowerCamelCase ,					(str, list, tuple)   ):
       UpperCAmelCase__  :List[Any]            =			data_files
       if isinstance(__lowerCamelCase ,					__lowerCamelCase   ):
          UpperCAmelCase__  :str            =			[files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       UpperCAmelCase__  :Optional[int]            =			[dl_manager.iter_files(__lowerCamelCase   ) for file in files]
       return [datasets.SplitGenerator(name=datasets.Split.TRAIN ,					gen_kwargs={'''files''': files}   )]
    UpperCAmelCase__  :Union[str, Any]            =			[]
    for split_name, files in data_files.items():
       if isinstance(__lowerCamelCase ,					__lowerCamelCase   ):
          UpperCAmelCase__  :List[str]            =			[files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       UpperCAmelCase__  :List[str]            =			[dl_manager.iter_files(__lowerCamelCase   ) 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(__lowerCamelCase   ):
             with open(__lowerCamelCase ,					'''rb'''   ) as f:
                UpperCAmelCase__  :Dict            =			datasets.Features.from_arrow_schema(pa.ipc.open_stream(__lowerCamelCase   ).schema   )
             break
       splits.append(datasets.SplitGenerator(name=__lowerCamelCase ,					gen_kwargs={'''files''': files}   )   )
    return splits


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any] ,					__lowerCamelCase       :  pa.Table   ):



    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
       UpperCAmelCase__  :Optional[Any]            =			table_cast(__lowerCamelCase ,					self.info.features.arrow_schema   )
    return pa_table





 def 					__SCREAMING_SNAKE_CASE				(			self       :  str ,					__lowerCamelCase       :  Any   ):



    for file_idx, file in enumerate(itertools.chain.from_iterable(__lowerCamelCase   )   ):
       with open(__lowerCamelCase ,					'''rb'''   ) as f:
          try:
             for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__lowerCamelCase   )   ):
                UpperCAmelCase__  :Any            =			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(__lowerCamelCase   )
          except ValueError as e:
             logger.error(f'''Failed to read file \'{file}\' with error {type(__lowerCamelCase   )}: {e}'''   )
             raise







'''simple docstring'''





from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
    import torch

__lowerCamelCase       				=	logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
    from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):


    def 					a__    (				cls	:			Optional[int]	):
       return cls


def 					a__    (				UpperCamelCase_	:			Union[Tuple, List],      UpperCamelCase_	:			Optional[Union[List["torch.Generator"], "torch.Generator"]] = None,      UpperCamelCase_	:			Optional["torch.device"] = None,      UpperCamelCase_	:			Optional["torch.dtype"] = None,      UpperCamelCase_	:			Optional["torch.layout"] = None,      ):
   UpperCAmelCase__  :Union[str, Any]            =			device
   UpperCAmelCase__  :List[Any]            =			shape[0]

   UpperCAmelCase__  :Tuple            =			layout or torch.strided
   UpperCAmelCase__  :Union[str, Any]            =			device or torch.device('''cpu'''	)

   if generator is not None:
      UpperCAmelCase__  :str            =			generator.device.type if not isinstance(UpperCamelCase_,      UpperCamelCase_	) else generator[0].device.type
      if gen_device_type != device.type and gen_device_type == "cpu":
         UpperCAmelCase__  :Tuple            =			'''cpu'''
         if device != "mps":
            logger.info(
                F'''The passed generator was created on \'cpu\' even though a tensor on {device} was expected.'''
                F''' Tensors will be created on \'cpu\' and then moved to {device}. Note that one can probably'''
                F''' slighly speed up this function by passing a generator that was created on the {device} device.'''	)
      elif gen_device_type != device.type and gen_device_type == "cuda":
         raise ValueError(F'''Cannot generate a {device} tensor from a generator of type {gen_device_type}.'''	)

   if isinstance(UpperCamelCase_,      UpperCamelCase_	):
      UpperCAmelCase__  :List[str]            =			(1,) + shape[1:]
      UpperCAmelCase__  :Optional[int]            =			[
          torch.randn(UpperCamelCase_,      generator=generator[i],      device=UpperCamelCase_,      dtype=UpperCamelCase_,      layout=UpperCamelCase_	)
          for i in range(UpperCamelCase_	)
      ]
      UpperCAmelCase__  :int            =			torch.cat(UpperCamelCase_,      dim=0	).to(UpperCamelCase_	)
   else:
      UpperCAmelCase__  :Optional[int]            =			torch.randn(UpperCamelCase_,      generator=UpperCamelCase_,      device=UpperCamelCase_,      dtype=UpperCamelCase_,      layout=UpperCamelCase_	).to(UpperCamelCase_	)

   return latents


def 					a__    (				UpperCamelCase_	:			Dict	):
   if is_torch_version('''<''',      '''2.0.0'''	) or not hasattr(UpperCamelCase_,      '''_dynamo'''	):
      return False
   return isinstance(UpperCamelCase_,      torch._dynamo.eval_frame.OptimizedModule	)



