Update main.py

main.py

@@ -23,8 +23,10 @@ os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface/datasets"
 os.environ["XDG_CACHE_HOME"] = "/tmp/huggingface"
 
-for path in os.environ.values():
-    os.makedirs(path, exist_ok=True)
+for key in ["HF_HOME", "TRANSFORMERS_CACHE", "HF_HUB_CACHE", "HF_DATASETS_CACHE", "XDG_CACHE_HOME"]:
+    path = os.environ.get(key)
+    if path:
+        os.makedirs(path, exist_ok=True)
 
 # Silence all transformers and huggingface logging
 logging.getLogger("transformers").setLevel(logging.ERROR)