first running benchmark

Author: generall

benchmark/dataset.py

@@ -53,7 +53,7 @@ def download(self):
 
     def get_reader(self) -> BaseReader:
         reader_class = READER_TYPE[self.config.type]
-        return reader_class(self.config.path)
+        return reader_class(DATASETS_DIR / self.config.path)
 
 
 if __name__ == "__main__":

benchmark/settings.py

@@ -1,6 +1,5 @@
 import os.path
-
+from pathlib import Path
 
 CODE_DIR = os.path.dirname(__file__)
-ROOT_DIR = os.path.dirname(CODE_DIR)
-DATASETS_DIR = os.path.join(ROOT_DIR, "datasets")
+ROOT_DIR = Path(os.path.dirname(CODE_DIR))

dataset_reader/json_reader.py

@@ -1,34 +1,66 @@
-from typing import Iterator
+from pathlib import Path
+from typing import Iterator, List, Optional
 
 import json
 
 from dataset_reader.base_reader import BaseReader, Record, Query
 
 
+VECTORS_FILE = 'vectors.jsonl'
+PAYLOADS_FILE = 'payloads.jsonl'
+QUERIES_FILE = 'queries.jsonl'
+NEIGHBOURS_FILE = 'neighbours.jsonl'
+
+
 class JSONReader(BaseReader):
-    def __init__(self, path):
+    def __init__(self, path: Path):
         self.path = path
 
-    def read_queries(self) -> Iterator[Query]:
-        with open(self.path, "r") as json_fp:
+    def read_payloads(self) -> Iterator[dict]:
+        if not (self.path / PAYLOADS_FILE).exists():
+            while True:
+                yield {}
+        with open(self.path / PAYLOADS_FILE, "r") as json_fp:
             for json_line in json_fp:
                 line = json.loads(json_line)
-                yield Query(
-                    vector=line, meta_conditions=None, expected_result=None,
-                )
+                yield line
 
-    def read_data(self) -> Iterator[Record]:
-        with open(self.path, "r") as json_fp:
-            for idx, json_line in enumerate(json_fp):
+    def read_vectors(self) -> Iterator[List[float]]:
+        with open(self.path / VECTORS_FILE, "r") as json_fp:
+            for json_line in json_fp:
+                line = json.loads(json_line)
+                yield line
+
+    def read_neighbours(self) -> Iterator[Optional[List[int]]]:
+        if not (self.path / NEIGHBOURS_FILE).exists():
+            while True:
+                yield None
+
+        with open(self.path / NEIGHBOURS_FILE, "r") as json_fp:
+            for json_line in json_fp:
+                line = json.loads(json_line)
+                yield line
+
+    def read_query_vectors(self) -> Iterator[List[float]]:
+        with open(self.path / QUERIES_FILE, "r") as json_fp:
+            for json_line in json_fp:
                 line = json.loads(json_line)
-                yield Record(id=idx, vector=line, metadata=None)
+                yield line
+
+    def read_queries(self) -> Iterator[Query]:
+        for idx, (vector, neighbours) in enumerate(zip(self.read_query_vectors(), self.read_neighbours())):
+            # ToDo: add meta_conditions
+            yield Query(vector=vector, meta_conditions=None, expected_result=neighbours)
+
+    def read_data(self) -> Iterator[Record]:
+        for idx, (vector, payload) in enumerate(zip(self.read_vectors(), self.read_payloads())):
+            yield Record(id=idx, vector=vector, metadata=payload)
 
 
 if __name__ == "__main__":
-    import os
-    from benchmark.settings import DATASET_DIR
+    from benchmark import DATASETS_DIR
 
-    test_path = os.path.join(DATASET_DIR, "random-100", "vectors.jsonl")
+    test_path = DATASETS_DIR / "random-100"
     record = next(JSONReader(test_path).read_data())
     print(record, end="\n\n")
 

datasets/datasets.json

@@ -2,23 +2,23 @@
   {
     "name": "glove-25-angular",
     "vector_size": 25,
-    "distance": "Cosine",
+    "distance": "cosine",
     "type": "h5",
     "path": "glove-25-angular/glove-25-angular.hdf5",
     "link": "http://ann-benchmarks.com/glove-25-angular.hdf5"
   },
   {
     "name": "glove-100-angular",
     "vector_size": 100,
-    "distance": "Cosine",
+    "distance": "cosine",
     "type": "h5",
     "path": "glove-100-angular/glove-100-angular.hdf5",
     "link": "http://ann-benchmarks.com/glove-100-angular.hdf5"
   },
   {
     "name": "random-100",
     "vector_size": 100,
-    "distance": "Cosine",
+    "distance": "cosine",
     "type": "jsonl",
     "path": "random-100/"
   }

engine/base_client/client.py

@@ -1,30 +1,47 @@
 import json
 from datetime import datetime
+from typing import List
 
 from benchmark.dataset import Dataset
+from benchmark.settings import ROOT_DIR
 from engine.base_client.configure import BaseConfigurator
 from engine.base_client.search import BaseSearcher
 from engine.base_client.upload import BaseUploader
 
+RESULTS_DIR = ROOT_DIR / "results"
+RESULTS_DIR.mkdir(exist_ok=True)
+
 
 class BaseClient:
     def __init__(
-        self,
-        name: str,  # name of the experiment
-        configurator: BaseConfigurator,
-        uploader: BaseUploader,
-        searcher: BaseSearcher,
+            self,
+            name: str,  # name of the experiment
+            configurator: BaseConfigurator,
+            uploader: BaseUploader,
+            searchers: List[BaseSearcher],
     ):
         self.name = name
         self.configurator = configurator
         self.uploader = uploader
-        self.searcher = searcher
+        self.searchers = searchers
 
-    def save_experiment_results(self, dataset_name: str, results: dict):
+    def save_search_results(self, dataset_name: str, results: dict, search_id: int, search_params: dict):
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
-        experiments_file = f"{self.name}-{dataset_name}-{timestamp}.json"
-        with open(experiments_file, "w") as out:
+        experiments_file = f"{self.name}-{dataset_name}-search-{search_id}-{timestamp}.json"
+        with open(RESULTS_DIR / experiments_file, "w") as out:
+            out.write(
+                json.dumps({
+                    "params": search_params,
+                    "results": results
+                }, indent=2)
+            )
+
+    def save_upload_results(self, dataset_name: str, results: dict):
+        now = datetime.now()
+        timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
+        experiments_file = f"{self.name}-{dataset_name}-upload-{timestamp}.json"
+        with open(RESULTS_DIR / experiments_file, "w") as out:
             out.write(
                 json.dumps(results, indent=2)
             )
@@ -37,12 +54,14 @@ def run_experiment(self, dataset: Dataset):
 
         reader = dataset.get_reader()
         upload_stats = self.uploader.upload(reader.read_data())
-        search_stats = self.searcher.search_all(reader.read_queries())
-
-        self.save_experiment_results(
-            dataset.config.name,
-            {
-                "upload": upload_stats,
-                "search": search_stats
-            }
-        )
+        self.save_upload_results(dataset.config.name, upload_stats)
+
+        for search_id, searcher in enumerate(self.searchers):
+            search_params = {**searcher.search_params}
+            search_stats = searcher.search_all(reader.read_queries())
+            self.save_search_results(
+                dataset.config.name,
+                search_stats,
+                search_id,
+                search_params
+            )

engine/base_client/configure.py

@@ -1,10 +1,7 @@
-import json
 import time
-from pathlib import Path
 
 
 class BaseConfigurator:
-    DEFAULT_CONFIG_PATH = Path("default.json")
     DISTANCE_MAPPING = {}
 
     def __init__(self, host, collection_params: dict, connection_params: dict):
@@ -23,8 +20,3 @@ def configure(self, distance, vector_size):
         start = time.perf_counter()
         self.recreate(distance, vector_size, self.collection_params)
         return time.perf_counter() - start
-
-    @classmethod
-    def read_default_config(cls):
-        with open(cls.DEFAULT_CONFIG_PATH, "r") as fp:
-            return json.load(fp)

engine/base_client/search.py

@@ -58,7 +58,7 @@ def search_all(
 
         if parallel == 1:
             self.init_client(self.host, self.connection_params, self.search_params)
-            precisions, latencies = list(zip([search_one(query) for query in queries]))
+            precisions, latencies = list(zip(*[search_one(query) for query in queries]))
         else:
             ctx = get_context(self.MP_CONTEXT)
 

engine/clients/client_factory.py

@@ -1,4 +1,5 @@
 from abc import ABC
+from typing import List, Type
 
 from engine.base_client.client import BaseClient
 from engine.base_client.configure import BaseConfigurator
@@ -44,19 +45,24 @@ def _create_uploader(self, experiment) -> BaseUploader:
         )
         return engine_uploader
 
-    def _create_searcher(self, experiment) -> BaseSearcher:
-        engine_searcher_class = ENGINE_SEARCHERS[experiment["engine"]]
-        engine_searcher = engine_searcher_class(
-            self.host,
-            experiment.get("connection_params", {}),
-            experiment.get("search_params", {}),
-        )
-        return engine_searcher
+    def _create_searchers(self, experiment) -> List[BaseSearcher]:
+        engine_searcher_class: Type[BaseSearcher] = ENGINE_SEARCHERS[experiment["engine"]]
+
+        engine_searchers = [
+            engine_searcher_class(
+                self.host,
+                connection_params=experiment.get("connection_params", {}),
+                search_params=search_params,
+            )
+            for search_params in experiment.get("search_params", [{}])
+        ]
+
+        return engine_searchers
 
     def build_client(self, experiment):
         return BaseClient(
             name=experiment["name"],
             configurator=self._create_configurator(experiment),
             uploader=self._create_uploader(experiment),
-            searcher=self._create_searcher(experiment),
+            searchers=self._create_searchers(experiment),
         )

engine/clients/qdrant/search.py

@@ -22,11 +22,12 @@ def conditions_to_filter(cls, _meta_conditions) -> Optional[rest.Filter]:
         return None
 
     @classmethod
-    def search_one(cls, vector, meta_conditions) -> List[Tuple[int, float]]:
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
         res = cls.client.search(
             collection_name=QDRANT_COLLECTION_NAME,
             query_vector=vector,
             query_filter=cls.conditions_to_filter(meta_conditions),
+            limit=top,
             **cls.search_params
         )