File size: 2,831 Bytes
253d988 | 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 | """
Pre-cache pile examples to local JSONL so SAE training doesn't depend on
HuggingFace streaming staying connected. v9b crashed at step 18,808 because
the streaming client closed mid-training; this script removes that risk.
Output: data/pile_cache.jsonl (one {"text": "..."} per line)
Usage: python scripts/cache_pile_examples.py [--n 300000]
"""
import argparse
import json
import time
from pathlib import Path
from datasets import load_dataset
REPO_ROOT = Path(__file__).resolve().parents[1]
CACHE_PATH = REPO_ROOT / "data" / "pile_cache.jsonl"
def main():
p = argparse.ArgumentParser()
p.add_argument("--n", type=int, default=300_000,
help="Number of examples to cache (~600 tok/example avg → 180M tok)")
args = p.parse_args()
CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
# If cache already has enough, skip.
existing = 0
if CACHE_PATH.exists():
with CACHE_PATH.open(encoding="utf-8") as f:
for _ in f:
existing += 1
if existing >= args.n:
print(f"[cache] {CACHE_PATH} already has {existing} >= {args.n}, skipping")
return
print(f"[cache] {existing} examples already in cache, need {args.n - existing} more")
# Robust streaming with retry on transient HF errors.
saved = existing
max_retries = 10
retry = 0
t0 = time.monotonic()
while saved < args.n and retry < max_retries:
try:
ds = load_dataset("monology/pile-uncopyrighted", split="train", streaming=True)
it = iter(ds)
# If we have existing, skip ahead. Streaming has no random access,
# so we just take from the top; some duplication on retry is OK.
with CACHE_PATH.open("a", encoding="utf-8") as f:
for ex in it:
if saved >= args.n:
break
text = ex.get("text", "")
if not text:
continue
f.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
saved += 1
if saved % 5_000 == 0:
elapsed = time.monotonic() - t0
rate = saved / max(elapsed, 1)
eta = (args.n - saved) / max(rate, 1)
print(f"[cache] {saved}/{args.n} rate={rate:.0f}/s ETA={eta/60:.1f}m", flush=True)
break # full pass succeeded
except Exception as e:
retry += 1
print(f"[cache] error: {type(e).__name__}: {str(e)[:160]} retry={retry}/{max_retries}", flush=True)
time.sleep(5)
print(f"[cache] done. {saved} examples in {CACHE_PATH} ({CACHE_PATH.stat().st_size/1024/1024:.0f} MB)")
if __name__ == "__main__":
main()
|