summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYurenHao0426 <Blackhao0426@gmail.com>2026-08-31 22:56:32 -0500
committerYurenHao0426 <Blackhao0426@gmail.com>2026-08-31 22:56:32 -0500
commit0a090e3f597ee0994ea16f0675135a33bf051441 (patch)
treea5ea835bd06175fc461b996c18bd4401cc692366
parentc87a419fc1318c800be10310d22f28a1060acf1a (diff)
prepare_data: write .partial then rename, so interrupted prep can't leave a half train.bin that gets mistaken for complete
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe
-rw-r--r--scripts/prepare_data.py5
1 files changed, 4 insertions, 1 deletions
diff --git a/scripts/prepare_data.py b/scripts/prepare_data.py
index 626e7a8..b7684b8 100644
--- a/scripts/prepare_data.py
+++ b/scripts/prepare_data.py
@@ -16,7 +16,9 @@ from datasets import load_dataset
enc = tiktoken.get_encoding("gpt2")
ds = load_dataset(a.dataset, name=a.name, split="train", streaming=True)
train_path, val_path = os.path.join(a.out, "train.bin"), os.path.join(a.out, "val.bin")
-ftr, fva = open(train_path, "wb"), open(val_path, "wb")
+# write to .partial and rename on completion, so an interrupted prep never leaves a
+# half-written train.bin that run_ladder.sh's existence check would mistake for done
+ftr, fva = open(train_path + ".partial", "wb"), open(val_path + ".partial", "wb")
n_tr = n_va = 0
target_tr, target_va = int(a.tokens), int(a.val_tokens)
buf = []
@@ -32,4 +34,5 @@ for i, ex in enumerate(ds):
if n_tr >= target_tr and n_va >= target_va:
break
ftr.close(); fva.close()
+os.rename(train_path + ".partial", train_path); os.rename(val_path + ".partial", val_path)
print(f"DONE train {n_tr/1e6:.1f}M val {n_va/1e6:.1f}M -> {a.out}")