From dd7b9f2ed7c7f648619c8d8bdd910e56a9795d28 Mon Sep 17 00:00:00 2001 From: Victor Morand Date: Thu, 30 Jul 2026 15:22:45 +0200 Subject: [PATCH 1/2] feat(huggingface): support local_path in download/data and bypass split verification for data_files --- src/datamaestro/data/huggingface.py | 12 +++++++++++- src/datamaestro/download/huggingface.py | 4 ++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/datamaestro/data/huggingface.py b/src/datamaestro/data/huggingface.py index 3383d5c..6c91f80 100644 --- a/src/datamaestro/data/huggingface.py +++ b/src/datamaestro/data/huggingface.py @@ -62,7 +62,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 +73,14 @@ 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 diff --git a/src/datamaestro/download/huggingface.py b/src/datamaestro/download/huggingface.py index 552fa21..01172ec 100644 --- a/src/datamaestro/download/huggingface.py +++ b/src/datamaestro/download/huggingface.py @@ -122,9 +122,9 @@ 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 From 80e2ba3279cf5f1af361d9df571b2dfbb3d9478f Mon Sep 17 00:00:00 2001 From: Victor Morand Date: Fri, 31 Jul 2026 10:14:34 +0200 Subject: [PATCH 2/2] fix: documentation, tests and ruff --- docs/source/api/data.md | 25 ++++++++--- docs/source/api/download.rst | 8 +++- src/datamaestro/data/huggingface.py | 10 +++++ src/datamaestro/download/huggingface.py | 6 ++- src/datamaestro/test/test_hf.py | 59 +++++++++++++++++++------ 5 files changed, 86 insertions(+), 22 deletions(-) diff --git a/docs/source/api/data.md b/docs/source/api/data.md index c756040..b5de937 100644 --- a/docs/source/api/data.md +++ b/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", ) ``` diff --git a/docs/source/api/download.rst b/docs/source/api/download.rst index 87febfc..ddf666e 100644 --- a/docs/source/api/download.rst +++ b/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 ----- diff --git a/src/datamaestro/data/huggingface.py b/src/datamaestro/data/huggingface.py index 6c91f80..ac4e6ea 100644 --- a/src/datamaestro/data/huggingface.py +++ b/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``).""" @@ -75,10 +83,12 @@ def download(self): 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: diff --git a/src/datamaestro/download/huggingface.py b/src/datamaestro/download/huggingface.py index 01172ec..3a0c2f6 100644 --- a/src/datamaestro/download/huggingface.py +++ b/src/datamaestro/download/huggingface.py @@ -124,7 +124,11 @@ def hf_download_and_prepare( # 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 or data_files is not None or split is not None) 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 diff --git a/src/datamaestro/test/test_hf.py b/src/datamaestro/test/test_hf.py index 235fd0b..e515df9 100644 --- a/src/datamaestro/test/test_hf.py +++ b/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,7 +86,7 @@ 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( @@ -93,11 +94,22 @@ def test_local_path_replaces_repo_id(self, fake_datasets, tmp_path): 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 --------------------------------------------