start weaviate

generall · generall · commit 5faa4cbf823d · 2022-07-27T17:42:26.000+02:00
diff --git a/benchmark/__init__.py b/benchmark/__init__.py
@@ -1,6 +1,10 @@
+import os
 from pathlib import Path
 
 # Base directory point to the main directory of the project, so all the data
 # loaded from files can refer to it as a root directory
+
 BASE_DIRECTORY = Path(__file__).parent.parent
 DATASETS_DIR = BASE_DIRECTORY / "datasets"
+CODE_DIR = os.path.dirname(__file__)
+ROOT_DIR = Path(os.path.dirname(CODE_DIR))
diff --git a/benchmark/settings.py b/benchmark/settings.py
diff --git a/dataset_reader/ann_h5_reader.py b/dataset_reader/ann_h5_reader.py
@@ -2,6 +2,7 @@
 
 import h5py
 
+from benchmark import DATASETS_DIR
 from dataset_reader.base_reader import BaseReader, Record, Query
 
 
@@ -33,7 +34,6 @@ def read_data(self) -> Iterator[Record]:
 
 if __name__ == '__main__':
     import os
-    from benchmark.settings import DATASET_DIR
 
     # h5py file 4 keys:
     # `train` - float vectors (num vectors 1183514)
@@ -42,7 +42,7 @@ def read_data(self) -> Iterator[Record]:
     # contains info about 100 nearest neighbors)
     # `distances` - float - distances for nearest neighbors for test vectors
 
-    test_path = os.path.join(DATASET_DIR, 'glove-100-angular', 'glove-100-angular.hdf5')
+    test_path = os.path.join(DATASETS_DIR, 'glove-100-angular', 'glove-100-angular.hdf5')
     record = next(AnnH5Reader(test_path).read_data())
     print(record, end='\n\n')
 
diff --git a/engine/base_client/client.py b/engine/base_client/client.py
@@ -3,7 +3,7 @@
 from typing import List
 
 from benchmark.dataset import Dataset
-from benchmark.settings import ROOT_DIR
+from benchmark import ROOT_DIR
 from engine.base_client.configure import BaseConfigurator
 from engine.base_client.search import BaseSearcher
 from engine.base_client.upload import BaseUploader
diff --git a/engine/base_client/search.py b/engine/base_client/search.py
@@ -54,6 +54,8 @@ def search_all(
         parallel = self.search_params.pop("parallel", 1)
         top = self.search_params.pop("top", None)
 
+        self.setup_search()
+
         search_one = functools.partial(self.__class__._search_one, top=top)
 
         if parallel == 1:
@@ -88,3 +90,6 @@ def search_all(
 
     def set_process_start_method(self, start_method):
         self.MP_CONTEXT = start_method
+
+    def setup_search(self):
+        pass
diff --git a/engine/clients/client_factory.py b/engine/clients/client_factory.py
@@ -9,7 +9,6 @@
 from engine.clients.qdrant.search import QdrantSearcher
 from engine.clients.qdrant.upload import QdrantUploader
 
-
 ENGINE_CONFIGURATORS = {
     "qdrant": QdrantConfigurator,
 }
@@ -31,17 +30,17 @@ def _create_configurator(self, experiment) -> BaseConfigurator:
         engine_configurator_class = ENGINE_CONFIGURATORS[experiment["engine"]]
         engine_configurator = engine_configurator_class(
             self.host,
-            experiment.get("collection_params", {}),
-            experiment.get("connection_params", {}),
+            collection_params={**experiment.get("collection_params", {})},
+            connection_params={**experiment.get("connection_params", {})},
         )
         return engine_configurator
 
     def _create_uploader(self, experiment) -> BaseUploader:
         engine_uploader_class = ENGINE_UPLOADERS[experiment["engine"]]
         engine_uploader = engine_uploader_class(
             self.host,
-            experiment.get("connection_params", {}),
-            experiment.get("upload_params", {}),
+            connection_params={**experiment.get("connection_params", {})},
+            upload_params={**experiment.get("upload_params", {})},
         )
         return engine_uploader
 
@@ -51,7 +50,7 @@ def _create_searchers(self, experiment) -> List[BaseSearcher]:
         engine_searchers = [
             engine_searcher_class(
                 self.host,
-                connection_params=experiment.get("connection_params", {}),
+                connection_params={**experiment.get("connection_params", {})},
                 search_params=search_params,
             )
             for search_params in experiment.get("search_params", [{}])
diff --git a/engine/clients/qdrant/__init__.py b/engine/clients/qdrant/__init__.py
@@ -0,0 +1 @@
+QDRANT_COLLECTION_NAME = "benchmark_collection"
diff --git a/engine/clients/qdrant/config.py b/engine/clients/qdrant/config.py
diff --git a/engine/clients/qdrant/configure.py b/engine/clients/qdrant/configure.py
@@ -3,7 +3,7 @@
 
 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 import QDRANT_COLLECTION_NAME
 
 
 class QdrantConfigurator(BaseConfigurator):
diff --git a/engine/clients/qdrant/search.py b/engine/clients/qdrant/search.py
@@ -4,7 +4,7 @@
 from qdrant_client.http import models as rest
 
 from engine.base_client.search import BaseSearcher
-from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
+from engine.clients.qdrant import QDRANT_COLLECTION_NAME
 
 
 class QdrantSearcher(BaseSearcher):
diff --git a/engine/clients/qdrant/upload.py b/engine/clients/qdrant/upload.py
@@ -5,7 +5,7 @@
 from qdrant_client.http.models import Batch, CollectionStatus
 
 from engine.base_client.upload import BaseUploader
-from engine.clients.qdrant.config import QDRANT_COLLECTION_NAME
+from engine.clients.qdrant import QDRANT_COLLECTION_NAME
 
 
 class QdrantUploader(BaseUploader):
diff --git a/engine/clients/weaviate/__init__.py b/engine/clients/weaviate/__init__.py
@@ -0,0 +1,2 @@
+WEAVIATE_CLASS_NAME = 'benchmark'
+WEAVIATE_DEFAULT_PORT = 8080
diff --git a/engine/clients/weaviate/configure.py b/engine/clients/weaviate/configure.py
@@ -0,0 +1,40 @@
+from weaviate import Client
+
+from engine.base_client.configure import BaseConfigurator
+from engine.base_client.distances import Distance
+from engine.clients.weaviate import WEAVIATE_CLASS_NAME, WEAVIATE_DEFAULT_PORT
+
+
+class WeaviateConfigurator(BaseConfigurator):
+    DISTANCE_MAPPING = {
+        Distance.L2_SQUARED: "l2-squared",
+        Distance.COSINE: "cosine",
+        Distance.DOT: "dot",
+    }
+
+    def __init__(self, host, collection_params: dict, connection_params: dict):
+        super().__init__(host, collection_params, connection_params)
+        url = f"http://{host}:{connection_params.pop('port', WEAVIATE_DEFAULT_PORT)}"
+        self.client = Client(url, **connection_params)
+
+    def clean(self):
+        classes = self.client.schema.get()
+        for cl in classes["classes"]:
+            if cl["class"] == WEAVIATE_CLASS_NAME:
+                self.client.schema.delete_class(WEAVIATE_CLASS_NAME)
+
+    def recreate(
+            self, distance, vector_size, collection_params,
+    ):
+        self.client.schema.create_class({
+            "class": WEAVIATE_CLASS_NAME,
+            "vectorizer": "none",
+            "properties": [],
+            "vectorIndexConfig": {
+                **{
+                    "vectorCacheMaxObjects": 1000000000,
+                    "distance": self.DISTANCE_MAPPING.get(distance),
+                },
+                **collection_params["vectorIndexConfig"]
+            },
+        })
diff --git a/engine/clients/weaviate/search.py b/engine/clients/weaviate/search.py
@@ -0,0 +1,43 @@
+from typing import Optional, Tuple, List
+
+from weaviate import Client
+
+from engine.base_client.search import BaseSearcher
+from engine.clients.qdrant import QDRANT_COLLECTION_NAME
+from engine.clients.weaviate import WEAVIATE_DEFAULT_PORT
+
+
+class QdrantSearcher(BaseSearcher):
+    search_params = {}
+    client: Client = None
+
+    @classmethod
+    def init_client(cls, host, connection_params: dict, search_params: dict):
+        url = f"http://{host}:{connection_params.pop('port', WEAVIATE_DEFAULT_PORT)}"
+        cls.client = Client(url, **connection_params)
+        cls.search_params = search_params
+
+    @classmethod
+    def conditions_to_filter(cls, _meta_conditions):
+        # ToDo: implement
+        return None
+
+    @classmethod
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
+        top = 10
+        near_vector = {"vector": vector}
+        res = (
+            cls.client.query.get(cls.collection, ["_additional {id certainty}"])
+            .with_near_vector(near_vector)
+            .with_limit(top)
+            .do()
+        )
+        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
+        )
+
+        return [(hit.id, hit.score) for hit in res]
diff --git a/engine/clients/weaviate/upload.py b/engine/clients/weaviate/upload.py
@@ -0,0 +1,29 @@
+from typing import Optional, List
+
+from weaviate import Client
+
+from engine.base_client.upload import BaseUploader
+from engine.clients.weaviate import WEAVIATE_DEFAULT_PORT
+
+
+class WeaviateUploader(BaseUploader):
+    client = None
+    upload_params = {}
+
+    @classmethod
+    def init_client(cls, host, connection_params, upload_params):
+        url = f"http://{host}:{connection_params.pop('port', WEAVIATE_DEFAULT_PORT)}"
+        cls.client = Client(url, **connection_params)
+
+        cls.upload_params = upload_params
+        cls.connection_params = connection_params
+
+    @classmethod
+    def upload_batch(
+            cls,
+            ids: List[int],
+            vectors: List[list],
+            metadata: Optional[List[dict]]
+    ):
+        ...
+
diff --git a/experiments/configurations/weaviate-single-node.json b/experiments/configurations/weaviate-single-node.json
@@ -0,0 +1,20 @@
+[
+  {
+    "name": "weaviate-default",
+    "engine": "weaviate",
+    "connection_params": {},
+    "collection_params": {
+      "vectorIndexConfig": {
+        "ef": 100,
+        "efConstruction": 100,
+        "maxConnections": 16
+      }
+    },
+    "search_params": [
+      {"parallel":  1},
+      {"parallel":  2},
+      {"parallel":  4},
+      {"parallel":  100}
+    ]
+  }
+]
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/run.py b/run.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+QDRANT_COLLECTION_NAME = "benchmark_collection"`