add more stuff

basedformer/dataset.py

@@ -7,6 +7,7 @@ from torch.utils import data
 from simplejpeg import decode_jpeg
 import simplejpeg
 import pickle
+from pathlib import Path
 
 # Does this work with other block_sizes? doesn't seem to.
 class FbDataset(data.Dataset):
@@ -98,10 +99,12 @@ class ShardedImageDataset(data.Dataset):
         return data
 
 class ImageDatasetBuilder():
-    def __init__(self, folder, name, threads=None):
-        self.folder = folder
-        self.name = name 
-        self.index_name = self.name + ".index"
+    def __init__(self, folder_path, name, threads=None):
+        self.folder_path = Path(folder_path)
+        self.dataset_name = name + ".ds"
+        self.index_name = name + ".index"
+        self.dataset_path = self.folder_path / self.dataset_name
+        self.index_path = self.folder_path / self.index_name
         self.dataset = None
         self.index = None
         self.threads = threads
@@ -119,6 +122,7 @@ class ImageDatasetBuilder():
         if self.is_open:
             raise Exception("Dataset already built")
         
+        self.folder_path.mkdir(parents=True, exist_ok=True)
         dataset = open(self.dataset_path, mode="ab+")
         dataset.flush()
         self.index = []
@@ -131,6 +135,9 @@ class ImageDatasetBuilder():
             self.close()
             print("Dataset closed and flushed.")
 
+        if not self.dataset_path.is_file() or not self.index_path.is_file():
+            raise Exception("Dataset files not found")
+
         self.dataset = open(self.dataset_path, mode="ab+")
         with open(self.index_name, 'rb') as f:
             self.index = pickle.load(f)