# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/15_training.arrow_utils.ipynb (unless otherwise specified).

__all__ = ['TextNoNewLine', 'TextNoNewLineDatasetReader']

# Cell
from typing import Optional

from datasets import Features, NamedSplit
from datasets.packaged_modules.text.text import Text
from datasets.utils.typing import NestedDataStructureLike, PathLike
from datasets.io.abc import AbstractDatasetReader
import datasets
from datasets.packaged_modules.text.text import TextConfig

import pyarrow as pa

# Cell
class TextNoNewLine(datasets.ArrowBasedBuilder):
    "An ArrowBasedBuilder dataset that doesn't split based on newlines"
    BUILDER_CONFIG_CLASS = TextConfig

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

    def _split_generators(self, dl_manager):
        """The `data_files` kwarg in load_dataset() can be a str, List[str], Dict[str,str], or Dict[str,List[str]].
        If str or List[str], then the dataset returns only the 'train' split.
        If dict, then keys should be from the `datasets.Split` enum.
        """
        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}")
        data_files = dl_manager.download_and_extract(self.config.data_files)
        if isinstance(data_files, (str, list, tuple)):
            files = data_files
            if isinstance(files, str):
                files = [files]
            return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"files": files})]
        splits = []
        for split_name, files in data_files.items():
            if isinstance(files, str):
                files = [files]
            splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
        return splits

    def _generate_tables(self, files):
        schema = pa.schema(self.config.features.type if self.config.features is not None else {"text": pa.string()})
        for file_idx, file in enumerate(files):
            batch_idx = 0
            with open(file, "r", encoding=self.config.encoding) as f:
                while True:
                    batch = f.read(self.config.chunksize)
                    if not batch:
                        break
                    batch = [batch]
                    pa_table = pa.Table.from_arrays([pa.array(batch)], schema=schema)
                    # 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), pa_table
                    batch_idx += 1

# Cell
class TextNoNewLineDatasetReader(AbstractDatasetReader):
    "A `DatasetReader` class that mimics `TextDatasetReader`, but uses `TextNoNewLine`"
    def __init__(
        self,
        path_or_paths: NestedDataStructureLike[PathLike], # Paths to texts
        split: Optional[NamedSplit] = None, # Potential dataset splits
        features: Optional[Features] = None, # Any features for the datasets
        cache_dir: str = None, # A cache directory
        keep_in_memory: bool = False, # Whether to keep the dataset completely in memory
        **kwargs,
    ):
        super().__init__(
            path_or_paths, split=split, features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs
        )
        path_or_paths = path_or_paths if isinstance(path_or_paths, dict) else {self.split: path_or_paths}
        self.builder = TextNoNewLine(
            cache_dir=cache_dir,
            data_files=path_or_paths,
            features=features,
            **kwargs,
        )

    def read(self):
        "Reads in a dataset"
        download_config = None # A download configuration
        download_mode = None # A download mode
        ignore_verifications = False # Whether to ignore any verification requests
        use_auth_token = None # Potential authorization tokens
        base_path = None # A path

        self.builder.download_and_prepare(
            download_config=download_config,
            download_mode=download_mode,
            ignore_verifications=ignore_verifications,
            # try_from_hf_gcs=try_from_hf_gcs,
            base_path=base_path,
            use_auth_token=use_auth_token,
        )

        # Build dataset for splits
        dataset = self.builder.as_dataset(
            split=self.split, ignore_verifications=ignore_verifications, in_memory=self.keep_in_memory
        )
        return dataset