Skip to content

Feat/integrate vector database#657

Open
Baduc81 wants to merge 4 commits into
OML-Team:mainfrom
Baduc81:feat/integrate_vector_database
Open

Feat/integrate vector database#657
Baduc81 wants to merge 4 commits into
OML-Team:mainfrom
Baduc81:feat/integrate_vector_database

Conversation

@Baduc81

@Baduc81 Baduc81 commented May 28, 2026

Copy link
Copy Markdown
  • I've checked contribution guide.

Summary

Closes #586.

Adds IVectorIndex — a common interface for vector index backends — and integrates it with RetrievalResults.from_index(). This lets users build a gallery index once and query it repeatedly across batches, without reprocessing the gallery each time.


Changes

oml/retrieval/vector_index.py (new file)

  • IVectorIndex — abstract base class with add(), search(), and __len__()
  • SklearnKNNIndex — exact search via sklearn.neighbors.NearestNeighbors
  • FaissIndex — exact/approximate search via faiss.IndexFlatL2 (swappable to IVF, HNSW, etc.)
  • QdrantIndex — in-memory Qdrant collection via qdrant-client, no server required

oml/retrieval/retrieval_results.py

  • Added RetrievalResults.from_index() classmethod
  • Automatically computes gt_ids when both datasets implement ILabeledDataset

oml/retrieval/__init__.py

  • Exported IVectorIndex, SklearnKNNIndex, FaissIndex, QdrantIndex

Usage

from oml.retrieval import RetrievalResults, FaissIndex

# Gallery is huge and fixed — process it once
dataset_gallery = ImageBaseDataset(galleries, transform=transform)
embeddings_gallery = inference(extractor, dataset_gallery, batch_size=4, num_workers=0)
index = FaissIndex(embeddings_gallery)

# Query in batches — gallery is never reprocessed
for queries in [queries1, queries2]:
    dataset_query = ImageBaseDataset(queries, transform=transform)
    embeddings_query = inference(extractor, dataset_query, batch_size=4, num_workers=0)
    rr = RetrievalResults.from_index(
        index=index,
        embeddings_query=embeddings_query,
        dataset_query=dataset_query,
        dataset_gallery=dataset_gallery,
    )
    rr.visualize_qg([0, 1], dataset_query=dataset_query, dataset_gallery=dataset_gallery, show=True)
    print(rr)

Notes

  • search() returns L2 distances (not squared), sorted ascending — consistent with the rest of OML
  • faiss, qdrant-client, and scikit-learn are treated as optional deps; ImportError is raised at runtime only if the relevant backend is instantiated without the package installed
  • QdrantIndex uses an in-memory client (:memory:) so no external server is needed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Integrate vector databases with RetrievalResults

1 participant