


'''simple docstring'''

import multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader
class   __lowerCAmelCase						(lowercase_					):





      '''simple docstring'''






      def __init__(self   :   Dict       ,						UpperCamelCase   :   NestedDataStructureLike[PathLike]       ,						UpperCamelCase   :   Optional[NamedSplit] = None       ,						UpperCamelCase   :   Optional[Features] = None       ,						UpperCamelCase   :   str = None       ,						UpperCamelCase   :   bool = False       ,						UpperCamelCase   :   bool = False       ,						UpperCamelCase   :   Optional[int] = None       ,						**UpperCamelCase   :   str       ,						):

             '''simple docstring'''







             super().__init__(
                 UpperCamelCase       ,						split=UpperCamelCase       ,						features=UpperCamelCase       ,						cache_dir=UpperCamelCase       ,						keep_in_memory=UpperCamelCase       ,						streaming=UpperCamelCase       ,						num_proc=UpperCamelCase       ,						**UpperCamelCase       ,						)
             lowercase__				      =			path_or_paths if isinstance(UpperCamelCase       ,						UpperCamelCase			) else {self.split: path_or_paths}
             lowercase__				      =			Csv(
                 cache_dir=UpperCamelCase       ,						data_files=UpperCamelCase       ,						features=UpperCamelCase       ,						**UpperCamelCase       ,						)






      def 				UpperCamelCase__		(self   :   List[Any]			):

             '''simple docstring'''







             if self.streaming:
                    lowercase__				      =			self.builder.as_streaming_dataset(split=self.split			)
             # Build regular (map-style) dataset
             else:
                    lowercase__				      =			None
                    lowercase__				      =			None
                    lowercase__				      =			None
                    lowercase__				      =			None

                    self.builder.download_and_prepare(
                        download_config=UpperCamelCase       ,						download_mode=UpperCamelCase       ,						verification_mode=UpperCamelCase       ,						base_path=UpperCamelCase       ,						num_proc=self.num_proc       ,						)
                    lowercase__				      =			self.builder.as_dataset(
                        split=self.split       ,						verification_mode=UpperCamelCase       ,						in_memory=self.keep_in_memory			)
             return dataset






class   __lowerCAmelCase						:





      '''simple docstring'''






      def __init__(self   :   Dict       ,						UpperCamelCase   :   Dataset       ,						UpperCamelCase   :   Union[PathLike, BinaryIO]       ,						UpperCamelCase   :   Optional[int] = None       ,						UpperCamelCase   :   Optional[int] = None       ,						**UpperCamelCase   :   Union[str, Any]       ,						):

             '''simple docstring'''







             if num_proc is not None and num_proc <= 0:
                    raise ValueError(f"num_proc {num_proc} must be an integer > 0."			)

             lowercase__				      =			dataset
             lowercase__				      =			path_or_buf
             lowercase__				      =			batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
             lowercase__				      =			num_proc
             lowercase__				      =			'''utf-8'''
             lowercase__				      =			to_csv_kwargs



      def 				UpperCamelCase__		(self   :   Tuple			):

             '''simple docstring'''







             lowercase__				      =			self.to_csv_kwargs.pop('''path_or_buf'''       ,						UpperCamelCase			)
             lowercase__				      =			self.to_csv_kwargs.pop('''header'''       ,						UpperCamelCase			)
             lowercase__				      =			self.to_csv_kwargs.pop('''index'''       ,						UpperCamelCase			)

             if isinstance(self.path_or_buf       ,						(str, bytes, os.PathLike)			):
                    with open(self.path_or_buf       ,						'''wb+'''			) as buffer:
                           lowercase__				      =			self._write(file_obj=UpperCamelCase       ,						header=UpperCamelCase       ,						index=UpperCamelCase       ,						**self.to_csv_kwargs			)
             else:
                    lowercase__				      =			self._write(file_obj=self.path_or_buf       ,						header=UpperCamelCase       ,						index=UpperCamelCase       ,						**self.to_csv_kwargs			)
             return written



      def 				UpperCamelCase__		(self   :   Optional[Any]       ,						UpperCamelCase   :   str			):

             '''simple docstring'''







             lowercase__      ,lowercase__      ,lowercase__      ,lowercase__				      =			args

             lowercase__				      =			query_table(
                 table=self.dataset.data       ,						key=slice(UpperCamelCase       ,						offset + self.batch_size			)       ,						indices=self.dataset._indices       ,						)
             lowercase__				      =			batch.to_pandas().to_csv(
                 path_or_buf=UpperCamelCase       ,						header=header if (offset == 0) else False       ,						index=UpperCamelCase       ,						**UpperCamelCase			)
             return csv_str.encode(self.encoding			)






      def 				UpperCamelCase__		(self   :   List[str]       ,						UpperCamelCase   :   BinaryIO       ,						UpperCamelCase   :   List[str]       ,						UpperCamelCase   :   Optional[int]       ,						**UpperCamelCase   :   str			):

             '''simple docstring'''







             lowercase__				      =			0

             if self.num_proc is None or self.num_proc == 1:
                    for offset in logging.tqdm(
                        range(0       ,						len(self.dataset			)       ,						self.batch_size			)       ,						unit='''ba'''       ,						disable=not logging.is_progress_bar_enabled()       ,						desc='''Creating CSV from Arrow format'''       ,						):
                           lowercase__				      =			self._batch_csv((offset, header, index, to_csv_kwargs)			)
                           written += file_obj.write(UpperCamelCase			)

             else:
                    lowercase__      ,lowercase__				      =			len(self.dataset			), self.batch_size
                    with multiprocessing.Pool(self.num_proc			) as pool:
                           for csv_str in logging.tqdm(
                               pool.imap(
                                   self._batch_csv       ,						[(offset, header, index, to_csv_kwargs) for offset in range(0       ,						UpperCamelCase       ,						UpperCamelCase			)]       ,						)       ,						total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size       ,						unit='''ba'''       ,						disable=not logging.is_progress_bar_enabled()       ,						desc='''Creating CSV from Arrow format'''       ,						):
                                  written += file_obj.write(UpperCamelCase			)

             return written








'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCamelCase							:       str  						=     {'configuration_nat': ['NAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NatConfig']}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase							:       Tuple  						=     [
     'NAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'NatForImageClassification',
     'NatModel',
     'NatPreTrainedModel',
     'NatBackbone',
 ]

if TYPE_CHECKING:
 from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nat import (
      NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      NatBackbone,
      NatForImageClassification,
      NatModel,
      NatPreTrainedModel,
  )

else:
 import sys

 lowerCamelCase							:       Optional[Any]  						=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




