Create prepare.py

prepare.py

@@ -0,0 +1,46 @@
+import os
+import tiktoken
+import numpy as np
+from datasets import load_dataset
+from tqdm import tqdm
+
+DATASET_NAME = "HuggingFaceFW/fineweb-edu"
+SAMPLE_NAME = "sample-10BT"
+TARGET_TOKENS = 100_000_000
+NUM_PROC = 8 
+
+enc = tiktoken.get_encoding("gpt2")
+
+def process(example):
+    ids = enc.encode_ordinary(example['text'])
+    ids.append(enc.eot_token)
+    return {'ids': ids, 'len': len(ids)}
+
+if __name__ == "__main__":
+    print(f"Loading streaming dataset {DATASET_NAME}...")
+    dataset = load_dataset(DATASET_NAME, name=SAMPLE_NAME, split='train', streaming=True)
+    
+    all_tokens = []
+    total_tokens = 0
+    
+    pbar = tqdm(total=TARGET_TOKENS, desc="Collecting tokens")
+    for example in dataset:
+        tokens = process(example)['ids']
+        all_tokens.extend(tokens)
+        total_tokens += len(tokens)
+        pbar.update(len(tokens))
+        if total_tokens >= TARGET_TOKENS:
+            break
+    pbar.close()
+
+    n = len(all_tokens)
+    train_data = all_tokens[:int(n*0.95)]
+    val_data = all_tokens[int(n*0.95):]
+
+    for name, d in [('train', train_data), ('val', val_data)]:
+        arr = np.array(d, dtype=np.uint16)
+        filename = f"{name}.bin"
+        arr.tofile(filename)
+        print(f"Saved {filename} with {len(d):,} tokens.")
+
+    print("\nDone!")