Configure Hugging Face cache directories for dataset preparation

zenml-io · Jan 27, 2025 · aecfe12 · aecfe12
1 parent df2de46
commit aecfe12
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/llm-finetuning/configs/generate_code_dataset.yaml b/llm-finetuning/configs/generate_code_dataset.yaml
@@ -2,6 +2,11 @@
 settings:
   docker:
     requirements: requirements.txt
+    apt_packages:
+      - git
+    environment:
+      HF_HOME: "/tmp/huggingface"
+      HF_HUB_CACHE: "/tmp/huggingface"
 
 # pipeline configuration
 parameters:
@@ -11,4 +16,4 @@ steps:
   mirror_repositories:
     parameters:
       repositories:
-        - zenml
+        - zenml
diff --git a/llm-finetuning/steps/prepare_dataset.py b/llm-finetuning/steps/prepare_dataset.py
@@ -6,8 +6,14 @@
 """
 
 import os
-from typing import Dict
+from pathlib import Path
+
+# Set cache directories before importing HF libraries
+os.environ["HF_HOME"] = "/tmp/huggingface"
+os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
+os.makedirs("/tmp/huggingface", exist_ok=True)
 
+from typing import Dict
 import pandas as pd
 from datasets import Dataset
 from huggingface_hub import HfApi