Initial commit

.dockerignore

@@ -0,0 +1,3 @@
+data
+*.key
+*.index

.gitignore

@@ -0,0 +1,49 @@
+
+# Gradle
+build
+dist
+.gradle
+
+# Maven
+target
+release.properties
+pom.xml.*
+
+# IntelliJ
+.idea/
+*.iml
+
+# VSCode
+.vscode
+
+# Eclipse
+.classpath
+.project
+.settings
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+.ipynb_checkpoints/
+# Egg
+*.egg-info
+
+# Virtualenv
+.venv
+
+# Java
+*.class
+
+# C extensions
+*.so
+
+# Etc
+*.DS_Store
+*.swp
+*.swo
+*.log
+*.out
+log
+data
+

Dockerfile

@@ -0,0 +1,16 @@
+FROM daangn/faiss
+
+ENV GRPC_PYTHON_VERSION 1.15.0
+RUN python -m pip install --upgrade pip
+RUN pip install grpcio==${GRPC_PYTHON_VERSION} grpcio-tools==${GRPC_PYTHON_VERSION}
+
+RUN pip install pandas
+
+# Set the working directory to /app
+RUN mkdir -p /app
+ADD . /app/
+WORKDIR /app
+
+ENTRYPOINT ["python"]
+CMD ["server.py"]
+

README.md

@@ -0,0 +1,20 @@
+
+## Protocol buffer compile for gRPC
+```
+$ python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. faiss_index.proto
+```
+
+## Build docker image
+```
+docker build -t cia/faiss-server .
+```
+
+## Run the faiss server
+```
+$ ./run [YOUR DOCKER IMAGE] [YOUR DOCKER CONTAINER NAME]
+```
+
+## Enter the docker container with shell command.
+```
+docker exec -it faiss-server /bin/bash
+```

faiss.proto

@@ -0,0 +1,62 @@
+syntax = "proto3";
+
+package faiss_index;
+
+// Service Definition
+service Server {
+    rpc Add (AddRequest) returns (SimpleResponse) {}
+    rpc Remove (IdRequest) returns (SimpleResponse) {}
+    rpc Search (SearchRequest) returns (SearchResponse) {}
+    rpc SearchByEmbedding (SearchByEmbeddingRequest) returns (SearchResponse) {}
+    rpc Restore (RestoreRequest) returns (SimpleResponse) {}
+    rpc Import (ImportRequest) returns (SimpleResponse) {}
+    rpc Total (EmptyRequest) returns (TotalResponse) {}
+}
+
+// Message Definition
+message AddRequest {
+    int64 id = 1;
+    repeated float embedding = 2;
+}
+
+message IdRequest {
+    int64 id = 1;
+}
+
+message SearchRequest {
+    int64 id = 1;
+    int32 count = 2;
+    string key = 3;
+}
+
+message SearchByEmbeddingRequest {
+    repeated float embedding = 1;
+    int32 count = 2;
+}
+
+message RestoreRequest {
+    string save_path = 1;
+}
+
+message ImportRequest {
+    string embs_path = 1;
+    string ids_path = 2;
+}
+
+message EmptyRequest {
+}
+
+message SimpleResponse {
+    string message = 1;
+}
+
+message SearchResponse {
+    repeated int64 ids = 1;
+    repeated float scores = 2;
+    repeated string keys = 3;
+}
+
+message TotalResponse {
+    int64 count = 1;
+}
+

faiss_index.py

@@ -0,0 +1,57 @@
+import logging
+from os.path import isfile
+
+import faiss
+import numpy as np
+
+
+class FaissIndex:
+    def __init__(self, dim, save_path):
+        if isfile(save_path):
+            logging.debug("restore: %s", save_path)
+            self._index = faiss.read_index(save_path)
+        else:
+            self._sub_index = faiss.IndexFlat(dim)
+            self._index = faiss.IndexIDMap2(self._sub_index)
+
+    def replace(self, xb, ids):
+        logging.debug("replace: %s", ids)
+        self.remove(ids)
+        return self._index.add_with_ids(xb, ids)
+
+    def add(self, xb, ids):
+        return self._index.add_with_ids(xb, ids)
+
+    def search(self, xq, k=10):
+        return self._index.search(xq, k)
+
+    def search_by_id(self, id, k=10):
+        try:
+            x = self._index.reconstruct(id)
+            xq = np.expand_dims(x, axis=0)
+        except RuntimeError as e:
+            if str(e).endswith("not found"):
+                return ([None], [None])
+            else:
+                raise e
+        return self.search(xq, k)
+
+    def ntotal(self):
+        return self._index.ntotal
+
+    def remove(self, ids):
+        return self._index.remove_ids(ids)
+
+    def restore(self, filepath):
+        pre_index = self._index
+        self._index = faiss.read_index(filepath)
+        if pre_index:
+            pre_index.reset()
+
+    def save(self, filepath):
+        if self.ntotal() > 0:
+            faiss.write_index(self._index, filepath)
+
+    def set_nprobe(self, nprobe):
+        faiss.ParameterSpace().set_index_parameter(self._index, "nprobe", nprobe)
+