chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource/huggingface_datasource.py

@@ -0,0 +1,194 @@
+import sys
+from typing import TYPE_CHECKING, Iterable, List, Optional, Union
+
+from ray.data._internal.tensor_extensions.arrow import pyarrow_table_from_pydict
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.block import Block, BlockAccessor, BlockMetadata
+from ray.data.dataset import Dataset
+from ray.data.datasource import Datasource, ReadTask
+
+if TYPE_CHECKING:
+    import datasets
+
+    from ray.data.context import DataContext
+
+
+TRANSFORMERS_IMPORT_ERROR: Optional[ImportError] = None
+
+try:
+    # Due to HF Dataset's dynamic module system, we need to dynamically import the
+    # datasets_modules module on every actor when training.
+    # We accomplish this by simply running the following bit of code directly
+    # in the module you are currently viewing. This ensures that when we
+    # unpickle the Dataset, it runs before pickle tries to
+    # import datasets_modules and prevents an exception from being thrown.
+    # Same logic is present inside HF Transformers Ray
+    # integration: https://github.com/huggingface/transformers/blob/\
+    # 7d5fde991d598370d961be8cb7add6541e2b59ce/src/transformers/integrations.py#L271
+    # Also see https://github.com/ray-project/ray/issues/28084
+    from transformers.utils import is_datasets_available
+
+    if "datasets_modules" not in sys.modules and is_datasets_available():
+        import importlib
+        import importlib.metadata
+        import os
+
+        import datasets.load
+        from packaging.version import parse
+
+        # Datasets >= 4.0 removed dataset scripts support and the dynamic-modules cache.
+        # Only initialize dynamic modules on <= 3.x where the initializer `init_dynamic_modules` exists.
+        DATASETS_VERSION = parse(importlib.metadata.version("datasets"))
+        DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT = parse("4.0.0")
+
+        if DATASETS_VERSION < DATASETS_VERSION_WITHOUT_SCRIPT_SUPPORT:
+            dynamic_modules_path = os.path.join(
+                datasets.load.init_dynamic_modules(), "__init__.py"
+            )
+            # load dynamic_modules from path
+            spec = importlib.util.spec_from_file_location(
+                "datasets_modules", dynamic_modules_path
+            )
+            datasets_modules = importlib.util.module_from_spec(spec)
+            sys.modules[spec.name] = datasets_modules
+            spec.loader.exec_module(datasets_modules)
+except ImportError as e:
+    TRANSFORMERS_IMPORT_ERROR = e
+
+
+class HuggingFaceDatasource(Datasource):
+    """Hugging Face Dataset datasource, for reading from a
+    `Hugging Face Datasets Dataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset/>`_.
+    This Datasource implements a streamed read using a
+    single read task, most beneficial for a
+    `Hugging Face Datasets IterableDataset <https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset/>`_
+    or datasets which are too large to fit in-memory.
+    For an in-memory Hugging Face Dataset (`datasets.Dataset`), use :meth:`~ray.data.from_huggingface`
+    directly for faster performance.
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        dataset: Union["datasets.Dataset", "datasets.IterableDataset"],
+        batch_size: int = 4096,
+    ):
+        if TRANSFORMERS_IMPORT_ERROR is not None:
+            raise TRANSFORMERS_IMPORT_ERROR
+
+        self._dataset = dataset
+        self._batch_size = batch_size
+
+    @classmethod
+    def list_parquet_urls_from_dataset(
+        cls, dataset: Union["datasets.Dataset", "datasets.IterableDataset"]
+    ) -> Dataset:
+        """Return list of Hugging Face hosted parquet file URLs if they
+        exist for the data (i.e. if the dataset is a public dataset that
+        has not been transformed) else return an empty list."""
+        import datasets
+
+        # We can use the dataset name, config name, and split name to load
+        # public hugging face datasets from the Hugging Face Hub. More info
+        # here: https://huggingface.co/docs/datasets-server/parquet
+        dataset_name = dataset.info.dataset_name
+        config_name = dataset.info.config_name
+        split_name = str(dataset.split)
+
+        # If a dataset is not an iterable dataset, we will check if the
+        # dataset with the matching dataset name, config name, and split name
+        # on the Hugging Face Hub has the same fingerprint as the
+        # dataset passed into this function. If it is not matching, transforms
+        # or other operations have been performed so we cannot use the parquet
+        # files on the Hugging Face Hub, so we return an empty list.
+        if not isinstance(dataset, datasets.IterableDataset):
+            from datasets import load_dataset
+
+            try:
+                ds = load_dataset(dataset_name, config_name, split=split_name)
+                if ds._fingerprint != dataset._fingerprint:
+                    return []
+            except Exception:
+                # If an exception is thrown when trying to reload the dataset
+                # we should exit gracefully by returning an empty list.
+                return []
+
+        import requests
+
+        public_url = (
+            f"https://huggingface.co/api/datasets/{dataset_name}"
+            f"/parquet/{config_name}/{split_name}"
+        )
+        resp = requests.get(public_url)
+        if resp.status_code == requests.codes["ok"]:
+            # dataset corresponds to a public dataset, return list of parquet_files
+            return resp.json()
+        else:
+            return []
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return self._dataset.dataset_size
+
+    def _read_dataset(self) -> Iterable[Block]:
+        # Note: This is a method instead of a higher level function because
+        # we need to capture `self`. This will trigger the try-import logic at
+        # the top of file to avoid import error of dataset_modules.
+        import numpy as np
+        import pandas as pd
+        import pyarrow
+
+        for batch in self._dataset.with_format("arrow").iter(
+            batch_size=self._batch_size
+        ):
+            # HuggingFace IterableDatasets do not fully support methods like
+            # `set_format`, `with_format`, and `formatted_as`, so the dataset
+            # can return whatever is the default configured batch type, even if
+            # the format is manually overridden before iterating above.
+            # Therefore, we limit support to batch formats which have native
+            # block types in Ray Data (pyarrow.Table, pd.DataFrame),
+            # or can easily be converted to such (dict, np.array).
+            # See: https://github.com/huggingface/datasets/issues/3444
+            if not isinstance(batch, (pyarrow.Table, pd.DataFrame, dict, np.array)):
+                raise ValueError(
+                    f"Batch format {type(batch)} isn't supported. Only the "
+                    f"following batch formats are supported: "
+                    f"dict (corresponds to `None` in `dataset.with_format()`), "
+                    f"pyarrow.Table, np.array, pd.DataFrame."
+                )
+            # Ensure np.arrays are wrapped in a dict
+            # (subsequently converted to a pyarrow.Table).
+            if isinstance(batch, np.ndarray):
+                batch = {"item": batch}
+            if isinstance(batch, dict):
+                batch = pyarrow_table_from_pydict(batch)
+            # Ensure that we return the default block type.
+            block = BlockAccessor.for_block(batch).to_default()
+            yield block
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        # Note: `parallelism` arg is currently not used by HuggingFaceDatasource.
+        # We always generate a single ReadTask to perform the read.
+        _check_pyarrow_version()
+
+        # TODO(scottjlee): IterableDataset doesn't provide APIs
+        # for getting number of rows, byte size, etc., so the
+        # BlockMetadata is currently empty. Properly retrieve
+        # or calculate these so that progress bars have meaning.
+        meta = BlockMetadata(
+            num_rows=None,
+            size_bytes=None,
+            input_files=None,
+            exec_stats=None,
+        )
+        read_tasks: List[ReadTask] = [
+            ReadTask(
+                self._read_dataset,
+                meta,
+                per_task_row_limit=per_task_row_limit,
+            )
+        ]
+        return read_tasks