Create prepare.py

The data prepare script for Fineweb-Edu-10BT. This script prepares the pretrainingdata binary files.

prepare.py

@@ -0,0 +1,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()