wehub-resource-sync
195 lines
8.4 KiB
Python
195 lines
8.4 KiB
Python
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
|