




'''simple docstring'''



import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
     import tensorflow as tf





class     lowerCAmelCase   ( TensorFormatter[Mapping, '''tf.Tensor''', Mapping]):


       def __init__(   self    ,  __SCREAMING_SNAKE_CASE=None    ,  **__SCREAMING_SNAKE_CASE  )      ->  List[str]:



            '''simple docstring'''




            super().__init__(features=__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  tf_tensor_kwargs
            import tensorflow as tf  # noqa: import tf at initialization


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




            import tensorflow as tf

            if isinstance(__SCREAMING_SNAKE_CASE    ,  __SCREAMING_SNAKE_CASE  ) and column:
                 if all(
                     isinstance(__SCREAMING_SNAKE_CASE    ,  tf.Tensor  ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column  ):
                      return tf.stack(__SCREAMING_SNAKE_CASE  )
                 elif all(
                     isinstance(__SCREAMING_SNAKE_CASE    ,  (tf.Tensor, tf.RaggedTensor)  ) and x.ndim == 1 and x.dtype == column[0].dtype
                     for x in column  ):
                      # only rag 1-D tensors, otherwise some dimensions become ragged even though they were consolidated
                      return tf.ragged.stack(__SCREAMING_SNAKE_CASE  )

            return column


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  int:



            '''simple docstring'''




            import tensorflow as tf

            if value is None:
                 return value

            __snake_case    					=  {}

            if isinstance(__SCREAMING_SNAKE_CASE    ,  (np.number, np.ndarray)  ) and np.issubdtype(value.dtype    ,  np.integer  ):
                 __snake_case    					=  {'''dtype''': tf.intaa}
            elif isinstance(__SCREAMING_SNAKE_CASE    ,  (np.number, np.ndarray)  ) and np.issubdtype(value.dtype    ,  np.floating  ):
                 __snake_case    					=  {'''dtype''': tf.floataa}
            elif config.PIL_AVAILABLE and "PIL" in sys.modules:
                 import PIL.Image

                 if isinstance(__SCREAMING_SNAKE_CASE    ,  PIL.Image.Image  ):
                      __snake_case    					=  np.asarray(__SCREAMING_SNAKE_CASE  )

            return tf.convert_to_tensor(__SCREAMING_SNAKE_CASE    ,  **{**default_dtype, **self.tf_tensor_kwargs}  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Tuple:



            '''simple docstring'''




            import tensorflow as tf

            # support for torch, tf, jax etc.
            if config.TORCH_AVAILABLE and "torch" in sys.modules:
                 import torch

                 if isinstance(__SCREAMING_SNAKE_CASE    ,  torch.Tensor  ):
                      return self._tensorize(data_struct.detach().cpu().numpy()[()]  )
            if hasattr(__SCREAMING_SNAKE_CASE    ,  '''__array__'''  ) and not isinstance(__SCREAMING_SNAKE_CASE    ,  tf.Tensor  ):
                 __snake_case    					=  data_struct.__array__()
            # support for nested types like struct of list of struct
            if isinstance(__SCREAMING_SNAKE_CASE    ,  np.ndarray  ):
                 if data_struct.dtype == object:  # tf tensors cannot be instantied from an array of objects
                      return self._consolidate([self.recursive_tensorize(__SCREAMING_SNAKE_CASE  ) for substruct in data_struct]  )
            elif isinstance(__SCREAMING_SNAKE_CASE    ,  (list, tuple)  ):
                 return self._consolidate([self.recursive_tensorize(__SCREAMING_SNAKE_CASE  ) for substruct in data_struct]  )
            return self._tensorize(__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[int]:



            '''simple docstring'''




            return map_nested(self._recursive_tensorize    ,  __SCREAMING_SNAKE_CASE    ,  map_list=__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Mapping:



            '''simple docstring'''




            __snake_case    					=  self.numpy_arrow_extractor().extract_row(__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.python_features_decoder.decode_row(__SCREAMING_SNAKE_CASE  )
            return self.recursive_tensorize(__SCREAMING_SNAKE_CASE  )


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  "tf.Tensor":



            '''simple docstring'''




            __snake_case    					=  self.numpy_arrow_extractor().extract_column(__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.python_features_decoder.decode_column(__SCREAMING_SNAKE_CASE    ,  pa_table.column_names[0]  )
            __snake_case    					=  self.recursive_tensorize(__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self._consolidate(__SCREAMING_SNAKE_CASE  )
            return column




       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Mapping:



            '''simple docstring'''




            __snake_case    					=  self.numpy_arrow_extractor().extract_batch(__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.python_features_decoder.decode_batch(__SCREAMING_SNAKE_CASE  )
            __snake_case    					=  self.recursive_tensorize(__SCREAMING_SNAKE_CASE  )
            for column_name in batch:
                 __snake_case    					=  self._consolidate(batch[column_name]  )
            return batch









'''simple docstring'''



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_       :  str      =		datasets.utils.logging.get_logger(__name__)





@dataclass
class     lowerCAmelCase   ( datasets.BuilderConfig):
       __lowercase  :      Optional[datasets.Features]        =			None
       __lowercase  :      str        =			"utf-8"
       __lowercase  :      Optional[str]        =			None
       __lowercase  :      int        =			10 << 20  # 10MB
       __lowercase  :      bool        =			False
       __lowercase  :      str        =			"line"





class     lowerCAmelCase   ( datasets.ArrowBasedBuilder):
       __lowercase  :      Dict        =			TextConfig


       def       lowerCAmelCase      (   self  )      ->  Dict:



            '''simple docstring'''




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


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


       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  pa.Table:



            '''simple docstring'''




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




       def       lowerCAmelCase      (   self    ,  __SCREAMING_SNAKE_CASE  )      ->  Optional[Any]:



            '''simple docstring'''




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






