src/storage.py

"""Download embedding files from Hugging Face Hub Dataset at startup if not present locally."""

import os
from pathlib import Path

DATA_DIR = Path(__file__).resolve().parent.parent / "data"


def _is_valid_npy(path: Path) -> bool:
    """Check if file is a real numpy array (not a git LFS pointer)."""
    try:
        with open(path, "rb") as f:
            return f.read(6) == b"\x93NUMPY"
    except Exception:
        return False


def ensure_embeddings(filenames: list[str]) -> None:
    """Download any missing or corrupted embedding files from HF Hub Dataset."""
    missing = [f for f in filenames if not (DATA_DIR / f).exists() or not _is_valid_npy(DATA_DIR / f)]
    if not missing:
        print("Embeddings already present locally, skipping download.")
        return

    from huggingface_hub import hf_hub_download

    repo_id = os.environ.get("HF_DATASET_REPO")
    token = os.environ.get("HF_TOKEN")

    if not repo_id:
        raise RuntimeError(
            "Missing HF_DATASET_REPO env var. Set it to your HF dataset repo, e.g. 'yourname/nichefind-data'"
        )

    print(f"Downloading {len(missing)} embedding file(s) from HF dataset '{repo_id}' ...")
    for filename in missing:
        dest = DATA_DIR / filename
        print(f"  ↓ {filename}")
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            repo_type="dataset",
            token=token,
            local_dir=str(DATA_DIR),
        )
        print(f"  ✓ {filename} ({dest.stat().st_size / 1e6:.1f} MB)")