[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

benchmark/dataset.py

@@ -1,10 +1,9 @@
 import os
 import shutil
-from dataclasses import dataclass
-from typing import Optional
 import tarfile
-
 import urllib.request
+from dataclasses import dataclass
+from typing import Optional
 
 from benchmark import DATASETS_DIR
 from dataset_reader.ann_h5_reader import AnnH5Reader

dataset_reader/ann_h5_reader.py

@@ -1,10 +1,10 @@
-from typing import Iterator, Iterable, Any
+from typing import Any, Iterable, Iterator
 
 import h5py
 import numpy as np
 
 from benchmark import DATASETS_DIR
-from dataset_reader.base_reader import BaseReader, Record, Query
+from dataset_reader.base_reader import BaseReader, Query, Record
 
 
 class AnnH5Reader(BaseReader):
@@ -54,5 +54,3 @@ def read_data(self) -> Iterator[Record]:
 
     query = next(AnnH5Reader(test_path).read_queries())
     print(query)
-
-

dataset_reader/json_reader.py

@@ -1,12 +1,10 @@
+import json
 from pathlib import Path
 from typing import Iterator, List, Optional
 
-import json
-
 import numpy as np
 
-from dataset_reader.base_reader import BaseReader, Record, Query
-
+from dataset_reader.base_reader import BaseReader, Query, Record
 
 VECTORS_FILE = "vectors.jsonl"
 PAYLOADS_FILE = "payloads.jsonl"

engine/base_client/__init__.py

@@ -1,4 +1,4 @@
+from engine.base_client.client import BaseClient
 from engine.base_client.configure import BaseConfigurator
-from engine.base_client.upload import BaseUploader
 from engine.base_client.search import BaseSearcher
-from engine.base_client.client import BaseClient
+from engine.base_client.upload import BaseUploader

engine/base_client/client.py

@@ -2,8 +2,8 @@
 from datetime import datetime
 from typing import List
 
-from benchmark.dataset import Dataset
 from benchmark import ROOT_DIR
+from benchmark.dataset import Dataset
 from engine.base_client.configure import BaseConfigurator
 from engine.base_client.search import BaseSearcher
 from engine.base_client.upload import BaseUploader
@@ -48,10 +48,11 @@ def save_upload_results(self, dataset_name: str, results: dict):
     def run_experiment(self, dataset: Dataset):
         print("Experiment stage: Configure")
         execution_params = self.configurator.configure(
-            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))
+        reader = dataset.get_reader(execution_params.get("normalize", False))
         print("Experiment stage: Upload")
         upload_stats = self.uploader.upload(reader.read_data())
         self.save_upload_results(dataset.config.name, upload_stats)

engine/base_client/search.py

@@ -1,13 +1,12 @@
-import time
 import functools
+import time
 from multiprocessing import get_context
-from typing import List, Tuple, Iterable, Optional
+from typing import Iterable, List, Optional, Tuple
 
 import numpy as np
 
 from dataset_reader.base_reader import Query
 
-
 DEFAULT_TOP = 10
 
 
@@ -50,7 +49,8 @@ def _search_one(cls, query, top: Optional[int] = None):
         return precision, end - start
 
     def search_all(
-        self, queries: Iterable[Query],
+        self,
+        queries: Iterable[Query],
     ):
         start = time.perf_counter()
         parallel = self.search_params.pop("parallel", 1)

engine/base_client/upload.py

@@ -1,6 +1,6 @@
 import time
 from multiprocessing import get_context
-from typing import List, Optional, Iterable
+from typing import Iterable, List, Optional
 
 import tqdm
 
@@ -21,7 +21,10 @@ def __init__(self, host, connection_params, upload_params):
     def init_client(cls, host, connection_params: dict, upload_params: dict):
         raise NotImplementedError()
 
-    def upload(self, records: Iterable[Record],) -> dict:
+    def upload(
+        self,
+        records: Iterable[Record],
+    ) -> dict:
         latencies = []
         start = time.perf_counter()
         parallel = self.upload_params.pop("parallel", 1)

engine/clients/client_factory.py

@@ -7,14 +7,13 @@
     BaseSearcher,
     BaseUploader,
 )
-from engine.clients.qdrant import QdrantConfigurator, QdrantUploader, QdrantSearcher
+from engine.clients.milvus import MilvusConfigurator, MilvusSearcher, MilvusUploader
+from engine.clients.qdrant import QdrantConfigurator, QdrantSearcher, QdrantUploader
 from engine.clients.weaviate import (
     WeaviateConfigurator,
-    WeaviateUploader,
     WeaviateSearcher,
+    WeaviateUploader,
 )
-from engine.clients.milvus import MilvusConfigurator, MilvusUploader, MilvusSearcher
-
 
 ENGINE_CONFIGURATORS = {
     "qdrant": QdrantConfigurator,

engine/clients/milvus/__init__.py

@@ -1,3 +1,3 @@
 from engine.clients.milvus.configure import MilvusConfigurator
-from engine.clients.milvus.upload import MilvusUploader
 from engine.clients.milvus.search import MilvusSearcher
+from engine.clients.milvus.upload import MilvusUploader

engine/clients/milvus/configure.py

@@ -1,16 +1,15 @@
 import pymilvus.client.exceptions as milvus_exceptions
+from pymilvus import Collection, CollectionSchema, DataType, FieldSchema, connections
 from pymilvus.orm import utility
 
 from engine.base_client.configure import BaseConfigurator
 from engine.base_client.distances import Distance
 from engine.clients.milvus.config import (
     MILVUS_COLLECTION_NAME,
-    MILVUS_DEFAULT_PORT,
     MILVUS_DEFAULT_ALIAS,
+    MILVUS_DEFAULT_PORT,
 )
 
-from pymilvus import connections, DataType, FieldSchema, CollectionSchema, Collection
-
 
 class MilvusConfigurator(BaseConfigurator):
     DISTANCE_MAPPING = {
@@ -39,18 +38,29 @@ def clean(self):
             pass
 
     def recreate(
-        self, distance, vector_size, collection_params,
+        self,
+        distance,
+        vector_size,
+        collection_params,
     ):
-        idx = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True,)
+        idx = FieldSchema(
+            name="id",
+            dtype=DataType.INT64,
+            is_primary=True,
+        )
         vector = FieldSchema(
-            name="vector", dtype=DataType.FLOAT_VECTOR, dim=vector_size,
+            name="vector",
+            dtype=DataType.FLOAT_VECTOR,
+            dim=vector_size,
         )
         schema = CollectionSchema(
             fields=[idx, vector], description=MILVUS_COLLECTION_NAME
         )
 
         collection = Collection(
-            name=MILVUS_COLLECTION_NAME, schema=schema, using=MILVUS_DEFAULT_ALIAS,
+            name=MILVUS_COLLECTION_NAME,
+            schema=schema,
+            using=MILVUS_DEFAULT_ALIAS,
         )
 
         for index in collection.indexes: