diff options
Diffstat (limited to 'data/openwebtext')
| -rw-r--r-- | data/openwebtext/prepare.py | 81 | ||||
| -rw-r--r-- | data/openwebtext/readme.md | 15 |
2 files changed, 96 insertions, 0 deletions
diff --git a/data/openwebtext/prepare.py b/data/openwebtext/prepare.py new file mode 100644 index 0000000..2a9b975 --- /dev/null +++ b/data/openwebtext/prepare.py @@ -0,0 +1,81 @@ +# saves the openwebtext dataset to a binary file for training. following was helpful: +# https://github.com/HazyResearch/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py + +import os +from tqdm import tqdm +import numpy as np +import tiktoken +from datasets import load_dataset # huggingface datasets + +# number of workers in .map() call +# 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 + +enc = tiktoken.get_encoding("gpt2") + +if __name__ == '__main__': + # takes 54GB in huggingface .cache dir, about 8M documents (8,013,769) + 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) + split_dataset['val'] = split_dataset.pop('test') # rename the test split to val + + # this results in: + # >>> split_dataset + # DatasetDict({ + # train: Dataset({ + # features: ['text'], + # num_rows: 8009762 + # }) + # val: Dataset({ + # features: ['text'], + # num_rows: 4007 + # }) + # }) + + # we now want to tokenize the dataset. first define the encoding function (gpt2 bpe) + def process(example): + ids = enc.encode_ordinary(example['text']) # encode_ordinary ignores any special tokens + ids.append(enc.eot_token) # add the end of text token, e.g. 50256 for gpt2 bpe + # note: I think eot should be prepended not appended... hmm. it's called "eot" though... + out = {'ids': ids, 'len': len(ids)} + return out + + # tokenize the dataset + tokenized = split_dataset.map( + process, + remove_columns=['text'], + desc="tokenizing the splits", + num_proc=num_proc, + ) + + # concatenate all the ids in each dataset into one large file we can use for training + for split, dset in tokenized.items(): + arr_len = np.sum(dset['len'], dtype=np.uint64) + filename = os.path.join(os.path.dirname(__file__), f'{split}.bin') + dtype = np.uint16 # (can do since enc.max_token_value == 50256 is < 2**16) + arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,)) + total_batches = 1024 + + idx = 0 + for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'): + # Batch together samples for faster write + batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy') + arr_batch = np.concatenate(batch['ids']) + # Write into mmap + arr[idx : idx + len(arr_batch)] = arr_batch + idx += len(arr_batch) + arr.flush() + + # train.bin is ~17GB, val.bin ~8.5MB + # train has ~9B tokens (9,035,582,198) + # val has ~4M tokens (4,434,897) + + # to read the bin files later, e.g. with numpy: + # m = np.memmap('train.bin', dtype=np.uint16, mode='r') diff --git a/data/openwebtext/readme.md b/data/openwebtext/readme.md new file mode 100644 index 0000000..95eb1bf --- /dev/null +++ b/data/openwebtext/readme.md @@ -0,0 +1,15 @@ + +## openwebtext dataset + +after running `prepare.py` (preprocess) we get: + +- train.bin is ~17GB, val.bin ~8.5MB +- train has ~9B tokens (9,035,582,198) +- val has ~4M tokens (4,434,897) + +this came from 8,013,769 documents in total. + +references: + +- OpenAI's WebText dataset is discussed in [GPT-2 paper](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) +- [OpenWebText](https://skylion007.github.io/OpenWebTextCorpus/) dataset |