Loading partial dataset when debugging

[reachtarunhere]

I am using PyTorch Lightning along with datasets (thanks for so many datasets already prepared and the great splits).

Every time I execute load_dataset for the imdb dataset it takes some time even if I specify a split involving very few samples. I guess this due to hashing as per the other issues.

Is there a way to only load part of the dataset on load_dataset? This would really speed up my workflow.
Something like a debug mode would really help. Thanks!

[lhoestq]

Hi ! load_dataset downloads the full dataset once and caches it, so that subsequent calls to load_dataset just reloads the dataset from your disk.
Then when you specify a split in load_dataset, it will just load the requested split from the disk. If your specified split is a sliced split (e.g. "train[:10]"), then it will load the 10 first rows of the train split that you have on disk.

Therefore, as long as you don't delete your cache, all your calls to load_dataset will be very fast. Except the first call that downloads the dataset of course ^^

[thomwolf]

That’s a use case for the new streaming feature, no?

[albertvillanova]

Hi @reachtarunhere.

Besides the above insights provided by @lhoestq and @thomwolf, there is also a Dataset feature in progress (I plan to finish it this week): #2249, which will allow you, when calling load_dataset, to pass the option to download/preprocess/cache only some specific split(s), which will definitely speed up your workflow.

If this feature is interesting for you, I can ping you once it will be merged into the master branch.

[reachtarunhere]

Thanks all for responding.

Hey @albertvillanova

Thanks. Yes, I would be interested.

@lhoestq I think even if a small split is specified it loads up the full dataset from the disk (please correct me if this is not the case). Because it does seem to be slow to me even on subsequent calls. There is no repeated downloading so it seems that the cache is working.

I am not aware of the streaming feature @thomwolf mentioned. So I might need to read up on it.

[PierreColombo]

@reshinthadithyan I use the .select function to have a fraction of indices.

[ndvbd]

If I want to create a dataset, containing only the 10 elements of a given dataset (slice it), how do I do that?

[lhoestq]

small_ds = ds.select(range(10))

[ndvbd]

small_ds = ds.select(range(10))

Thanks, but this doesn't help me to save time during initial loading, right?

[lhoestq]

Indeed by default load_dataset would download and prepare everything as Arrow files. And passing split=train[:10] memory maps only the beginning of the full dataset that has been prepared on disk.

If you don't want to download everything, you can use streaming :

ids = load_dataset(..., streaming=True)
first_samples = list(ids["train"].take(10))

To get a Dataset you can use

ds = Dataset.from_generator(ids.take(10).__iter__)

edit: fixed small bug

[ndvbd]

Thanks @lhoestq, but I don't think it is 100% accurate, as it doesn't keep the dataset structure exactly the same.
To load the full dataset, I do:

data = load_dataset("json", data_files="a.json")
train_data = data["train"].shuffle()

But when I am changing it as per your instructions:

ids = load_dataset("json", data_files="a.json", streaming=True)
data = Dataset.from_generator(ids["train"].take(1).__iter__)
train_data = data["train"].shuffle()

It throws KeyError.
I need a simple way, like you suggested, to have a subset of a Dataset, which exactly the same attributes.

[lhoestq]

Whoops I fixed my code sorry

- ds = Dataset.from_generator(ids["train"].take(10).__iter__)
+ ds = Dataset.from_generator(ids.take(10).__iter__)

in your case that means running

train_data = data.shuffle()

without ["train"]