Preallocate token buffer in prepare_data to cut peak RAM ~20x

[Mantissagithub]

tokenize_documents builds the token stream in a Python list of ints, then converts to uint16 at the end. A Python int in a list costs ~40 bytes (28 B object + 8 B pointer + list overallocation) versus the 2 bytes the output array needs, so the default 100M-token split peaks over 4 GB. That makes data prep the peak-memory step of the repo, and it's easy to OOM or swap on a machine that would otherwise be fine.

Writing straight into a preallocated buffer avoids it:

-    tokens = []
+    tokens = np.empty(token_budget, dtype=np.uint16)
+    n = 0
     for doc in dataset_iter:
         ...
-        tokens.extend(doc_tokens[:keep])
+        tokens[n:n + keep] = doc_tokens[:keep]
+        n += keep

Peak RSS in tokenize_documents, measured with ru_maxrss in separate processes over an identical document stream (RTX 4060 laptop, numpy 2.x):

tokens	before	after
2M	73.7 MB	0.0 MB
5M	196.2 MB	2.2 MB
10M	402.1 MB	12.0 MB
100M	4116.4 MB	198.8 MB

100M is the default --train_tokens budget, measured on the real FineWeb stream (100,000,000 tokens, 143,651 documents).

Output is unchanged. doc_starts is still recorded before truncation, so a document truncated at the budget boundary still gets its start. Checked against the old implementation across 641 configurations — budget landing on and around document boundaries, empty documents, a first document larger than the budget, stream running out early, and real GPT-2 tokenization at budgets up to 5M — plus the full 100M split. All give bit-identical tokens and doc_starts.

[Mantissagithub]

cc @akshayvegesna @saaaaaamip