

"""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  (					_UpperCAmelCase      ):


  '''simple docstring'''







  def __init__(self       ,					_lowerCamelCase       ,					_lowerCamelCase = None       ,					_lowerCamelCase = None       ,					_lowerCamelCase = None       ,					_lowerCamelCase = False       ,					_lowerCamelCase = False       ,					_lowerCamelCase = None       ,					**_lowerCamelCase       ,					):


        """simple docstring"""

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




  def    _a			(self  ):


        """simple docstring"""

        if self.streaming:
              UpperCAmelCase__ :       List[Any]						=						self.builder.as_streaming_dataset(split=self.split  )
        # Build regular (map-style) dataset
        else:
              UpperCAmelCase__ :       Optional[Any]						=						None
              UpperCAmelCase__ :       str						=						None
              UpperCAmelCase__ :       Union[str, Any]						=						None
              UpperCAmelCase__ :       Optional[int]						=						None

              self.builder.download_and_prepare(
                  download_config=_lowerCamelCase       ,					download_mode=_lowerCamelCase       ,					verification_mode=_lowerCamelCase       ,					base_path=_lowerCamelCase       ,					num_proc=self.num_proc       ,					)
              UpperCAmelCase__ :       List[str]						=						self.builder.as_dataset(
                  split=self.split       ,					verification_mode=_lowerCamelCase       ,					in_memory=self.keep_in_memory  )
        return dataset




class 			lowerCamelCase  :


  '''simple docstring'''







  def __init__(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase = None       ,					_lowerCamelCase = None       ,					**_lowerCamelCase       ,					):


        """simple docstring"""

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

        UpperCAmelCase__ :       str						=						dataset
        UpperCAmelCase__ :       Optional[Any]						=						path_or_buf
        UpperCAmelCase__ :       str						=						batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        UpperCAmelCase__ :       Tuple						=						num_proc
        UpperCAmelCase__ :       int						=						"""utf-8"""
        UpperCAmelCase__ :       Optional[int]						=						to_csv_kwargs



  def    _a			(self  ):


        """simple docstring"""

        UpperCAmelCase__ :       str						=						self.to_csv_kwargs.pop("""path_or_buf"""       ,					_lowerCamelCase  )
        UpperCAmelCase__ :       Union[str, Any]						=						self.to_csv_kwargs.pop("""header"""       ,					_lowerCamelCase  )
        UpperCAmelCase__ :       List[str]						=						self.to_csv_kwargs.pop("""index"""       ,					_lowerCamelCase  )

        if isinstance(self.path_or_buf       ,					(str, bytes, os.PathLike)  ):
              with open(self.path_or_buf       ,					"""wb+"""  ) as buffer:
                    UpperCAmelCase__ :       List[str]						=						self._write(file_obj=_lowerCamelCase       ,					header=_lowerCamelCase       ,					index=_lowerCamelCase       ,					**self.to_csv_kwargs  )
        else:
              UpperCAmelCase__ :       Optional[Any]						=						self._write(file_obj=self.path_or_buf       ,					header=_lowerCamelCase       ,					index=_lowerCamelCase       ,					**self.to_csv_kwargs  )
        return written



  def    _a			(self       ,					_lowerCamelCase  ):


        """simple docstring"""

        UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__				,		UpperCAmelCase__ :       Any						=						args

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




  def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					**_lowerCamelCase  ):


        """simple docstring"""

        UpperCAmelCase__ :       Tuple						=						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"""       ,					):
                    UpperCAmelCase__ :       Dict						=						self._batch_csv((offset, header, index, to_csv_kwargs)  )
                    written += file_obj.write(_lowerCamelCase  )

        else:
              UpperCAmelCase__				,		UpperCAmelCase__ :       Optional[int]						=						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       ,					_lowerCamelCase       ,					_lowerCamelCase  )]       ,					)       ,					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(_lowerCamelCase  )

        return written




"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_keras_nlp_available,
    is_tensorflow_text_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_A						=     {
    """configuration_gpt2""": ["""GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPT2Config""", """GPT2OnnxConfig"""],
    """tokenization_gpt2""": ["""GPT2Tokenizer"""],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A						=     ["""GPT2TokenizerFast"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A						=     [
        """GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """GPT2DoubleHeadsModel""",
        """GPT2ForQuestionAnswering""",
        """GPT2ForSequenceClassification""",
        """GPT2ForTokenClassification""",
        """GPT2LMHeadModel""",
        """GPT2Model""",
        """GPT2PreTrainedModel""",
        """load_tf_weights_in_gpt2""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A						=     [
        """TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """TFGPT2DoubleHeadsModel""",
        """TFGPT2ForSequenceClassification""",
        """TFGPT2LMHeadModel""",
        """TFGPT2MainLayer""",
        """TFGPT2Model""",
        """TFGPT2PreTrainedModel""",
    ]

try:
    if not is_keras_nlp_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A						=     ["""TFGPT2Tokenizer"""]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _A						=     ["""FlaxGPT2LMHeadModel""", """FlaxGPT2Model""", """FlaxGPT2PreTrainedModel"""]

if TYPE_CHECKING:
    from .configuration_gpta import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTaConfig, GPTaOnnxConfig
    from .tokenization_gpta import GPTaTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_gpta_fast import GPTaTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_gpta import (
            GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
            GPTaDoubleHeadsModel,
            GPTaForQuestionAnswering,
            GPTaForSequenceClassification,
            GPTaForTokenClassification,
            GPTaLMHeadModel,
            GPTaModel,
            GPTaPreTrainedModel,
            load_tf_weights_in_gpta,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_gpta import (
            TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFGPTaDoubleHeadsModel,
            TFGPTaForSequenceClassification,
            TFGPTaLMHeadModel,
            TFGPTaMainLayer,
            TFGPTaModel,
            TFGPTaPreTrainedModel,
        )

    try:
        if not is_keras_nlp_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_gpta_tf import TFGPTaTokenizer

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel, FlaxGPTaPreTrainedModel

else:
    import sys

    _A						=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

