| """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 |
|
|
|
|
| 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 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, |
| ) |
|
|
| |
| src_faiss = staging_path / "astroparse_fp16.faiss" |
| faiss_path.parent.mkdir(parents=True, exist_ok=True) |
| shutil.move(str(src_faiss), str(faiss_path)) |
|
|
| |
| src_corpus = staging_path / "corpus" |
| corpus_dir.parent.mkdir(parents=True, exist_ok=True) |
| shutil.move(str(src_corpus), str(corpus_dir)) |
|
|
