Spaces:
Running
Running
File size: 3,668 Bytes
5959a78 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
Compare similarity-search speed: FAISS (exact, flat) vs ChromaDB (HNSW, approximate).
python scripts/benchmark_search.py # 2,000 and 20,000 random vectors
python scripts/benchmark_search.py --sizes 5000 50000 --queries 300
Measures search only (no embedding model needed), using random 384-d unit vectors like MiniLM's.
Also reports recall: how many of Chroma's top-k results match the exact FAISS answer.
Numbers depend on your CPU - run it on your own machine before quoting results.
"""
import argparse
import os
import sys
import tempfile
import time
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from backend.services.faiss_index import FaissIndex # noqa: E402
def run(n: int, dim: int, queries: int, top_k: int, seed: int = 0) -> dict:
import chromadb
from chromadb.config import Settings as ChromaSettings
rng = np.random.default_rng(seed)
data = rng.normal(size=(n, dim)).astype("float32")
data /= np.linalg.norm(data, axis=1, keepdims=True)
qs = rng.normal(size=(queries, dim)).astype("float32")
qs /= np.linalg.norm(qs, axis=1, keepdims=True)
ids = [f"c{i}" for i in range(n)]
t = time.perf_counter()
faiss_index = FaissIndex()
faiss_index.add(ids, data, ["doc"] * n)
faiss_build = time.perf_counter() - t
with tempfile.TemporaryDirectory() as tmp:
client = chromadb.PersistentClient(path=tmp, settings=ChromaSettings(anonymized_telemetry=False))
collection = client.get_or_create_collection("bench", metadata={"hnsw:space": "cosine"},
embedding_function=None)
t = time.perf_counter()
for start in range(0, n, 5000):
end = min(start + 5000, n)
collection.add(ids=ids[start:end], embeddings=data[start:end].tolist())
chroma_build = time.perf_counter() - t
for q in qs[:10]: # warm-up
faiss_index.search(q, top_k)
collection.query(query_embeddings=[q.tolist()], n_results=top_k)
t = time.perf_counter()
exact = [{c for c, _ in faiss_index.search(q, top_k)} for q in qs]
faiss_ms = (time.perf_counter() - t) * 1000 / queries
t = time.perf_counter()
approx = [set(collection.query(query_embeddings=[q.tolist()], n_results=top_k)["ids"][0]) for q in qs]
chroma_ms = (time.perf_counter() - t) * 1000 / queries
recall = float(np.mean([len(a & e) / top_k for a, e in zip(approx, exact)]))
return {"n": n, "faiss_ms": faiss_ms, "chroma_ms": chroma_ms, "recall": recall,
"faiss_build": faiss_build, "chroma_build": chroma_build}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--sizes", type=int, nargs="+", default=[2000, 20000])
parser.add_argument("--dim", type=int, default=384)
parser.add_argument("--queries", type=int, default=200)
parser.add_argument("--top-k", type=int, default=5)
args = parser.parse_args()
print(f"{'vectors':>9} | {'FAISS ms/query':>15} | {'Chroma ms/query':>16} | {'speed-up':>8} | "
f"{'Chroma recall@k':>15} | {'build FAISS/Chroma (s)':>23}")
print("-" * 100)
for n in args.sizes:
r = run(n, args.dim, args.queries, args.top_k)
print(f"{r['n']:>9,} | {r['faiss_ms']:>15.3f} | {r['chroma_ms']:>16.3f} | "
f"{r['chroma_ms'] / r['faiss_ms']:>7.1f}x | {r['recall']:>14.1%} | "
f"{r['faiss_build']:>10.2f} / {r['chroma_build']:<10.2f}")
if __name__ == "__main__":
main()
|