diff options
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}") |