dataset in dockerfile

Dockerfile

@@ -15,4 +15,5 @@ RUN pip install datasets
 COPY . .
 RUN mkdir -p /code/.cache
 RUN chmod -R 777 /code
+RUN python prepare.py
 CMD ["python", "run.py"]

prepare.py

@@ -0,0 +1,22 @@
+import pickle
+import datasets
+import os
+
+if __name__ == "__main__":
+    cache_file = "dataset_cache.pkl"
+    if os.path.exists(cache_file):
+        # Load dataset from cache
+        with open(cache_file, "rb") as file:
+            dataset = pickle.load(file)
+        print("Dataset loaded from cache.")
+    else:
+        # Load dataset using datasets.load_dataset()
+        dataset = datasets.load_dataset("renumics/cifar100-enriched", split="train")
+        print("Dataset loaded using datasets.load_dataset().")
+
+        # Save dataset to cache
+        with open(cache_file, "wb") as file:
+            pickle.dump(dataset, file)
+
+        print("Dataset saved to cache.")
+