backend/scripts/embed_dataset.py

"""Generate embeddings for a dataset's images with the *active* embedder, with terminal progress.

Use this instead of the UI "Embed" button for large datasets / the real (CPU) HF model, where the
synchronous HTTP endpoint would time out. Idempotent: skips pictures that already have an embedding
for the active model, so it's safe to re-run / resume.

Usage:
    python -m scripts.embed_dataset --dataset-id 3        # one dataset
    python -m scripts.embed_dataset --all                 # every dataset
    python -m scripts.embed_dataset --all --limit 5       # quick smoke (first 5 pictures each)
"""
from __future__ import annotations

import argparse
import time

from sqlalchemy import select

from app.db import SessionLocal, engine
from app.ml import get_embedder
from app.models import Base, Dataset, Picture
from app.services.datasets import _known_ids, _picture_ids, _unknown_ids
from app.services.images import embed_picture

BATCH = 100


def run(dataset_id: int | None, do_all: bool, limit: int | None) -> None:
    Base.metadata.create_all(bind=engine)
    db = SessionLocal()
    try:
        embedder = get_embedder()
        print(f"Active embedder: {embedder.name}/{embedder.version} (dim {embedder.dim})")

        if do_all:
            datasets = db.execute(select(Dataset).order_by(Dataset.id)).scalars().all()
        elif dataset_id is not None:
            d = db.get(Dataset, dataset_id)
            datasets = [d] if d else []
            if not d:
                print(f"Dataset #{dataset_id} not found.")
        else:
            print("Pass --dataset-id N or --all.")
            return

        for ds in datasets:
            pic_ids = _picture_ids(db, _known_ids(db, ds.id), _unknown_ids(db, ds.id))
            if limit:
                pic_ids = pic_ids[:limit]
            print(f"\nDataset #{ds.id} {ds.name!r}: {len(pic_ids)} picture(s) to consider")
            embedded = skipped = errors = 0
            start = time.time()
            for i, pid in enumerate(pic_ids, 1):
                pic = db.get(Picture, pid)
                if pic is None:
                    continue
                try:
                    if embed_picture(db, pic, skip_if_exists=True):
                        embedded += 1
                    else:
                        skipped += 1
                except Exception as exc:  # noqa: BLE001
                    errors += 1
                    print(f"  ! picture {pid}: {exc}")
                if i % BATCH == 0:
                    db.commit()
                    rate = i / max(time.time() - start, 1e-6)
                    print(f"  {i}/{len(pic_ids)}  embedded={embedded} skipped={skipped} "
                          f"errors={errors}  ({rate:.1f} img/s)")
            db.commit()
            print(f"  done: embedded={embedded} skipped={skipped} errors={errors} "
                  f"in {time.time() - start:.0f}s")
    finally:
        db.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="Embed a dataset's images with the active embedder.")
    parser.add_argument("--dataset-id", type=int, default=None)
    parser.add_argument("--all", action="store_true", help="Embed every dataset")
    parser.add_argument("--limit", type=int, default=None, help="Only the first N pictures (smoke test)")
    args = parser.parse_args()
    run(args.dataset_id, args.all, args.limit)


if __name__ == "__main__":
    main()