diff --git a/data/openwebtext/prepare.py b/data/openwebtext/prepare.py index aad8170..eebfe53 100644 --- a/data/openwebtext/prepare.py +++ b/data/openwebtext/prepare.py @@ -11,8 +11,13 @@ from datasets import load_dataset # huggingface datasets # good number to use is ~order number of cpu cores // 2 num_proc = 8 +# number of workers in load_dataset() call +# best number might be different from num_proc above as it also depends on NW speed. +# it is better than 1 usually though +num_proc_load_dataset = num_proc + # takes 54GB in huggingface .cache dir, about 8M documents (8,013,769) -dataset = load_dataset("openwebtext") +dataset = load_dataset("openwebtext", num_proc=num_proc_load_dataset) # owt by default only contains the 'train' split, so create a test split split_dataset = dataset["train"].train_test_split(test_size=0.0005, seed=2357, shuffle=True)