backend/astroparse_api/bootstrap.py

"""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))