feat(huggingface): support local_path in download/data and bypass split verification for data_files

docs/source/api/data.md

@@ -135,14 +135,29 @@ return Supervised(
 
 Package: `datamaestro.data.huggingface`
 
-For datasets from the HuggingFace Hub:
+For datasets from the HuggingFace Hub or local disk mirrors:
+
+```{eval-rst}
+.. autoxpmconfig:: datamaestro.data.huggingface.HuggingFaceDataset
+```
+
+Example usage:
 
 ```python
-from datamaestro.data.huggingface import DatasetDict
+from datamaestro.data.huggingface import HuggingFaceDataset
+
+# Load from HuggingFace Hub
+ds = HuggingFaceDataset.C(
+    id="squad_dataset",
+    repo_id="squad",
+    split="train",
+)
 
-return DatasetDict(
-    dataset_id="squad",
-    config=None,  # Optional config name
+# Load from a local disk mirror or task output directory
+ds_local = HuggingFaceDataset.C(
+    id="squad_local",
+    repo_id="squad",
+    local_path="/path/to/local/dataset",
 )
 ```
 

docs/source/api/download.rst

@@ -168,7 +168,7 @@ HuggingFace Integration
 
 Package: ``datamaestro.download.huggingface``
 
-For datasets hosted on HuggingFace Hub:
+For datasets hosted on HuggingFace Hub or local disk mirrors:
 
 .. autoclass:: datamaestro.download.huggingface.HFDownloader
 
@@ -178,7 +178,11 @@ For datasets hosted on HuggingFace Hub:
 
     @dataset(url="https://huggingface.co/datasets/squad")
     class Squad(QADataset):
-        HF_DATA = HFDownloader("squad_data", "squad")
+        # Download from HuggingFace Hub
+        HF_DATA = HFDownloader("squad_data", repo_id="squad")
+
+        # Or specify a local mirror path to bypass Hub download
+        LOCAL_DATA = HFDownloader("local_data", repo_id="squad", local_path="/path/to/mirror")
 
 Links
 -----

src/datamaestro/data/huggingface.py

@@ -21,6 +21,14 @@
 
 
 class HuggingFaceDataset(Base):
+    """Adapter for datasets from HuggingFace Hub or local disk mirrors.
+
+    Supports loading datasets via HuggingFace ``datasets`` with support for
+    specific configs, data files, splits, streaming mode, or loading directly
+    from a local mirror/disk path (e.g. saved via ``Dataset.save_to_disk`` or local folder,
+    This can be useful for storing preprocessed versions of the dataset e.g shuffling and filtering).
+    """
+
     repo_id: Param[str]
     """The HuggingFace repository id (e.g. ``user/dataset``)."""
 
@@ -62,7 +70,9 @@ def download(self):
         super().download()
 
         # Streaming mode never materialises anything locally.
-        if self.streaming:
+        # When local_path is set (e.g. from task output directory or local mirror),
+        # there is nothing to download from HF Hub.
+        if self.streaming or self.local_path is not None:
             return
 
         hf_download_and_prepare(
@@ -71,6 +81,16 @@ def download(self):
 
     @cached_property
     def data(self):
+        if self.local_path is not None:
+            from datasets import load_from_disk
+
+            try:
+                return load_from_disk(str(self.local_path))
+            except Exception:
+                from datasets import load_dataset
+
+                return load_dataset(str(self.local_path))
+
         if self.streaming:
             try:
                 from datasets import load_dataset

src/datamaestro/download/huggingface.py

@@ -122,9 +122,13 @@ def hf_download_and_prepare(
     """
     builder, restricted = hf_builder(source, name, data_files, split)
 
-    # A split-restricted build records fewer splits than the dataset
+    # A split-restricted or data_files-restricted build records fewer splits than the dataset
     # metadata declares, which trips ``verify_splits``.
-    kwargs = {"verification_mode": "no_checks"} if restricted else {}
+    kwargs = (
+        {"verification_mode": "no_checks"}
+        if (restricted or data_files is not None or split is not None)
+        else {}
+    )
     builder.download_and_prepare(**kwargs)
     return builder
 

src/datamaestro/test/test_hf.py

@@ -37,6 +37,7 @@ def fake_datasets(monkeypatch):
     """
     fake = types.ModuleType("datasets")
     fake.load_dataset = MagicMock(return_value=MagicMock(name="FakeDataset"))
+    fake.load_from_disk = MagicMock(return_value=MagicMock(name="FakeDiskDataset"))
 
     # Splits the fake Hub repo advertises; tests may override.
     fake.SPLITS = ["train", "validation", "test"]
@@ -85,19 +86,30 @@ def test_passes_name_split_streaming(self, fake_datasets):
             streaming=True,
         )
 
-    def test_local_path_replaces_repo_id(self, fake_datasets, tmp_path):
+    def test_local_path_loads_from_disk(self, fake_datasets, tmp_path):
         local = tmp_path / "mirror"
         local.mkdir()
         ds = HuggingFaceDataset.C(
             id="test.hf.2",
             repo_id="user/dataset",
             local_path=local,
         )
-        _ = ds.data
-        # Source is the local path, not the repo id.
-        positional = fake_datasets.load_dataset_builder.call_args.args
-        assert positional[0] == str(local)
-        assert positional[1] is None  # no name
+        data = ds.data
+        fake_datasets.load_from_disk.assert_called_once_with(str(local))
+        assert data is fake_datasets.load_from_disk.return_value
+
+    def test_local_path_fallback_to_load_dataset(self, fake_datasets, tmp_path):
+        local = tmp_path / "mirror"
+        local.mkdir()
+        fake_datasets.load_from_disk.side_effect = Exception("Not a disk dataset")
+        ds = HuggingFaceDataset.C(
+            id="test.hf.2b",
+            repo_id="user/dataset",
+            local_path=local,
+        )
+        data = ds.data
+        fake_datasets.load_dataset.assert_called_once_with(str(local))
+        assert data is fake_datasets.load_dataset.return_value
 
     def test_default_args(self, fake_datasets):
         """Non-streaming access goes through the builder, and returns what
@@ -149,11 +161,13 @@ def test_download_prepares_builder(self, fake_datasets):
             "config-a",
             data_files="train.jsonl.gz",
         )
-        prepared(fake_datasets).download_and_prepare.assert_called_once_with()
+        prepared(fake_datasets).download_and_prepare.assert_called_once_with(
+            verification_mode="no_checks"
+        )
         # No in-RAM instantiation of the dataset.
         fake_datasets.load_dataset.assert_not_called()
 
-    def test_download_uses_local_path(self, fake_datasets, tmp_path):
+    def test_download_with_local_path_is_noop(self, fake_datasets, tmp_path):
         local = tmp_path / "mirror"
         local.mkdir()
         ds = HuggingFaceDataset.C(
@@ -162,7 +176,7 @@ def test_download_uses_local_path(self, fake_datasets, tmp_path):
             local_path=local,
         )
         ds.download()
-        assert fake_datasets.load_dataset_builder.call_args.args[0] == str(local)
+        fake_datasets.load_dataset_builder.assert_not_called()
 
     def test_download_streaming_is_noop(self, fake_datasets):
         ds = HuggingFaceDataset.C(
@@ -202,10 +216,12 @@ def test_restricts_on_a_sliced_split(self, fake_datasets):
         assert list(builder.config.data_files) == ["train"]
 
     def test_no_restriction_for_compound_split(self, fake_datasets):
-        """``train+test`` spans several splits: prepare everything."""
+        """``train+test`` spans several splits: prepare everything with verification disabled."""
         builder = self._download(fake_datasets, split="train+test")
         assert list(builder.config.data_files) == ["train", "validation", "test"]
-        builder.download_and_prepare.assert_called_once_with()
+        builder.download_and_prepare.assert_called_once_with(
+            verification_mode="no_checks"
+        )
 
     def test_no_restriction_without_split(self, fake_datasets):
         builder = self._download(fake_datasets)
@@ -221,9 +237,11 @@ def test_no_restriction_for_unknown_split(self, fake_datasets):
     def test_no_restriction_when_already_single_split(self, fake_datasets):
         fake_datasets.SPLITS = ["train"]
         builder = self._download(fake_datasets, split="train")
-        # Nothing to gain: one builder, prepared with checks left on.
+        # Split is set: verification is bypassed.
         assert fake_datasets.load_dataset_builder.call_count == 1
-        builder.download_and_prepare.assert_called_once_with()
+        builder.download_and_prepare.assert_called_once_with(
+            verification_mode="no_checks"
+        )
 
     def test_no_restriction_for_script_builder(self, fake_datasets):
         """A script-based builder exposes no per-split ``data_files``."""
@@ -237,7 +255,20 @@ def no_data_files(source, name=None, data_files=None):
         fake_datasets.load_dataset_builder.side_effect = no_data_files
         builder = self._download(fake_datasets, split="train")
         assert fake_datasets.load_dataset_builder.call_count == 1
-        builder.download_and_prepare.assert_called_once_with()
+        builder.download_and_prepare.assert_called_once_with(
+            verification_mode="no_checks"
+        )
+
+    def test_data_files_bypasses_split_verification(self, fake_datasets):
+        ds = HuggingFaceDataset.C(
+            id="test.hf.df",
+            repo_id="user/dataset",
+            data_files="data.parquet",
+        )
+        ds.download()
+        prepared(fake_datasets).download_and_prepare.assert_called_once_with(
+            verification_mode="no_checks"
+        )
 
 
 # ---- Identity: Param vs Meta --------------------------------------------