Configure Hugging Face cache directories for dataset preparation

Author: strickvl

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

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