"""simple docstring"""



import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


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







@dataclass
class _lowerCAmelCase       (   datasets.BuilderConfig			):
   __UpperCAmelCase				: Union[str, Any]													=						10000
   __UpperCAmelCase				: Tuple													=						None
   __UpperCAmelCase				: str													=						None







class _lowerCAmelCase       (   datasets.ArrowBasedBuilder			):
   __UpperCAmelCase				: Optional[int]													=						ParquetConfig







   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




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







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	int:

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







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	pa.Table:

       '''simple docstring'''




       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
           snake_case					:   Dict         =					table_cast(_lowercase		,			self.info.features.arrow_schema		)
       return pa_table







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		)			->	str:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					self.info.features.arrow_schema if self.info.features is not None else None
       if self.info.features is not None and self.config.columns is not None:
           if sorted(field.name for field in schema		) != sorted(self.config.columns		):
               raise ValueError(
                   F'Tried to load parquet data with columns \'{self.config.columns}\' with mismatching features \'{self.info.features}\''		)
       for file_idx, file in enumerate(itertools.chain.from_iterable(_lowercase		)		):
           with open(_lowercase		,			"rb"		) as f:
               snake_case					:   int         =					pq.ParquetFile(_lowercase		)
               try:
                   for batch_idx, record_batch in enumerate(
                       parquet_file.iter_batches(batch_size=self.config.batch_size		,			columns=self.config.columns		)		):
                       snake_case					:   Tuple         =					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(_lowercase		)
               except ValueError as e:
                   logger.error(F'Failed to read file \'{file}\' with error {type(_lowercase		)}: {e}'		)
                   raise









"""simple docstring"""



from typing import TYPE_CHECKING

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


__snake_case         =					{
    """configuration_distilbert""": [
        """DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """DistilBertConfig""",
        """DistilBertOnnxConfig""",
    ],
    """tokenization_distilbert""": ["""DistilBertTokenizer"""],
}

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

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __snake_case         =					[
      """DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DistilBertForMaskedLM""",
      """DistilBertForMultipleChoice""",
      """DistilBertForQuestionAnswering""",
      """DistilBertForSequenceClassification""",
      """DistilBertForTokenClassification""",
      """DistilBertModel""",
      """DistilBertPreTrainedModel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __snake_case         =					[
      """TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFDistilBertForMaskedLM""",
      """TFDistilBertForMultipleChoice""",
      """TFDistilBertForQuestionAnswering""",
      """TFDistilBertForSequenceClassification""",
      """TFDistilBertForTokenClassification""",
      """TFDistilBertMainLayer""",
      """TFDistilBertModel""",
      """TFDistilBertPreTrainedModel""",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __snake_case         =					[
      """FlaxDistilBertForMaskedLM""",
      """FlaxDistilBertForMultipleChoice""",
      """FlaxDistilBertForQuestionAnswering""",
      """FlaxDistilBertForSequenceClassification""",
      """FlaxDistilBertForTokenClassification""",
      """FlaxDistilBertModel""",
      """FlaxDistilBertPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_distilbert import (
      DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      DistilBertConfig,
      DistilBertOnnxConfig,
  )
  from .tokenization_distilbert import DistilBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_distilbert_fast import DistilBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_distilbert import (
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_distilbert import (
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_distilbert import (
        FlaxDistilBertForMaskedLM,
        FlaxDistilBertForMultipleChoice,
        FlaxDistilBertForQuestionAnswering,
        FlaxDistilBertForSequenceClassification,
        FlaxDistilBertForTokenClassification,
        FlaxDistilBertModel,
        FlaxDistilBertPreTrainedModel,
    )

else:
  import sys

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





