diff --git a/packages/timenet-connectors/src/timenet_connectors/bases/huggingface.py b/packages/timenet-connectors/src/timenet_connectors/bases/huggingface.py index 4ba7ba727..ca720900c 100644 --- a/packages/timenet-connectors/src/timenet_connectors/bases/huggingface.py +++ b/packages/timenet-connectors/src/timenet_connectors/bases/huggingface.py @@ -17,8 +17,8 @@ from timenet.connectors import BaseConnector from timenet.dataset import TimeFDataset -from timenet.errors import TimeFValidationError -from timenet_connectors.sources.huggingface_hub import download_hub_files +from timenet.errors import TimeFValidationError, TimeNetDownloadError +from timenet_connectors.sources.huggingface_hub import hub_snapshot _COMMIT_SHA = re.compile(r"[0-9a-f]{40}") @@ -26,7 +26,7 @@ @dataclass(frozen=True) class HubFileRef: - """One local source file and the part of the dataset it holds.""" + """One local source file, or the directory of a pattern's files, and the part of the dataset it holds.""" path: Path part: str @@ -46,7 +46,9 @@ class BaseHuggingFaceConnector(BaseConnector[HubFileRef], ABC): SOURCE_REVISION: ClassVar[str] """An immutable commit: a mutable branch can silently change rebuilt datasets.""" SOURCE_FILES: ClassVar[Mapping[str, str]] - """Each part's repository-relative filename, in the order ``download()`` returns them.""" + """Each part's repository-relative filename, in the order ``download()`` returns them. A filename can end + in a glob pattern, for example ``data/*.parquet``. Then the path of the part is the directory of the + files that match the pattern.""" def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) @@ -97,8 +99,12 @@ async def download_async(self, cache_dir: Path) -> list[HubFileRef]: Raises: ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing. TimeNetDownloadError: If a declared file is absent from the revision or cannot be downloaded. - """ # ruff: ignore[docstring-extraneous-exception] - both are raised by download_hub_files - paths = await download_hub_files( - self.HF_REPO, list(self.SOURCE_FILES.values()), revision=self.SOURCE_REVISION, cache_dir=cache_dir - ) - return [HubFileRef(path, part) for part, path in zip(self.SOURCE_FILES, paths, strict=True)] + """ # ruff: ignore[docstring-extraneous-exception] - also raised by hub_snapshot + root = await hub_snapshot(self.HF_REPO, self.SOURCE_REVISION, cache_dir, tuple(self.SOURCE_FILES.values())) + refs = [] + for part, name in self.SOURCE_FILES.items(): + path = root / name if "*" not in name else root / Path(name).parent + if not path.exists(): + raise TimeNetDownloadError(f"{self.HF_REPO}@{self.SOURCE_REVISION} has no {name}") + refs.append(HubFileRef(path, part)) + return refs diff --git a/packages/timenet-connectors/src/timenet_connectors/sources/huggingface_hub.py b/packages/timenet-connectors/src/timenet_connectors/sources/huggingface_hub.py index 46f120f77..23e988d06 100644 --- a/packages/timenet-connectors/src/timenet_connectors/sources/huggingface_hub.py +++ b/packages/timenet-connectors/src/timenet_connectors/sources/huggingface_hub.py @@ -1,17 +1,16 @@ """Helpers for downloading files from pinned Hugging Face dataset revisions. -:func:`download_hub_files` downloads with ``huggingface_hub.snapshot_download``, limited to the -requested files. The Hub client runs the downloads in its own thread pool. Its progress bars stay -hidden: the bar that counts file bytes reaches the ambient :mod:`~timenet_connectors.download.progress` -sink instead, so the build CLI shows the download like any other connector download. -``huggingface_hub`` is declared in the connector's ``requirements.txt`` and imported lazily, which -keeps ``--no-isolation`` usable. The Hub client reads ``HF_TOKEN`` from the environment, so gated -datasets work with no extra wiring. +:func:`hub_snapshot` downloads with ``huggingface_hub.snapshot_download``. It downloads only the files +that match the requested patterns, for example a filename or ``data/*.parquet``. The Hub client runs the +downloads in its own thread pool. Its progress bars stay hidden. The bar that counts file bytes reaches +the ambient :mod:`~timenet_connectors.download.progress` sink instead. Thus the build CLI shows the +download like any other connector download. ``huggingface_hub`` is declared in the connector's ``requirements.txt`` +and imported lazily, which keeps ``--no-isolation`` usable. The Hub client reads ``HF_TOKEN`` from the +environment, so gated datasets work with no extra wiring. """ import asyncio from collections.abc import Sequence -from fnmatch import fnmatchcase from pathlib import Path import threading from typing import Any @@ -30,21 +29,22 @@ ) -async def download_hub_files(repo: str, filenames: Sequence[str], *, revision: str, cache_dir: Path) -> list[Path]: - """Download files from a dataset revision with one ``snapshot_download`` call. +async def hub_snapshot(repo: str, revision: str, cache_dir: Path, patterns: Sequence[str]) -> Path: + """Download the files of a dataset revision that match the patterns, with one ``snapshot_download`` call. Args: repo: The Hub dataset id, such as ``"BASH-Lab/SensorCaps"``. - filenames: Repository-relative paths of the files to download. revision: The revision to read, normally a pinned commit. cache_dir: Directory where the connector caches Hub files. + patterns: Repository-relative filenames or glob patterns, for example ``"data/*.parquet"``. Returns: - The local paths, in the order of ``filenames``. + The local snapshot directory. It holds the matching files at their repository paths. + The caller must check that the files it needs are there. Raises: ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing. - TimeNetDownloadError: If the download fails, or a file is absent from the revision. + TimeNetDownloadError: If the download fails. """ try: # Optional connector dependency. @@ -83,42 +83,11 @@ def update(self, n: float | None = 1) -> bool | None: repo_type="dataset", revision=revision, cache_dir=cache_dir, - allow_patterns=list(filenames), + allow_patterns=list(patterns), tqdm_class=_SinkProgress, ) ) ) except Exception as exc: raise TimeNetDownloadError(f"Could not download {repo}@{revision}: {exc}") from exc - paths = [snapshot / name for name in filenames] - missing = [name for name, path in zip(filenames, paths, strict=True) if not path.is_file()] - if missing: - raise TimeNetDownloadError(f"{repo}@{revision} has no {', '.join(missing)}") - return paths - - -async def hub_snapshot(repo: str, revision: str, cache_dir: Path, patterns: tuple[str, ...]) -> Path: - """Download files that match the patterns from a pinned dataset revision. - - Returns: - The local snapshot directory. - - Raises: - ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing. - TimeNetDownloadError: If the revision cannot be listed or no files match. - """ - try: - # Optional connector dependency. - from huggingface_hub import list_repo_files # ruff: ignore[import-outside-top-level] - except ImportError as exc: - raise ImportError(_MISSING_HUB_CLIENT) from exc - - try: - names = await asyncio.to_thread(list_repo_files, repo, repo_type="dataset", revision=revision) - except Exception as exc: - raise TimeNetDownloadError(f"Could not list {repo}@{revision}: {exc}") from exc - selected = [name for name in names if any(fnmatchcase(name, pattern) for pattern in patterns)] - if not selected: - raise TimeNetDownloadError(f"{repo}@{revision} has no files matching {patterns}") - paths = await download_hub_files(repo, selected, revision=revision, cache_dir=cache_dir) - return paths[0].parents[len(Path(selected[0]).parts) - 1] + return snapshot diff --git a/packages/timenet-connectors/tests/test_huggingface_hub.py b/packages/timenet-connectors/tests/test_huggingface_hub.py index 3379e429d..f8c853483 100644 --- a/packages/timenet-connectors/tests/test_huggingface_hub.py +++ b/packages/timenet-connectors/tests/test_huggingface_hub.py @@ -1,6 +1,7 @@ """Tests for the pinned Hub base and its download helper, driven against a fake huggingface_hub.""" import asyncio +from fnmatch import fnmatchcase from pathlib import Path import sys from types import SimpleNamespace @@ -11,7 +12,7 @@ from timenet.errors import TimeFValidationError, TimeNetDownloadError from timenet_connectors.bases.huggingface import BaseHuggingFaceConnector, HubFileRef from timenet_connectors.download.progress import DownloadProgress, progress_sink -from timenet_connectors.sources.huggingface_hub import download_hub_files, hub_snapshot +from timenet_connectors.sources.huggingface_hub import hub_snapshot _SHA = "0" * 40 @@ -37,8 +38,9 @@ def update(self, n=1): return None -def _fake_hub(monkeypatch, tmp_path, present=None, error=None, bars=None): - """Install a fake huggingface_hub whose snapshot_download writes the allowed files into a snapshot.""" +def _fake_hub(monkeypatch, tmp_path, files=("data/train.jsonl", "data/test.jsonl"), error=None, bars=None): + """Install a fake huggingface_hub. Its snapshot_download writes the revision's ``files`` that match the + allowed patterns into a snapshot, as the Hub does.""" calls = [] def snapshot_download(repo, **kwargs): @@ -48,8 +50,8 @@ def snapshot_download(repo, **kwargs): if error is not None: raise error snapshot = tmp_path / "snapshot" - for name in kwargs["allow_patterns"]: - if present is None or name in present: + for name in files: + if any(fnmatchcase(name, pattern) for pattern in kwargs["allow_patterns"]): path = snapshot / name path.parent.mkdir(parents=True, exist_ok=True) path.write_text(name) @@ -111,11 +113,9 @@ def test_download_failure_names_the_pinned_revision(monkeypatch, tmp_path): def test_a_file_absent_from_the_revision_fails(monkeypatch, tmp_path): - _fake_hub(monkeypatch, tmp_path, present={"data/train.jsonl"}) + _fake_hub(monkeypatch, tmp_path, files=("data/train.jsonl",)) with pytest.raises(TimeNetDownloadError, match=rf"org/repo@{_SHA} has no data/test\.jsonl$"): - asyncio.run( - download_hub_files("org/repo", ["data/train.jsonl", "data/test.jsonl"], revision=_SHA, cache_dir=tmp_path) - ) + asyncio.run(_Pinned().download_async(tmp_path / "cache")) def test_missing_hub_library_states_the_fix(monkeypatch, tmp_path): @@ -166,17 +166,40 @@ def test_convert_checks_refs_before_the_subclass_runs(): assert _Pinned.convert.__name__ == "convert" -def test_snapshot_selects_patterns_and_keeps_the_snapshot_root(monkeypatch, tmp_path): - calls = _fake_hub(monkeypatch, tmp_path) - hub = sys.modules["huggingface_hub"] - monkeypatch.setattr( - hub, "list_repo_files", lambda *args, **kwargs: ["meta.csv", "data/train.parquet", "README.md"], raising=False - ) +def test_snapshot_downloads_the_matching_files_and_returns_the_snapshot_root(monkeypatch, tmp_path): + calls = _fake_hub(monkeypatch, tmp_path, files=("meta.csv", "data/train.parquet", "README.md")) # asyncio.run gives the helper a running event loop, as download_async does in a build. root = asyncio.run(hub_snapshot("org/repo", _SHA, tmp_path, ("data/*.parquet",))) assert root == tmp_path / "snapshot" - assert (root / "data/train.parquet").is_file() - assert calls[0][1]["allow_patterns"] == ["data/train.parquet"] - with pytest.raises(TimeNetDownloadError, match="no files matching"): - asyncio.run(hub_snapshot("org/repo", _SHA, tmp_path, ("missing/*",))) + assert sorted(path.relative_to(root).as_posix() for path in root.rglob("*") if path.is_file()) == [ + "data/train.parquet" + ] + assert calls[0][1]["allow_patterns"] == ["data/*.parquet"] + + +class _Patterned(BaseHuggingFaceConnector): + HF_REPO = "org/repo" + SOURCE_REVISION = _SHA + SOURCE_FILES: ClassVar = {"questions": "qa.csv", "series": "data/*.parquet"} + + def convert(self, raw_refs): + return raw_refs + + +def test_a_pattern_part_resolves_to_the_directory_of_its_files(monkeypatch, tmp_path): + calls = _fake_hub(monkeypatch, tmp_path, files=("qa.csv", "data/a.parquet", "data/b.parquet", "README.md")) + + refs = asyncio.run(_Patterned().download_async(tmp_path / "cache")) + + assert refs == [ + HubFileRef(tmp_path / "snapshot" / "qa.csv", "questions"), + HubFileRef(tmp_path / "snapshot" / "data", "series"), + ] + assert calls[0][1]["allow_patterns"] == ["qa.csv", "data/*.parquet"] + + +def test_a_declared_file_or_pattern_without_a_match_fails(monkeypatch, tmp_path): + _fake_hub(monkeypatch, tmp_path, files=("qa.csv", "README.md")) + with pytest.raises(TimeNetDownloadError, match=rf"org/repo@{_SHA} has no data/\*\.parquet$"): + asyncio.run(_Patterned().download_async(tmp_path / "cache"))