# How to split Hugging Face dataset to train and test?
**URL:**
**Author:**  [@laro1](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Flaro1)%5C
**Post date:** [July 26, 2022, 12:24pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F1 "2022-07-26T12:24:29Z")
I have json file with data which I want to load and split to train and test (70% data for train).
I’m loading the records in this way:
```auto
full_path = "/home/ad/ds/fiction"
data_files = {
"DATA": os.path.join(full_path, "dev.json")
}
ds = load_dataset("json", data_files=data_files)
ds
DatasetDict({
DATA: Dataset({
features: ['premise', 'hypothesis', 'label'],
num_rows: 750
})
})
```
- How can I split this `ds` to train and test ?
- Can I change the `DATA` label to `TRAIN` and `TEST` with the relevant samples ?
---
**Author:**  [@stevhliu](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fstevhliu)%5C
**Post date:** [July 26, 2022, 4:13pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F2 "2022-07-26T16:13:09Z")
Hello and welcome @laro1!
You can use the `train_test_split()` function and specify the `test_size` parameter to determine the size of the split. For example:
```auto
ds.train_test_split(test_size=0.3)
DatasetDict({
train: Dataset({
features: ['premise', 'hypothesis', 'label'],
num_rows: 525
})
test: Dataset({
features: ['premise', 'hypothesis', 'label'],
num_rows: 225
})
})
```
Check out the docs [here](https://huggingface.co/docs/datasets/process#split) and let me know if that helps! 🤗
---
**Author:**  [@afriedman412](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fafriedman412)%5C
**Post date:** [August 13, 2022, 11:58pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F3 "2022-08-13T23:58:49Z")
is there anything like the “stratify” param in scikit-learn?
(or more generally a way to assure class balancing in train and test splits?)
---
**Author:**  [@lhoestq](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Flhoestq)%5C
**Post date:** [August 19, 2022, 10:26am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F4 "2022-08-19T10:26:26Z")
Yup, please check the `stratify_by_column` argument in [the docs](https://huggingface.co/docs/datasets/v2.4.0/en/package_reference/main_classes#datasets.Dataset.train_test_split.stratify_by_column)
```python
>>> ds = load_dataset("imdb",split="train")
Dataset({
features: ['text', 'label'],
num_rows: 25000
})
>>> ds = ds.train_test_split(test_size=0.2, stratify_by_column="label")
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 20000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 5000
})
})
```
---
**Author:**  [@YanaS](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FYanaS)%5C
**Post date:** [December 22, 2022, 7:39am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F5 "2022-12-22T07:39:03Z")
When I load my custom dataset from dictionary I get an error:
> ValueError: Stratifying by column is only supported for ClassLabel column, and column label is Sequence.
> with open(‘/content/drive/MyDrive/all.bio.pickle’, ‘rb’) as f:
> bio\_dict = pickle.load(f)
>
> ds = datasets.Dataset.from\_dict(bio\_dict)
Dataset({
features: [‘id’, ‘text’, ‘ner\_tags’, ‘input\_ids’, ‘attention\_mask’, ‘label’],
num\_rows: 8805
})
> train\_testvalid = ds.train\_test\_split(test\_size=0.5, shuffle=True, stratify\_by\_column=“label”)
>
> test\_valid = train\_testvalid[‘test’].train\_test\_split(test\_size=0.5, shuffle=True, stratify\_by\_column=“label”)
>
> ttv\_ds = datasets.DatasetDict({
> ‘train’: train\_testvalid[‘train’],
> ‘validation’: test\_valid[‘train’],
> ‘test’: test\_valid[‘test’]})
---
**Author:**  [@mkdeeperinsights](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmkdeeperinsights)%5C
**Post date:** [January 24, 2023, 4:56pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-split-hugging-face-dataset-to-train-and-test%2F20885%2F6 "2023-01-24T16:56:24Z")
Yes this is an annoying error, it looks like they are using `sklearn` in the background.
One way to overcome this (as long as your labels have at least 2 members per group) is to cast the label as a ClassDict first:
```auto
# column we want to stratify with respect to
stratify_column_name = "label"
# create class label column and stratify
dataset.class_encode_column(
stratify_column_name
).train_test_split(
test_size=0.2,
stratify_by_column=stratify_column_name
)
```