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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
"""FineWeb-Edu sample-10BT -> 32k ByteLevel BPE -> train.bin/val.bin (uint16) + meta.pkl + tokenizer.json.
Same bin format as tinystories_bpe; trainers consume it via --data fineweb_edu.
Stage-2 corpus (recipe order: TinyStories epoch -> FineWeb-Edu -> OLMo2-corpus class).
Phases (all resumable-ish, markers for the watcher):
1. snapshot_download of sample/10BT parquets (~28GB) -> DOWNLOAD_DONE
2. train 32k BPE on ~1.5GB sample from shard 0 -> TOKENIZER_DONE
3. tokenize shard 0: first 20M tokens -> val.bin, rest ->
train.bin (disjoint) -> SMOKE_READY
4. tokenize remaining shards, append to train.bin -> DONE_FINEWEB
Docs are joined with a <|eot|> separator (id 0). vocab 32768 fits uint16.
NFS note: peak disk = raw parquet ~28GB + bins ~20GB; keep raw/ for tokenizer reruns.
"""
import os, pickle, time
from pathlib import Path
import numpy as np
import pyarrow.parquet as pq
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDec
DATA_ROOT = Path(os.environ.get('EPT_DATA_ROOT', Path(__file__).resolve().parent / 'data'))
D = DATA_ROOT / 'fineweb_edu'
RAW = D / 'raw'
D.mkdir(parents=True, exist_ok=True)
VOCAB = 32768
EOT = '<|eot|>'
SAMPLE_BYTES = 1_500_000_000 # tokenizer-training text cap
VAL_TOKENS = 20_000_000
t0 = time.time()
snapshot_download(repo_id='HuggingFaceFW/fineweb-edu', repo_type='dataset',
allow_patterns=['sample/10BT/*'], local_dir=str(RAW))
shards = sorted((RAW / 'sample' / '10BT').glob('*.parquet'))
print(f"DOWNLOAD_DONE {len(shards)} shards {time.time()-t0:.0f}s", flush=True)
def docs(shard, batch_rows=2000):
for b in pq.ParquetFile(shard).iter_batches(batch_size=batch_rows, columns=['text']):
for t in b.column('text').to_pylist():
if t: yield t
tok_path = D / 'tokenizer.json'
if tok_path.exists():
tok = Tokenizer.from_file(str(tok_path))
print(f"TOKENIZER_DONE (cached) vocab={tok.get_vocab_size()}", flush=True)
else:
def sample_iter():
n = 0
for t in docs(shards[0]):
yield t
n += len(t)
if n > SAMPLE_BYTES: return
tok = Tokenizer(BPE(unk_token=None))
tok.pre_tokenizer = ByteLevel(add_prefix_space=False)
tok.decoder = ByteLevelDec()
tok.train_from_iterator(sample_iter(), BpeTrainer(vocab_size=VOCAB, special_tokens=[EOT], show_progress=False))
tok.save(str(tok_path))
print(f"TOKENIZER_DONE vocab={tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)
eot_id = tok.token_to_id(EOT)
assert tok.get_vocab_size() <= 65535 and eot_id is not None
pickle.dump({'vocab_size': tok.get_vocab_size()}, open(D / 'meta.pkl', 'wb'))
def encode_shard(shard, sink, head_sink=None, head_n=0):
"""tokenize one shard; first head_n tokens to head_sink (if given), rest to sink. Returns tokens written."""
wrote_head, n, buf = 0, 0, []
def flush(buf):
nonlocal wrote_head, n
ids = []
for e in tok.encode_batch(buf):
ids.extend(e.ids); ids.append(eot_id)
a = np.array(ids, dtype=np.uint16)
if head_sink is not None and wrote_head < head_n:
k = min(head_n - wrote_head, len(a))
a[:k].tofile(head_sink); wrote_head += k; a = a[k:]
a.tofile(sink); n += len(a)
for t in docs(shard):
buf.append(t)
if len(buf) >= 1000: flush(buf); buf = []
if buf: flush(buf)
return n + wrote_head
train_f = open(D / 'train.bin', 'wb')
val_f = open(D / 'val.bin', 'wb')
n0 = encode_shard(shards[0], train_f, head_sink=val_f, head_n=VAL_TOKENS)
val_f.close(); train_f.flush()
print(f"SMOKE_READY shard0 {n0/1e6:.0f}M tokens (val {VAL_TOKENS/1e6:.0f}M) {time.time()-t0:.0f}s", flush=True)
total = n0
for i, sh in enumerate(shards[1:], 1):
ns = encode_shard(sh, train_f)
total += ns
print(f"[shard {i}/{len(shards)-1}] +{ns/1e6:.0f}M -> total {total/1e9:.2f}B tokens {time.time()-t0:.0f}s", flush=True)
train_f.close()
print(f"DONE_FINEWEB total {total/1e9:.2f}B tokens (train {(total-VAL_TOKENS)/1e9:.2f}B / val {VAL_TOKENS/1e6:.0f}M) vocab {tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)
|
