File size: 1,628 Bytes
2d066db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """Corpus bootstrap — download from HF dataset if not already present."""
from __future__ import annotations
import shutil
import tempfile
from pathlib import Path
from huggingface_hub import snapshot_download # noqa: F401 — re-exported so tests can patch
def ensure_corpus(corpus_dir: Path, faiss_path: Path, repo_id: str) -> None:
"""Ensure corpus parquet shards and FAISS index exist locally.
If both *corpus_dir* (directory) and *faiss_path* (file) already exist,
return immediately without any network access. Otherwise download the
full HF dataset repo into a temporary staging directory and move the
expected artefacts into place:
<repo>/astroparse_fp16.faiss → faiss_path
<repo>/corpus/ → corpus_dir
"""
if corpus_dir.is_dir() and faiss_path.is_file():
return
# Import the module-level name so monkeypatch can replace it
import astroparse_api.bootstrap as _self
_snapshot_download = _self.snapshot_download
with tempfile.TemporaryDirectory() as staging:
staging_path = Path(staging)
_snapshot_download(
repo_id=repo_id,
repo_type="dataset",
local_dir=staging_path,
)
# Move FAISS index
src_faiss = staging_path / "astroparse_fp16.faiss"
faiss_path.parent.mkdir(parents=True, exist_ok=True)
shutil.move(str(src_faiss), str(faiss_path))
# Move corpus directory
src_corpus = staging_path / "corpus"
corpus_dir.parent.mkdir(parents=True, exist_ok=True)
shutil.move(str(src_corpus), str(corpus_dir))
|