File size: 1,100 Bytes
d9350d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import os
from tqdm import tqdm
import numpy as np
import tiktoken
from datasets import load_dataset
enc = tiktoken.get_encoding("gpt2")
if __name__ == '__main__':
dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train")
split_dataset = dataset.train_test_split(test_size=0.0005, seed=2357, shuffle=True)
split_dataset['val'] = split_dataset.pop('test')
def process(example):
ids = enc.encode_ordinary(example['text'])
ids.append(enc.eot_token)
return {'ids': ids, 'len': len(ids)}
tokenized = split_dataset.map(process, remove_columns=['text'], desc="tokenizing", num_proc=8)
for split, dset in tokenized.items():
arr_len = np.sum(dset['len'], dtype=np.int64)
filename = os.path.join(os.path.dirname(__file__), f'{split}.bin')
arr = np.memmap(filename, dtype=np.uint16, mode='w+', shape=(arr_len,))
idx = 0
for example in tqdm(dset, desc=f"writing {filename}"):
arr[idx : idx + example['len']] = example['ids']
idx += example['len']
arr.flush() |