add changes for running cohere wiki benchmark

Author: generall

datasets/datasets.json

@@ -354,5 +354,13 @@
     "type": "tar",
     "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/random_keywords_1m_vocab_10_no_filters.tgz",
     "path": "random-100-match-kw-small-vocab/random_keywords_1m_vocab_10_no_filters"
+  },
+  {
+    "name": "cohere-wiki-50m-test-only",
+    "vector_size": 768,
+    "distance": "cosine",
+    "type": "tar",
+    "link": "https://storage.googleapis.com/ann-filtered-benchmark/datasets/cohere-wiki-50m-test-only.tgz",
+    "path": "cohere-wiki-50m/cohere_wiki_50m"
   }
 ]

engine/clients/qdrant/config.py

@@ -1,3 +1,4 @@
 import os
 
 QDRANT_COLLECTION_NAME = os.getenv("QDRANT_COLLECTION_NAME", "benchmark")
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", None)

engine/clients/qdrant/configure.py

@@ -4,7 +4,7 @@
 from benchmark.dataset import Dataset
 from engine.base_client.configure import BaseConfigurator
 from engine.base_client.distances import Distance
-from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
+from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME, QDRANT_API_KEY
 
 
 class QdrantConfigurator(BaseConfigurator):
@@ -32,7 +32,7 @@ class QdrantConfigurator(BaseConfigurator):
     def __init__(self, host, collection_params: dict, connection_params: dict):
         super().__init__(host, collection_params, connection_params)
 
-        self.client = QdrantClient(host=host, **connection_params)
+        self.client = QdrantClient(url=host, api_key=QDRANT_API_KEY, **connection_params)
 
     def clean(self):
         self.client.delete_collection(collection_name=QDRANT_COLLECTION_NAME)

engine/clients/qdrant/search.py

@@ -8,7 +8,7 @@
 
 from dataset_reader.base_reader import Query
 from engine.base_client.search import BaseSearcher
-from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
+from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME, QDRANT_API_KEY
 from engine.clients.qdrant.parser import QdrantConditionParser
 
 
@@ -22,8 +22,9 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
         os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "true"
         os.environ["GRPC_POLL_STRATEGY"] = "epoll,poll"
         cls.client: QdrantClient = QdrantClient(
-            host,
+            url=host,
             prefer_grpc=True,
+            api_key=QDRANT_API_KEY,
             limits=httpx.Limits(max_connections=None, max_keepalive_connections=0),
             **connection_params,
         )

engine/clients/qdrant/upload.py

@@ -13,7 +13,7 @@
 
 from dataset_reader.base_reader import Record
 from engine.base_client.upload import BaseUploader
-from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
+from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME, QDRANT_API_KEY
 
 
 class QdrantUploader(BaseUploader):
@@ -24,7 +24,7 @@ class QdrantUploader(BaseUploader):
     def init_client(cls, host, distance, connection_params, upload_params):
         os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "true"
         os.environ["GRPC_POLL_STRATEGY"] = "epoll,poll"
-        cls.client = QdrantClient(host=host, prefer_grpc=True, **connection_params)
+        cls.client = QdrantClient(url=host, prefer_grpc=True, api_key=QDRANT_API_KEY, **connection_params)
         cls.upload_params = upload_params
 
     @classmethod

experiments/configurations/qdrant-single-node.json

@@ -39,7 +39,25 @@
         "config": {
           "hnsw_ef": 256,
           "quantization": {
-            "oversampling": 2.0
+            "oversampling": 2.0,
+            "rescore": true
+          }
+        }
+      }
+    ],
+    "upload_params": { "parallel": 16, "batch_size": 1024 }
+  },
+  {
+    "name": "qdrant-rescore-only",
+    "engine": "qdrant",
+    "connection_params": { "timeout": 30 },
+    "search_params": [
+      {
+        "parallel": 8,
+        "config": {
+          "hnsw_ef": 128,
+          "quantization": {
+            "rescore": true
           }
         }
       }