# Splitting dataset from generator

**URL:** <https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336>\
**Category:** 🤗Datasets\
**Created:** [January 24, 2023, 6:32pm UTC](https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336 "2023-01-24T18:32:57Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![LuizBromfman](https://avatars.discourse-cdn.com/v4/letter/l/3e96dc/32.png) [@LuizBromfman](https://discuss.huggingface.co/u/LuizBromfman)\
**Post date:** [January 24, 2023, 6:32pm UTC](https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336/1 "2023-01-24T18:32:57Z")

</div>

Hello all!

I am making a dataset in a python generator which gets data from AWS’ s3 buckets. I am using boto3 in order to get the images and data I need. The yielded dict looks something like this:

```auto
{
            'bboxes': [],
            'image': Image.open(get_file(s3, 'bucket_name', img_file)).resize((1000, 1000)),
            'ner_tags': [],
            'tokens': [],
            'ids': []
}

```

Since my AWS bucket has around 14.000 documents, I thought a generator would be my best bet in order to load this data using HF datasets.

I am already aware of the `from_generator()` method from the Dataset class and I want to load the dataset as such:

```auto
dataset = Dataset.from_generator(s3_file_generator)

```

However, I also want to split the dataset into train and test datasets (80:20 ratio), and, unlike the `load_dataset()` function, the `from_generator()` method does not accept the `split="train"` parameter. I have tried this already:

```auto
dataset = Dataset.from_generator(s3_file_generator).train_test_split(test_size=0.2)

```

But this didn’t work.

I need a way to load this dataset from my python generator while still being able to split it into train and test. Is there any solution I’m missing?

I really appreciate any help you can provide. Thanks in advance!

---

<div class="post-metadata">

**Author:** ![LuizBromfman](https://avatars.discourse-cdn.com/v4/letter/l/3e96dc/32.png) [@LuizBromfman](https://discuss.huggingface.co/u/LuizBromfman)\
**Post date:** [January 24, 2023, 9:14pm UTC](https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336/2 "2023-01-24T21:14:11Z")

</div>

Hello!  
Just an update:  
After 170 minutes the following code successfully worked:

```auto
dataset = Dataset.from_generator(s3_file_generator).train_test_split(test_size=0.2)

```

However, due to the long run time, I suspect it downloaded every 14k dicts, instead of showing one at a time like the generator is supposed to…

So I still need to figure out some way of optimizing this code or if I am making a mistake when loading the dataset from my generator.

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [January 25, 2023, 12:59pm UTC](https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336/3 "2023-01-25T12:59:04Z")

</div>

Hi ! A `Dataset` stores the data in Arrow format, so it downloads everything. If you want your dataset to download progressively as your iterate over the dataset, you can use `IterableDataset.from_generator` instead.

`IterableDataset` doesn’t implement train\_test\_split though, because you can’t split a generator in the middle. Instead, I’d recommend defining two generators: one for train and one for test

---

<div class="post-metadata">

**Author:** ![LuizBromfman](https://avatars.discourse-cdn.com/v4/letter/l/3e96dc/32.png) [@LuizBromfman](https://discuss.huggingface.co/u/LuizBromfman)\
**Post date:** [January 26, 2023, 12:36pm UTC](https://discuss.huggingface.co/t/splitting-dataset-from-generator/30336/4 "2023-01-26T12:36:27Z")

</div>

Hello! Thank you for your response! This worked but in the end it would return an error every epoch from the training since it would try to iterate `next()` from the dataset, which wouldn’t work. It ended up being more efficient to download the dataset as a `Dataset` object.

Thank you very much for your time!