File size: 1,587 Bytes
29870b7 | 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 | # stream_data.py - rolling-shard streaming tokenizer (v13/v15 Stage 0).
# Streams Cosmopedia + FineWeb-Edu, tokenizes with the project BPE, writes rolling
# shards of SHARD_TOK tokens each; keeps at most KEEP shards on disk (delete oldest).
# UNIQUE tokens forever - never re-reads (measured overfit wall, v12 8.2).
import numpy as np, itertools, sys
from pathlib import Path
from datasets import load_dataset
from tokenizers import Tokenizer
OUT=Path('/root/dna/data'); OUT.mkdir(parents=True,exist_ok=True)
tok=Tokenizer.from_file('/root/dna/fineweb-tokenizer.json')
SHARD_TOK=25_000_000
KEEP=40
def streams():
fw=load_dataset('HuggingFaceFW/fineweb-edu','sample-100BT',split='train',streaming=True)
co=load_dataset('HuggingFaceTB/cosmopedia','web_samples_v2',split='train',streaming=True)
fi=iter(fw); ci=iter(co)
while True:
for _ in range(3):
yield next(fi)['text']
try: yield next(ci)['text']
except Exception: pass
def main():
it=streams(); si=0; buf=np.empty(SHARD_TOK,np.uint16); n=0; total=0
for txt in it:
ids=tok.encode(txt).ids
for x in ids:
buf[n]=x; n+=1
if n==SHARD_TOK:
tmp=OUT/f'shard_{si:05d}.tmp'; buf.tofile(tmp)
tmp.rename(OUT/f'shard_{si:05d}.u16')
total+=n; print(f'SHARD {si} done total_tokens {total:,}',flush=True)
si+=1; n=0
old=sorted(OUT.glob('shard_*.u16'))
while len(old)>KEEP:
old[0].unlink(); old=old[1:]
if __name__=='__main__': main()
|