huggingface / huggingface/datasets
Download + preparation speed of datasets.load_dataset is 20x slower than huggingface hub snapshot and manual loding
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
I am working with a dataset I am trying to publish.
The path is Antreas/TALI.
It's a fairly large dataset, and contains images, video, audio and text.
I have been having multiple problems when the dataset is being downloaded using the load_dataset function -- even with 64 workers taking more than 7 days to process.
With snapshot download it takes 12 hours, and that includes the dataset preparation done using load_dataset and passing the dataset parquet file paths.
Find the script I am using below:
```python
import multiprocessing as mp
import pathlib
from typing import Optional
import datasets
from rich import print
from tqdm import tqdm
def download_dataset_via_hub(
dataset_name: str,
dataset_download_path: pathlib.Path,
num_download_workers: int = mp.cpu_count(),
):
import huggingface_hub as hf_hub
download_folder = hf_hub.snapshot_download(
repo_id=dataset_name,
repo_type="dataset",
cache_dir=dataset_download_path,
resume_download=True,
max_workers=num_download_workers,
ignore_patterns=[],
)
return pathlib.Path(download_folder) / "data"
def load_dataset_via_hub(
dataset_download_path: pathlib.Path,
num_download_workers: int = mp.cpu_count(),
dataset_name: Optional[str] = None,
):
from dataclasses import dataclass, field
from datasets import ClassLabel, Features, Image, Sequence, Value
dataset_path = download_dataset_via_hub(
dataset_download_path=dataset_download_path,
num_download_workers=num_download_workers,
dataset_name=dataset_name,
)
# Building a list of file paths for validation set
train_files = [
file.as_posix()
for file in pathlib.Path(dataset_path).glob("*.parquet")
if "train" in file.as_posix()
]
val_files = [
file.as_posix()
for file in pathlib.Path(dataset_path).glob("*.parquet")
if "val" in file.as_posix()
]
test_files = [
file.as_posix()
for file in pathlib.Path(dataset_path).glob("*.parquet")
if "test" in file.as_posix()
]
print(
f"Found {len(test_files)} files for testing set, {len(train_files)} for training set and {len(val_files)} for validation set"
)
data_files = {
"test": test_files,
"val": val_files,
"train": train_files,
}
features = Features(
{
"image": Image(
decode=True
), # Set `decode=True` if you want to decode the images, otherwise `decode=False`
"image_url": Value("string"),
"item_idx": Value("int64"),
"wit_features": Sequence(
{
"attribution_passes_lang_id": Value("bool"),
"caption_alt_text_description": Value("string"),
"caption_reference_description": Value("string"),
"caption_title_and_reference_description": Value("string"),
"context_page_description": Value("string"),
"context_section_description": Value("string"),
"hierarchical_section_title": Value("string"),
"is_main_image": Value("bool"),
"language": Value("string"),
"page_changed_recently": Value("bool"),
"page_title": Value("string"),
"page_url": Value("string"),
"section_title": Value("string"),
}
),
"wit_idx": Value("int64"),
"youtube_title_text": Value("string"),
"youtube_description_text": Value("string"),
"youtube_video_content": Value("binary"),
"youtube_video_starting_time": Value("string"),
"youtube_subtitle_text": Value("string"),
"youtube_video_size": Value("int64"),
"youtube_video_file_path": Value("string"),
}
)
dataset = datasets.load_dataset(
"parquet" if dataset_name is None else dataset_name,
data_files=data_files,
features=features,
num_proc=1,
cache_dir=dataset_download_path / "cache",
)
return dataset
if __name__ == "__main__":
dataset_cache = pathlib.Path("/disk/scratch_fast0/tali/")
dataset = load_dataset_via_hub(dataset_cache, dataset_name="Antreas/TALI")[
"test"
]
for sample in tqdm(dataset):
print(list(sample.keys()))
```
Also, streaming this dataset has been a very painfully slow process. Streaming the train set takes 15m to start, and streaming the test and val sets takes 3 hours to start!
### Steps to reproduce the bug
1. Run the code I provided to get a sense of how fast snapshot + manual is
2. Run datasets.load_dataset("Antreas/TALI") to get a sense of the speed of that OP.
3. You should now have an appreciation of how long these things take.
### Expected behavior
The load dataset function should be at least as fast as the huggingface snapshot download function in terms of downloading dataset files. Not 20 times slower.
### Environment info
- `datasets` version: 2.14.5
- Platform: Linux-5.15.0-76-generic-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- PyArrow version: 13.0.0
- Pandas version: 2.1.1
Contributor guide
Assessment
This issue has not been assessed yet.