schihei/llama2.c
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add assert message to give better warning

Browse Source
This commit is contained in:
rahulschand
2023-08-19 13:05:26 +05:30
parent bd182289c5
commit fbefeec1b1
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinystories.py
+1
View File
@@ -196,6 +196,7 @@ class PretokDataset(torch.utils.data.IterableDataset):
shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin"))) shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
# train/test split. let's use only shard 0 for test split, rest train # train/test split. let's use only shard 0 for test split, rest train
shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1] shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
assert len(shard_filenames)>0, f"No bin files found in {bin_dir}"
while True: while True:
rng.shuffle(shard_filenames) rng.shuffle(shard_filenames)
for shard in shard_filenames: for shard in shard_filenames: