# How to split main dataset into train, dev, test as DatasetDict

**URL:** <https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090>\
**Category:** 🤗Datasets\
**Created:** [September 11, 2020, 7:30pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090 "2020-09-11T19:30:30Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [September 11, 2020, 7:30pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/1 "2020-09-11T19:30:30Z")

</div>

It seems that a single dataset can be split up into different partitions but in such a way that the connection between them is still clear (by using a DatasetDict), which is neat. I am having difficulties trying to figure out how I can create them, and use them, though. I’ve been going through the documentation [[1]](https://huggingface.co/docs/datasets/splits.html),[[2]](https://huggingface.co/docs/datasets/package_reference/main_classes.html?highlight=datasetdict#datasets.DatasetDict) and the source code [[1]](https://github.com/huggingface/datasets/blob/master/src/datasets/splits.py),[[2]](https://github.com/huggingface/datasets/blob/537de11963005d9ae1a7e560bdcde4f179056f15/src/datasets/dataset_dict.py) but it hasn’t become any clearer.

In some parts you speak of only a [train, test split](https://huggingface.co/docs/datasets/package_reference/main_classes.html?highlight=train_test#datasets.Dataset.train_test_split) other times you include [validation](https://huggingface.co/docs/datasets/package_reference/builder_classes.html#datasets.Split). It is not clear how to split an existing Dataset into a DatasetDict containing train, dev (validation), and test keys, even though I think this is a common scenario.

The focus of the library seems to be on the use of existing datasets (which is a great feature!), but time permitting I would like to see more information and use cases for creating and managing your own datasets. I’m usually pretty good at figuring out a code base by its documentation and source code (e.g. `transformers`), but `dataset` seems to give me a harder time. That is not to say that I think it is bad, of course. But it does show, I think, that even though it seems simple from the examples, it is actually quite complex behind the scenes. For people who want do custom stuff it is then super important to have clear, verbose documentation about as many aspects of the library as possible.

---

<div class="post-metadata">

**Author:** ![thomwolf](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/thomwolf/32/2198_2.png) [@thomwolf](https://discuss.huggingface.co/u/thomwolf)\
**Post date:** [September 11, 2020, 8:50pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/2 "2020-09-11T20:50:24Z")

</div>

Hi Bram,  
Yes the documentation of `train_test_split` that you link to is the right one. The `train_test_split` method currently provided is just a copy of the famous sklearn `train_test_split` (that we kinda assume people to be familiar with), we just removed the stratified split options which are quite complex.  
We could add an option to split in three with a validation split indeed, feel free to open a PR on this if you would like to have this feature fast.  
Right now what you can do is splitting two times:

```python
# 90% train, 10% test + validation
train_testvalid = dataset.train_test_split(test=0.1)
# Split the 10% test + valid in half test, half valid
test_valid = train_test_dataset['test'].train_test_split(test=0.5)
# gather everyone if you want to have a single DatasetDict
train_test_valid_dataset = DatasetDict({
    'train': train_testvalid['train'],
    'test': test_valid['test'],
    'valid': test_valid['train']})

```

The mention of a `validation` split that you point to is just an enum provided for dataset creators who would like to include a standard name for a validation split in there dataset.

We will try to add more doc on the code organization but bear in mind that (1) this library is still very young and (2) we are a lot less working on it (it’s really mostly one person, Quentin, that I try to help as much as I can), so it will definitely take some time before we have “have clear, verbose documentation about as many aspects of the library as possible”.

Basically, to give you an idea, the code is organized in two main parts:

1. the **dataset building part** which is defined in part by the people writing datasets and very open (hence the many options for splits in this part) =\> this is most of the complex code because it a wrapper around script provided externally. This includes files like `builder.py`, `load.py`, `arrow_dataset.py`.
2. the **dataset processing part** (after the dataset has been build) which is mostly contained in the `arrow_dataset.py` file and contains most of what the users will actually interact with =\> this is probably the part you need to read the most. The main complex part here is that we are deeply integrated with Apache Arrow which is very efficient but definitely not the easiest framework to understand.

You can also read this part in the doc where I tried to make a graph and give some information on how datasets are created (the first part in my list above): [https://huggingface.co/docs/datasets/add\_dataset.html](https://huggingface.co/docs/datasets/add_dataset.html)

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [September 12, 2020, 9:59am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/3 "2020-09-12T09:59:36Z")

</div>

I tried your approach, but I end up with empty batches here. Creating a dataloader for the whole dataset works:

```auto
dataloaders = {"train": DataLoader(dataset, batch_size=8)}

for batch in dataloaders["train"]:
    print(batch.keys())
    # prints the expected keys

```

But when I split the dataset as you suggest, I run into issues; the batches are empty.

```auto
# dataset is already `map`'d and already has `set_format`
# 90% train, 10% test + validation
train_testvalid = dataset.train_test_split(test_size=0.1)
# Split the 10% test + valid in half test, half valid
test_valid = train_testvalid["test"].train_test_split(test_size=0.5)
# gather everyone if you want to have a single DatasetDict
datasets = DatasetDict({
    "train": train_testvalid["train"],
    "test": test_valid["test"],
    "valid": test_valid["train"]})

dataloaders = {partition: DataLoader(ds, batch_size=8) for partition, ds in datasets.items()}

for batch in dataloaders["train"]:
    print(batch.keys())
    # dict_keys([])

```

Even when I just split into train and test, the batches are empty.

```python
train_test = dataset.train_test_split(test_size=0.1)
print(next(iter(train_test["train"]))) # empty

```

During those two last snippets, I see a lot of warnings/loggings. Using a dummy dataset of 100 entries, I see this:

```auto
PyTorch version 1.6.0+cu101 available.
Testing the mapped function outputs
Testing finished, running the mapping function on the dataset
100%|██████████| 100/100 [00:00<00:00, 4999.95ex/s]
Done writing 100 examples in 426851 bytes .
Testing the mapped function outputs
Testing finished, running the mapping function on the dataset
100%|██████████| 1/1 [00:00<00:00, 15.38ba/s]
Done writing 100 examples in 1656863 bytes .
Set __getitem__ (key) output type to torch for ['input_ids', 'sembedding'] columns (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 90 indices in 720 bytes .
Set __getitem__ (key) output type to torch for ['input_ids', 'sembedding'] columns (when key is int or slice) and don't output other (un-formatted) columns.
Done writing 10 indices in 80 bytes .
Set __getitem__ (key) output type to torch for ['input_ids', 'sembedding'] columns (when key is int or slice) and don't output other (un-formatted) columns.
Set __getitem__ (key) output type to torch for [] columns (when key is int or slice) and don't output other (un-formatted) columns.
Set __getitem__ (key) output type to torch for [] columns (when key is int or slice) and don't output other (un-formatted) columns.

```

It seems that for each split, `set_format` is called again, but after that it is called again with empty column names - meaning that no columns will be included in the end. I am not sure why it is calling it with empty columns there. I thought that this might be a bug with the fingerprinting, but after clearing the cache this still occurs.

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [September 12, 2020, 11:38am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/4 "2020-09-12T11:38:43Z")

</div>

The solution seems to be to _first_ doing the partitioning and then doing the `map` and `set_format` on the `DatasetDict`. That being said, I still feel like the behaviour that I mentioned before is a bug - but I am not sure.

---

<div class="post-metadata">

**Author:** ![thomwolf](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/thomwolf/32/2198_2.png) [@thomwolf](https://discuss.huggingface.co/u/thomwolf)\
**Post date:** [September 15, 2020, 6:56am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/5 "2020-09-15T06:56:32Z")

</div>

Thanks for the feedback @BramVanroy

There are a few things here indeed and I think the first is that we should rationalize more the logging of the library.

For instance these `Set __getitem__ (key) output type to torch` logs are too much information for the users, in particular since they are actually only used internally by the library in these calls.

Now the fact that you have to use the method in a specific order seems to be a bug, if you want to open an issue about we would be happy to investigate.

---

<div class="post-metadata">

**Author:** ![Roberto](https://avatars.discourse-cdn.com/v4/letter/r/8baadc/32.png) [@Roberto](https://discuss.huggingface.co/u/Roberto)\
**Post date:** [March 25, 2021, 12:35pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/6 "2021-03-25T12:35:11Z")

</div>

train\_testvalid = dataset.train\_test\_split(test=0.1)  
Split the 10% test + valid in half test, half valid  
test\_valid = **train\_testvalid** [‘test’].train\_test\_split(test=0.5)

I guess there is a typo here (in bold), works for me like that, thanks a lot !

---

<div class="post-metadata">

**Author:** ![cla](https://avatars.discourse-cdn.com/v4/letter/c/77aa72/32.png) [@cla](https://discuss.huggingface.co/u/cla)\
**Post date:** [April 1, 2021, 12:25pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/7 "2021-04-01T12:25:10Z")

</div>

Did you manage to create a DatasetDict with these partitions?

\*Solved. It was an import error.

Thanks!

---

<div class="post-metadata">

**Author:** ![nehanawalkar](https://avatars.discourse-cdn.com/v4/letter/n/7ab992/32.png) [@nehanawalkar](https://discuss.huggingface.co/u/nehanawalkar)\
**Post date:** [May 19, 2021, 5:18pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/8 "2021-05-19T17:18:02Z")

</div>

> [@thomwolf](#):
>
> ```auto
> train_testvalid = dataset.train_test_split(test=0.1)
> # Split the 10% test + valid in half test, half valid
> test_valid = train_test_dataset['test'].train_test_split(test=0.5)
> # gather everyone if you want to have a single DatasetDict
> train_test_valid_dataset = DatasetDict({
> 'train': train_testvalid['train'],
> 'test': test_valid['test'],
> 'valid': test_valid['train']})
> 
> ```

I am getting error like AttributeError: ‘DatasetDict’ object has no attribute ‘train\_test\_split’. Could you please tell me how to solve this?

---

<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:** [May 25, 2021, 9:17am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/9 "2021-05-25T09:17:06Z")

</div>

Hi ! This is because your `dataset` is probably a DatasetDict and not a Dataset object.

You might have to change the first line:

```auto
train_testvalid = dataset["train"].train_test_split(test=0.1)

```

---

<div class="post-metadata">

**Author:** ![baumstan](https://avatars.discourse-cdn.com/v4/letter/b/f05b48/32.png) [@baumstan](https://discuss.huggingface.co/u/baumstan)\
**Post date:** [September 22, 2021, 3:54am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/10 "2021-09-22T03:54:11Z")

</div>

Thanks for this, but I’m getting an error for the last piece of the code:

```auto
posts = load_dataset('peixian/rtGender', 'posts', split = 'train')
print(posts)
posts[0]

# save out a list of unique post_ids for the training and then dev sets
# split train 70%/ dev 15% / test 15%

# split twice and combine
train_devtest = posts.train_test_split(shuffle = True, seed = 200, test_size=0.3)
posts_dev_test = train_devtest['test'].train_test_split(shuffle = True, seed = 200, test_size=0.50)
posts_train_dev_test_dataset = DatasetDict({
    'train': train_devtest['train'],
    'test': posts_dev_test['test'],
    'dev': posts_dev_test['train']})

```

NameError: name ‘DatasetDict’ is not defined

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [September 22, 2021, 10:49am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/11 "2021-09-22T10:49:22Z")

</div>

Well, then you forgot to import it from `datasets`.

---

<div class="post-metadata">

**Author:** ![baumstan](https://avatars.discourse-cdn.com/v4/letter/b/f05b48/32.png) [@baumstan](https://discuss.huggingface.co/u/baumstan)\
**Post date:** [September 22, 2021, 4:56pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/12 "2021-09-22T16:56:32Z")

</div>

Hm. Could you specify which module to load? Currently this is what I have:

```auto
from datasets import load_dataset, dataset_dict

```

But the error persists.

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [September 22, 2021, 6:32pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/13 "2021-09-22T18:32:12Z")

</div>

You are missing

> [@baumstan](#):
>
> DatasetDict

---

<div class="post-metadata">

**Author:** ![aclifton314](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/aclifton314/32/55465_2.png) [@aclifton314](https://discuss.huggingface.co/u/aclifton314)\
**Post date:** [July 5, 2022, 9:08pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/14 "2022-07-05T21:08:03Z")

</div>

Just curious if this is still the recommended way to create a train, validation, and test set?

---

<div class="post-metadata">

**Author:** ![jorgenhw](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/jorgenhw/32/5005_2.png) [@jorgenhw](https://discuss.huggingface.co/u/jorgenhw)\
**Post date:** [October 18, 2022, 10:58am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/15 "2022-10-18T10:58:47Z")

</div>

The original `train_test_split` from `sklearn` also has the `random=state=True` which is really useful for when comparing different models. It doesn’t seem as this is included with huggingface?

---

<div class="post-metadata">

**Author:** ![jorgenhw](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/jorgenhw/32/5005_2.png) [@jorgenhw](https://discuss.huggingface.co/u/jorgenhw)\
**Post date:** [October 18, 2022, 11:33am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/16 "2022-10-18T11:33:12Z")

</div>

I’m sorry, I just realized that this function is just calls `seed = [int]`

---

<div class="post-metadata">

**Author:** ![BbKWlf](https://avatars.discourse-cdn.com/v4/letter/b/c67d28/32.png) [@BbKWlf](https://discuss.huggingface.co/u/BbKWlf)\
**Post date:** [December 25, 2022, 3:04am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/17 "2022-12-25T03:04:19Z")

</div>

I tried this, but now it shows a new problem.

 ![image](https://us1.discourse-cdn.com/hellohellohello/original/2X/6/66678c793fe8e9a27a9ad7a62d217ec8ec6aa23d.png)  
Can you please help me?

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [December 27, 2022, 5:37pm UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/18 "2022-12-27T17:37:04Z")

</div>

When encountering errors like these, just have a look at the documentation. [There](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Dataset.train_test_split) you’ll find that the correct argument is `test_size`.

---

<div class="post-metadata">

**Author:** ![BbKWlf](https://avatars.discourse-cdn.com/v4/letter/b/c67d28/32.png) [@BbKWlf](https://discuss.huggingface.co/u/BbKWlf)\
**Post date:** [December 31, 2022, 1:10am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/19 "2022-12-31T01:10:30Z")

</div>

Thanks! I already solved it with your method but forgot to delete the question.

---

<div class="post-metadata">

**Author:** ![Resi](https://avatars.discourse-cdn.com/v4/letter/r/c37758/32.png) [@Resi](https://discuss.huggingface.co/u/Resi)\
**Post date:** [March 27, 2024, 9:43am UTC](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090/21 "2024-03-27T09:43:50Z")

</div>

From now on it is written like this:

from datasets import DatasetDict  
#split into test, train, val

# 90% train, 10% test + validation

train\_testvalid = dataset.train\_test\_split(test\_size=0.1)

# Split the 10% test + valid in half test, half valid

test\_valid = train\_testvalid[‘test’].train\_test\_split(test\_size=0.5)

# gather everyone if you want to have a single DatasetDict

train\_test\_valid\_dataset = DatasetDict({  
‘train’: train\_testvalid[‘train’],  
‘test’: test\_valid[‘test’],  
‘valid’: test\_valid[‘train’]})

train\_test\_valid\_dataset

[Next page](https://discuss.huggingface.co/t/how-to-split-main-dataset-into-train-dev-test-as-datasetdict/1090.md?page=2)