Merge pull request #28 from qdrant/pre-commit-ci-update-config

[pre-commit.ci] pre-commit suggestions

Author: kacperlukawski

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: check-added-large-files
 
   - repo: https://github.com/psf/black
-    rev: 22.6.0
+    rev: 22.8.0
     hooks:
       - id: black
         name: "Black: The uncompromising Python code formatter"

benchmark/config_read.py

@@ -2,7 +2,7 @@
 import json
 import os
 
-from benchmark import ROOT_DIR, DATASETS_DIR
+from benchmark import DATASETS_DIR, ROOT_DIR
 
 
 def read_engine_configs() -> dict:

engine/base_client/client.py

@@ -14,19 +14,19 @@
 
 class BaseClient:
     def __init__(
-            self,
-            name: str,  # name of the experiment
-            configurator: BaseConfigurator,
-            uploader: BaseUploader,
-            searchers: List[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.searchers = searchers
 
     def save_search_results(
-            self, dataset_name: str, results: dict, search_id: int, search_params: dict
+        self, dataset_name: str, results: dict, search_id: int, search_params: dict
     ):
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
@@ -38,7 +38,9 @@ def save_search_results(
                 json.dumps({"params": search_params, "results": results}, indent=2)
             )
 
-    def save_upload_results(self, dataset_name: str, results: dict, upload_params: dict):
+    def save_upload_results(
+        self, dataset_name: str, results: dict, upload_params: dict
+    ):
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%d-%H-%M-%S")
         experiments_file = f"{self.name}-{dataset_name}-upload-{timestamp}.json"
@@ -51,8 +53,8 @@ def save_upload_results(self, dataset_name: str, results: dict, upload_params: d
 
     def run_experiment(self, dataset: Dataset, skip_upload: bool = False):
         execution_params = self.configurator.execution_params(
-            distance=dataset.config.distance,
-            vector_size=dataset.config.vector_size)
+            distance=dataset.config.distance, vector_size=dataset.config.vector_size
+        )
 
         reader = dataset.get_reader(execution_params.get("normalize", False))
 
@@ -65,18 +67,23 @@ def run_experiment(self, dataset: Dataset, skip_upload: bool = False):
 
             print("Experiment stage: Upload")
             upload_stats = self.uploader.upload(
-                distance=dataset.config.distance,
-                records=reader.read_data()
+                distance=dataset.config.distance, records=reader.read_data()
+            )
+            self.save_upload_results(
+                dataset.config.name,
+                upload_stats,
+                upload_params={
+                    **self.uploader.upload_params,
+                    **self.configurator.collection_params,
+                },
             )
-            self.save_upload_results(dataset.config.name, upload_stats, upload_params={
-                **self.uploader.upload_params,
-                **self.configurator.collection_params
-            })
 
         print("Experiment stage: Search")
         for search_id, searcher in enumerate(self.searchers):
             search_params = {**searcher.search_params}
-            search_stats = searcher.search_all(dataset.config.distance, reader.read_queries())
+            search_stats = searcher.search_all(
+                dataset.config.distance, reader.read_queries()
+            )
             self.save_search_results(
                 dataset.config.name, search_stats, search_id, search_params
             )

engine/base_client/search.py

@@ -20,7 +20,9 @@ def __init__(self, host, connection_params, search_params):
         self.search_params = search_params
 
     @classmethod
-    def init_client(cls, host: str, distance, connection_params: dict, search_params: dict):
+    def init_client(
+        cls, host: str, distance, connection_params: dict, search_params: dict
+    ):
         raise NotImplementedError()
 
     @classmethod
@@ -63,20 +65,29 @@ def search_all(
         top = self.search_params.pop("top", None)
 
         # setup_search may require initialized client
-        self.init_client(self.host, distance, self.connection_params, self.search_params)
+        self.init_client(
+            self.host, distance, self.connection_params, self.search_params
+        )
         self.setup_search()
 
         search_one = functools.partial(self.__class__._search_one, top=top)
 
         if parallel == 1:
-            precisions, latencies = list(zip(*[search_one(query) for query in tqdm.tqdm(queries)]))
+            precisions, latencies = list(
+                zip(*[search_one(query) for query in tqdm.tqdm(queries)])
+            )
         else:
             ctx = get_context(self.get_mp_start_method())
 
             with ctx.Pool(
                 processes=parallel,
                 initializer=self.__class__.init_client,
-                initargs=(self.host, distance, self.connection_params, self.search_params),
+                initargs=(
+                    self.host,
+                    distance,
+                    self.connection_params,
+                    self.search_params,
+                ),
             ) as pool:
                 precisions, latencies = list(
                     zip(*pool.imap_unordered(search_one, iterable=tqdm.tqdm(queries)))
@@ -101,4 +112,4 @@ def setup_search(self):
         pass
 
     def post_search(self):
-        pass
+        pass

engine/base_client/upload.py

@@ -25,31 +25,40 @@ def init_client(cls, host, distance, connection_params: dict, upload_params: dic
         raise NotImplementedError()
 
     def upload(
-            self,
-            distance,
-            records: Iterable[Record],
+        self,
+        distance,
+        records: Iterable[Record],
     ) -> dict:
         latencies = []
         start = time.perf_counter()
         parallel = self.upload_params.pop("parallel", 1)
         batch_size = self.upload_params.pop("batch_size", 64)
 
-        self.init_client(self.host, distance, self.connection_params, self.upload_params)
+        self.init_client(
+            self.host, distance, self.connection_params, self.upload_params
+        )
 
         if parallel == 1:
             for batch in iter_batches(tqdm.tqdm(records), batch_size):
                 latencies.append(self._upload_batch(batch))
         else:
             ctx = get_context(self.get_mp_start_method())
             with ctx.Pool(
-                    processes=int(parallel),
-                    initializer=self.__class__.init_client,
-                    initargs=(self.host, distance, self.connection_params, self.upload_params),
+                processes=int(parallel),
+                initializer=self.__class__.init_client,
+                initargs=(
+                    self.host,
+                    distance,
+                    self.connection_params,
+                    self.upload_params,
+                ),
             ) as pool:
-                latencies = list(pool.imap(
-                    self.__class__._upload_batch,
-                    iter_batches(tqdm.tqdm(records), batch_size),
-                ))
+                latencies = list(
+                    pool.imap(
+                        self.__class__._upload_batch,
+                        iter_batches(tqdm.tqdm(records), batch_size),
+                    )
+                )
 
         upload_time = time.perf_counter() - start
 
@@ -68,7 +77,7 @@ def upload(
 
     @classmethod
     def _upload_batch(
-            cls, batch: Tuple[List[int], List[list], List[Optional[dict]]]
+        cls, batch: Tuple[List[int], List[list], List[Optional[dict]]]
     ) -> float:
         ids, vectors, metadata = batch
         start = time.perf_counter()
@@ -81,6 +90,6 @@ def post_upload(cls, distance):
 
     @classmethod
     def upload_batch(
-            cls, ids: List[int], vectors: List[list], metadata: List[Optional[dict]]
+        cls, ids: List[int], vectors: List[list], metadata: List[Optional[dict]]
     ):
         raise NotImplementedError()

engine/clients/elasticsearch/configure.py

@@ -36,7 +36,9 @@ def __init__(self, host, collection_params: dict, connection_params: dict):
 
     def clean(self):
         try:
-            self.client.indices.delete(index=ELASTIC_INDEX, timeout="5m", master_timeout="5m")
+            self.client.indices.delete(
+                index=ELASTIC_INDEX, timeout="5m", master_timeout="5m"
+            )
         except NotFoundError:
             pass
 

engine/clients/elasticsearch/search.py

@@ -1,6 +1,6 @@
+import multiprocessing as mp
 import uuid
 from typing import List, Tuple
-import multiprocessing as mp
 
 from elasticsearch import Elasticsearch
 
@@ -14,7 +14,6 @@
 
 
 class ClosableElastic(Elasticsearch):
-
     def __del__(self):
         self.close()
 
@@ -25,7 +24,7 @@ class ElasticSearcher(BaseSearcher):
 
     @classmethod
     def get_mp_start_method(cls):
-        return 'forkserver' if 'forkserver' in mp.get_all_start_methods() else 'spawn'
+        return "forkserver" if "forkserver" in mp.get_all_start_methods() else "spawn"
 
     @classmethod
     def init_client(cls, host, distance, connection_params: dict, search_params: dict):

engine/clients/elasticsearch/upload.py

@@ -1,5 +1,5 @@
-import uuid
 import multiprocessing as mp
+import uuid
 from typing import List, Optional
 
 from elasticsearch import Elasticsearch
@@ -14,7 +14,6 @@
 
 
 class ClosableElastic(Elasticsearch):
-
     def __del__(self):
         self.close()
 
@@ -25,7 +24,7 @@ class ElasticUploader(BaseUploader):
 
     @classmethod
     def get_mp_start_method(cls):
-        return 'forkserver' if 'forkserver' in mp.get_all_start_methods() else 'spawn'
+        return "forkserver" if "forkserver" in mp.get_all_start_methods() else "spawn"
 
     @classmethod
     def init_client(cls, host, distance, connection_params, upload_params):

engine/clients/milvus/config.py

@@ -10,4 +10,4 @@
     # Milvus does not support cosine. Cosine is equal to IP of normalized vectors
     Distance.COSINE: "IP"
     # Jaccard, Tanimoto, Hamming distance, Superstructure and Substructure are also available
-}
+}

engine/clients/milvus/configure.py

@@ -18,7 +18,6 @@
 
 
 class MilvusConfigurator(BaseConfigurator):
-
     def __init__(self, host, collection_params: dict, connection_params: dict):
         super().__init__(host, collection_params, connection_params)
         self.client = connections.connect(