Repository: PolicyEngine/policyengine-core

Issue:
Ensure Hugging Face datasets download to correct location
Currently, in cases where the repo hasn't already downloaded a dataset from Hugging Face, it downloads the dataset to a temp folder instead of the intended destination. This causes `Simulation.build_dataset()` to fail, as in the tests at https://github.com/PolicyEngine/policyengine-api/actions/runs/12402725316/job/34625177069.

Discussion:

Solution patch:
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
index e69de29b..531d0bb9 100644
--- a/changelog_entry.yaml
+++ b/changelog_entry.yaml
@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+    - Explicitly set local directory when downloading datasets from Hugging Face
\ No newline at end of file
diff --git a/policyengine_core/data/dataset.py b/policyengine_core/data/dataset.py
index 566a45a8..c6cec09f 100644
--- a/policyengine_core/data/dataset.py
+++ b/policyengine_core/data/dataset.py
@@ -501,4 +501,5 @@ class Dataset:
             repo=f"{owner_name}/{model_name}",
             repo_filename=file_name,
             version=version,
+            local_dir=self.file_path.parent,
         )
diff --git a/policyengine_core/tools/hugging_face.py b/policyengine_core/tools/hugging_face.py
index 2c18585a..d92f4447 100644
--- a/policyengine_core/tools/hugging_face.py
+++ b/policyengine_core/tools/hugging_face.py
@@ -13,7 +13,10 @@ with warnings.catch_warnings():
 
 
 def download_huggingface_dataset(
-    repo: str, repo_filename: str, version: str = None
+    repo: str,
+    repo_filename: str,
+    version: str = None,
+    local_dir: str | None = None,
 ):
     """
     Download a dataset from the Hugging Face Hub.
@@ -22,6 +25,7 @@ def download_huggingface_dataset(
         repo (str): The Hugging Face repo name, in format "{org}/{repo}".
         repo_filename (str): The filename of the dataset.
         version (str, optional): The version of the dataset. Defaults to None.
+        local_dir (str, optional): The local directory to save the dataset to. Defaults to None.
     """
     # Attempt connection to Hugging Face model_info endpoint
     # (https://huggingface.co/docs/huggingface_hub/v0.26.5/en/package_reference/hf_api#huggingface_hub.HfApi.model_info)
@@ -52,6 +56,7 @@ def download_huggingface_dataset(
         filename=repo_filename,
         revision=version,
         token=authentication_token,
+        local_dir=local_dir,
     )
 
 


Test patch:
diff --git a/tests/core/tools/test_hugging_face.py b/tests/core/tools/test_hugging_face.py
index 813a1278..54f04bd8 100644
--- a/tests/core/tools/test_hugging_face.py
+++ b/tests/core/tools/test_hugging_face.py
@@ -15,6 +15,7 @@ class TestHuggingFaceDownload:
         test_repo = "test_repo"
         test_filename = "test_filename"
         test_version = "test_version"
+        test_dir = "test_dir"
 
         with patch(
             "policyengine_core.tools.hugging_face.hf_hub_download"
@@ -29,7 +30,7 @@ class TestHuggingFaceDownload:
                 )
 
                 download_huggingface_dataset(
-                    test_repo, test_filename, test_version
+                    test_repo, test_filename, test_version, test_dir
                 )
 
                 mock_download.assert_called_with(
@@ -37,6 +38,7 @@ class TestHuggingFaceDownload:
                     repo_type="model",
                     filename=test_filename,
                     revision=test_version,
+                    local_dir=test_dir,
                     token=None,
                 )
 
@@ -45,6 +47,7 @@ class TestHuggingFaceDownload:
         test_repo = "test_repo"
         test_filename = "test_filename"
         test_version = "test_version"
+        test_dir = "test_dir"
 
         with patch(
             "policyengine_core.tools.hugging_face.hf_hub_download"
@@ -61,7 +64,7 @@ class TestHuggingFaceDownload:
                     mock_token.return_value = "test_token"
 
                     download_huggingface_dataset(
-                        test_repo, test_filename, test_version
+                        test_repo, test_filename, test_version, test_dir
                     )
                     mock_download.assert_called_with(
                         repo_id=test_repo,
@@ -69,6 +72,7 @@ class TestHuggingFaceDownload:
                         filename=test_filename,
                         revision=test_version,
                         token=mock_token.return_value,
+                        local_dir=test_dir,
                     )
 
     def test_download_private_repo_no_token(self):
@@ -76,6 +80,7 @@ class TestHuggingFaceDownload:
         test_repo = "test_repo"
         test_filename = "test_filename"
         test_version = "test_version"
+        test_dir = "test_dir"
 
         with patch(
             "policyengine_core.tools.hugging_face.hf_hub_download"
@@ -93,7 +98,7 @@ class TestHuggingFaceDownload:
 
                     with pytest.raises(Exception):
                         download_huggingface_dataset(
-                            test_repo, test_filename, test_version
+                            test_repo, test_filename, test_version, test_dir
                         )
                         mock_download.assert_not_called()