diff options
| author | YurenHao0426 <Blackhao0426@gmail.com> | 2026-08-31 22:56:32 -0500 |
|---|---|---|
| committer | YurenHao0426 <Blackhao0426@gmail.com> | 2026-08-31 22:56:32 -0500 |
| commit | 0a090e3f597ee0994ea16f0675135a33bf051441 (patch) | |
| tree | a5ea835bd06175fc461b996c18bd4401cc692366 /scripts/prepare_data.py | |
| parent | c87a419fc1318c800be10310d22f28a1060acf1a (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
Diffstat (limited to 'scripts/prepare_data.py')
| -rw-r--r-- | scripts/prepare_data.py | 5 |
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}") |