File size: 8,067 Bytes
2ae6691 | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | """Offline data preparation (parallel tokenization).
Streams raw text from a Hugging Face dataset, tokenizes it across many CPU
worker processes, and appends the token ids to a per-source .bin file (uint32).
Documents are separated by the tokenizer's end-of-text id.
Code sources are read one language at a time (e.g. "stack_python" reads only
Python) and pass a quality filter that drops data dumps and minified files.
--skip_docs skips documents an earlier run already used, so batches never
overlap. Each run prints the value to pass as the next run's --skip_docs.
"""
import os
# each worker is its own process, so the tokenizer must NOT also spawn internal
# threads — those Rust threads race Python's shutdown and crash on exit
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import argparse
import itertools
import multiprocessing as mp
import re
import numpy as np
from datasets import load_dataset
from .tokenizer import Tokenizer
# stop cleanly (keeping everything written) if the shared disk gets this full
MIN_FREE_BYTES = 300 * 1024**3
DTYPE = np.uint32 # vocab > 65535, so 2 bytes is not enough
def free_bytes(path: str) -> int:
st = os.statvfs(path)
return st.f_bavail * st.f_frsize
# each source: (HF repo, config, data_dir, text column, split, is_code)
SOURCES = {
"fineweb": ("HuggingFaceFW/fineweb-edu", "sample-350BT", None, "text", "train", False),
"finemath": ("HuggingFaceTB/finemath", "finemath-4plus", None, "text", "train", False),
"finemath3": ("HuggingFaceTB/finemath", "finemath-3plus", None, "text", "train", False),
"cosmopedia": ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", None, "text", "train", False),
"infiwebmath": ("HuggingFaceTB/finemath", "infiwebmath-4plus", None, "text", "train", False),
"fineweb_hq": ("HuggingFaceFW/fineweb-edu", "sample-350BT", None, "text", "train", False),
"stack": ("bigcode/the-stack-dedup", None, None, "content", "train", True),
}
# one source per language: "stack_python" reads only the data/python folder
CODE_LANGS = ["python", "javascript", "java", "cpp", "typescript",
"go", "rust", "shell", "sql"]
for _lang in CODE_LANGS:
SOURCES[f"stack_{_lang}"] = ("bigcode/the-stack-dedup", None, f"data/{_lang}",
"content", "train", True)
# keep only documents with at least this quality score (fineweb-edu's 0-5 "int_score")
MIN_INT_SCORE = {"fineweb_hq": 4}
# --- code quality filter (StarCoder-style rules) --------------------------
_HEX = re.compile(r"0x[0-9a-fA-F]{2,}")
def _looks_like_good_code(text: str) -> bool:
lines = text.split("\n")
if max(len(l) for l in lines) > 1000: # minified / one giant line
return False
if len(text) / len(lines) > 100: # very long lines on average
return False
if sum(c.isalnum() for c in text) / len(text) < 0.25: # mostly symbols
return False
if sum(len(m) for m in _HEX.findall(text)) / len(text) > 0.2: # hex data tables
return False
if sum(c.isdigit() for c in text) / len(text) > 0.4: # numeric data tables
return False
return True
# --- worker side ---------------------------------------------------------
# each worker builds its own tokenizer once (tokenizers aren't shareable
# across processes), then reuses it for every document it receives
_worker_tok = None
_worker_is_code = False
def _init_worker(is_code):
global _worker_tok, _worker_is_code
_worker_tok = Tokenizer()
_worker_is_code = is_code
def _encode_one(text: str) -> np.ndarray:
# a rejected document comes back empty; the main process skips it
if not text or (_worker_is_code and not _looks_like_good_code(text)):
return np.empty(0, dtype=DTYPE)
ids = _worker_tok.encode(text)
ids.append(_worker_tok.eot_id)
return np.array(ids, dtype=DTYPE)
# --- main side -----------------------------------------------------------
SCORE_STATS = {"seen": 0, "low": 0} # docs streamed / docs below min score
def _doc_texts(ds, text_col, min_score=0):
# pull the text out of each document, skipping low-quality ones if asked
for doc in ds:
SCORE_STATS["seen"] += 1
if min_score and (doc.get("int_score") or 0) < min_score:
SCORE_STATS["low"] += 1
continue
yield doc[text_col] or ""
def prepare(source: str, out_dir: str, target_tokens: int,
num_workers: int, chunk_docs: int = 10_000, skip_docs: int = 0,
data_files=None):
repo, config, data_dir, text_col, split, is_code = SOURCES[source]
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, f"{source}.bin")
# fresh data appended to an old file would mix two different runs
if skip_docs == 0 and os.path.exists(out_path) and os.path.getsize(out_path) > 0:
raise SystemExit(f"[refuse] {out_path} already has data: delete it first")
if data_files:
# read one specific file directly (e.g. a late shard for a held-out set)
ds = load_dataset(repo, data_files=data_files, split=split, streaming=True)
else:
ds = load_dataset(repo, config, data_dir=data_dir, split=split, streaming=True)
print(f"[read] {repo} files={data_files} config={config} data_dir={data_dir}", flush=True)
texts = _doc_texts(ds, text_col, MIN_INT_SCORE.get(source, 0))
# skip documents an earlier run already consumed, so batches never overlap
if skip_docs > 0:
print(f" [skip] skipping first {skip_docs:,} documents", flush=True)
texts = itertools.islice(texts, skip_docs, None)
written = n_docs = n_kept = 0
stopped_early = False
# imap keeps results IN ORDER and pulls lazily, so memory stays small
pool = mp.Pool(num_workers, initializer=_init_worker, initargs=(is_code,))
try:
with open(out_path, "ab") as f:
for arr in pool.imap(_encode_one, texts, chunksize=64):
n_docs += 1 # counts every streamed doc (for --skip_docs)
if arr.size:
arr.tofile(f) # main process is the only writer
written += arr.size
n_kept += 1
if written >= target_tokens:
break
if n_docs % chunk_docs == 0:
if free_bytes(out_dir) < MIN_FREE_BYTES:
print(f"\n[stop] free space below {MIN_FREE_BYTES/1e9:.0f} GB",
flush=True)
stopped_early = True
break
finally:
pool.terminate()
pool.join()
if MIN_INT_SCORE.get(source):
low = 100 * SCORE_STATS["low"] / max(SCORE_STATS["seen"], 1)
print(f"\n[score] {low:.0f}% of streamed docs had int_score < {MIN_INT_SCORE[source]}")
status = "STOPPED EARLY (disk)" if stopped_early else "done"
dropped = 100 * (1 - n_kept / max(n_docs, 1))
print(f"\n[{status}] {source}: {written/1e6:.1f}M tokens, "
f"{n_docs:,} docs read, {dropped:.0f}% dropped by code filter")
print(f"[next] for the following run, pass --skip_docs {skip_docs + n_docs}",
flush=True)
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--source", required=True, choices=list(SOURCES))
p.add_argument("--out_dir", default="/fsx/ajay-stu/ar-data/bin")
p.add_argument("--target_tokens", type=int, default=100_000_000)
p.add_argument("--skip_docs", type=int, default=0)
p.add_argument("--num_workers", type=int, default=max(1, mp.cpu_count() - 4))
p.add_argument("--data_files", default=None)
args = p.parse_args()
prepare(args.source, args.out_dir, args.target_tokens, args.num_workers,
skip_docs=args.skip_docs, data_files=args.data_files or None)
# HF's fast tokenizer leaves native threads that crash normal interpreter
# shutdown; the file is closed above, so exit immediately
os._exit(0) |