diff --git a/requirements.txt b/requirements.txt index 3b51e88..5ba2c8a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,4 @@ torch==2.2.2 tqdm==4.65.0 transformers==4.37.0 chromadb==0.4.24 +qdrant-client diff --git a/tests/benchtest.py b/tests/benchtest.py index 676bc34..0ffb4a8 100644 --- a/tests/benchtest.py +++ b/tests/benchtest.py @@ -12,6 +12,11 @@ import chromadb from chromadb.utils import embedding_functions +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, VectorParams, PointStruct +from sentence_transformers import SentenceTransformer + + def main(queries, corpuss, top_k, token_counts) -> pd.DataFrame: """Run the VLite benchmark. @@ -144,6 +149,87 @@ def main(queries, corpuss, top_k, token_counts) -> pd.DataFrame: "stddev_time": np.std(times), } ) + + ################################################# + # Qdrant # + ################################################# + print("Begin Qdrant benchmark.") + print("Adding vectors to Qdrant instance...") + t0 = time.time() + + qdrant_client = QdrantClient( + ":memory:" # Use in-memory storage for benchmarking + ) + qdrant_client.recreate_collection( + collection_name="my_collection", + vectors_config=VectorParams(size=1024, distance=Distance.COSINE), + ) + + model = SentenceTransformer('mixedbread-ai/mxbai-embed-large-v1') + vectors = model.encode(corpus).tolist() + + print("Adding vectors to Qdrant instance...") + + try: + qdrant_client.upsert( + collection_name="my_collection", + points=[ + PointStruct( + id=idx, + vector=model.encode(vector).tolist(), + payload={"text": corpus[idx]} + ) + for idx, vector in enumerate(corpus) + ] + ) + except Exception as e: + print(e) + print("Failed to upsert vectors to Qdrant instance.") + t0 = time.time() + + t1 = time.time() + print(f"Took {t1 - t0:.3f}s to add vectors.") + indexing_times.append( + { + "num_tokens": token_count, + "lib": "Qdrant", + "num_embeddings": len(vectors), + "indexing_time": t1 - t0, + } + ) + + print("Starting Qdrant trials...") + times = [] + + for i in range(len(query)): + t0 = time.time() + query_vector = model.encode(query[i]).tolist() + try: + hits = qdrant_client.search( + collection_name="my_collection", + query_vector=query_vector, + limit=top_k + ) + # print(f"Top {top_k} results: {[hit.payload['text'] for hit in hits]}") + t1 = time.time() + except Exception as e: + print(e) + print("Failed to query Qdrant instance.") + t0 = time.time() + times.append(t1 - t0) + + results.append( + { + "num_embeddings": len(vectors), + "lib": "Qdrant", + "k": top_k, + "avg_time": np.mean(times), + "stddev_time": np.std(times), + } + ) + + print(json.dumps(results[-1], indent=2)) + print("Done Qdrant benchmark.") results_df = pd.DataFrame(results) indexing_times_df = pd.DataFrame(indexing_times)