dtata fox maybe

Dockerfile

@@ -23,9 +23,7 @@ RUN mkdir -p /app/out && chmod 777 /app/out
 COPY . .
 
 # Prepare the OpenWebText dataset
-RUN mkdir -p /app/data/openwebtext && \
-    cd /app/data/openwebtext && \
-    python prepare.py
+RUN cd /app/data/openwebtext && python prepare.py
 
 # Command to run training
 CMD ["python", "train.py", "--wandb_log=True"]

data/openwebtext/prepare.py

@@ -19,6 +19,8 @@ num_proc_load_dataset = num_proc
 enc = tiktoken.get_encoding("gpt2")
 
 if __name__ == '__main__':
+    # Load dataset with trust_remote_code=True to allow custom code execution
+    dataset = load_dataset("openwebtext", trust_remote_code=True, num_proc=num_proc_load_dataset)
     # takes 54GB in huggingface .cache dir, about 8M documents (8,013,769)
     dataset = load_dataset("openwebtext", num_proc=num_proc_load_dataset)