Improved search function by key

.gitignore

@@ -45,7 +45,8 @@ __pycache__/
 *.log
 *.out
 log
-data
+#data
 *.ini
 conf.yaml
+faiss-server.yaml
 

Dockerfile

@@ -13,6 +13,9 @@ RUN pip install azure-storage-blob
 
 # Set the working directory to /app
 RUN mkdir -p /app
+RUN mkdir -p /app/data
+RUN mkdir -p /app/data/index
+RUN mkdir -p /app/log
 ADD . /app/
 WORKDIR /app
 

README.md

@@ -10,20 +10,20 @@ $ docker build -t [YOUR DOCKER IMAGE NAME] .
 ```
 or 
 ```
-$ ./build [YOUR DOCKER IMAGE NAME]
+$ ./build.sh [YOUR DOCKER IMAGE NAME]
 ```
 eg.
 ```
-$ ./build cia/faiss-server
+$ ./build.sh cia/faiss-server
 ```
 
 ## Run the faiss server
 ```
-$ ./run [YOUR DOCKER IMAGE] [YOUR DOCKER CONTAINER NAME]
+$ ./run.sh [YOUR DOCKER IMAGE] [YOUR DOCKER CONTAINER NAME]
 ```
 eg. 
 ```
-$ ./run cia/faiss-server:latest faiss-server
+$ ./run.sh cia/faiss-server:latest faiss-server
 ```
 
 
@@ -34,4 +34,9 @@ $ docker exec -it faiss-server /bin/bash
 ```
 ```
 $ python client_sample.py test --dim 200 --host localhost --port 50051
+```
+
+```
+$ python client_sample.py import data/embeds.csv data/ids.csv data/keys.csv --host localhost:50051
+$ python client_sample.py search-by-key a2 --host localhost:50051 --count 2
 ```

client_sample.py

@@ -76,16 +76,17 @@ def test(host, port, dim):
 @click.command('import')
 @click.argument('embs-path')
 @click.argument('ids-path')
+@click.argument('keys_path')
 @click.option('-h', '--host', default='localhost:50051', help='server host:port')
-def imports(host, embs_path, ids_path):
+def imports(host, embs_path, ids_path, keys_path):
     print("host: %s" % host)
     channel = grpc.insecure_channel(host)
     stub = pb2_grpc.ServerStub(channel)
 
     response = stub.Total(pb2.EmptyRequest())
     print("total: %d" % response.count)
 
-    response = stub.Import(pb2.ImportRequest(embs_path=embs_path, ids_path=ids_path))
+    response = stub.Import(pb2.ImportRequest(embs_path=embs_path, ids_path=ids_path, keys_path=keys_path))
     print("response: %s" % response.message)
 
     response = stub.Total(pb2.EmptyRequest())

data/embeds.csv

@@ -0,0 +1,2 @@
+0.1890639	0.9702392	0.013680172	0.092980236	0.41790208	0.45386073	0.42763093	0.18946815	0.99329203	0.23457207	0.20104194	0.44114107	0.31397247	0.5472484	0.49339047	0.055004463	0.82110214	0.59724337	0.3030562	0.5604773	0.86728984	0.14949812	0.5419759	0.4310291	0.6337217	0.0032894297	0.4160197	0.8318942	0.23526901	0.72281927	0.26539928	0.35501707	0.29356173	0.46452603	0.35151333	0.8819671	0.27393293	0.24769767	0.8637873	0.8788942	0.089255996	0.43524617	0.44403502	0.16971748	0.8279367	0.6133808	0.5200635	0.40707567	0.06089727	0.024314305	0.17863126	0.42186296	0.91130155	0.62434167	0.52470624	0.44659892	0.80466014	0.80982417	0.14247158	0.5217549	0.34561506	0.24736822	0.6818239	0.78890514	0.84892565	0.13298817	0.30310392	0.9118455	0.912892	0.59955186	0.22453378	0.70343715	0.99294907	0.9670585	0.29036522	0.33896813	0.9677055	0.622309	0.035264466	0.8804555	0.83223575	0.41757375	0.654066	0.76559556	0.6438391	0.14858608	0.7449582	0.14724912	0.72489494	0.37603164	0.6330364	0.06426713	0.5521647	0.38839877	0.09412226	0.7804446	0.24404515	0.3232764	0.020562846	0.27378193	0.14074397	0.38344228	0.51933336	0.3208555	0.33973262	0.7230768	0.39680433	0.9440062	0.40030423	0.13196363	0.54955983	0.2602228	0.98967415	0.19574209	0.4396908	0.15687948	0.6528181	0.16210744	0.6821134	0.9838601	0.6030431	0.76025194	0.5603035	0.30253878	0.95268965	0.95803154	0.88824934	0.39214066	0.97051394	0.98006177	0.5589214	0.89685917	0.49186754	0.10588156	0.99579036	0.9320394	0.16306546	0.9771742	0.4154118	0.73195076	0.9704186	0.08134132	0.44648233	0.26538768	0.24470824	0.4622599	0.5520417	0.9491697	0.1569141	0.13408567	0.0723426	0.09442261	0.6284843	0.7555779	0.09291987	0.5963138	0.6409636	0.34491828	0.55458796	0.3242194	0.47418925	0.49232462	0.8206485	0.17203958	0.70475864	0.8180915	0.47923547	0.64192975	0.27658936	0.9921742	0.98313874	0.64295906	0.8723445	0.9711968	0.33215186	0.90582335	0.3844488	0.8416527	0.23931529	0.07138017	0.35610893	0.9391459	0.1473125	0.33508018	0.19288318	0.5700721	0.8662711	0.7108658	0.9739855	0.18008451	0.49275267	0.0756371	0.30579612	0.8620466	0.22285913	0.37046522	0.03124313	0.022965709	0.5979236	0.62304807
+0.46847782	0.5724462	0.35274652	0.62228197	0.94845986	0.82393634	0.36760718	0.9329617	0.25462407	0.03489479	0.6343733	0.89104235	0.0088841375	0.434985	0.45548746	0.39551702	0.29033276	0.08154021	0.39270833	0.09074037	0.49404362	0.47930077	0.6845382	0.25939927	0.63000745	0.34634835	0.6582647	0.111826	0.7104301	0.83143544	0.5640379	0.096034795	0.42972058	0.19082488	0.1321226	0.024501482	0.725668	0.6897104	0.4250858	0.1251062	0.5091508	0.6207674	0.8351553	0.20070034	0.72079927	0.19090374	0.46887815	0.39961165	0.2627553	0.7282705	0.32426667	0.6449381	0.8896842	0.5562882	0.14014505	0.7301511	0.020518007	0.9928512	0.92046094	0.49451125	0.37993443	0.33404455	0.59406567	0.56489307	0.3928042	0.18487878	0.6123481	0.50253797	0.48198527	0.33171695	0.6836534	0.9708121	0.41574663	0.2875891	0.2252129	0.12852979	0.0787948	0.47629118	0.7980404	0.20445277	0.19953942	0.965661	0.025018165	0.7956314	0.13388361	0.43737453	0.45942026	0.25558746	0.20246962	0.48806503	0.9837458	0.29225153	0.71916616	0.9526122	0.57358354	0.14428927	0.49091274	0.36894524	0.3147875	0.96187484	0.39117917	0.35355663	0.28674498	0.23659502	0.916376	0.99271977	0.25626796	0.82681125	0.6115973	0.46268886	0.41223294	0.18178888	0.11981806	0.63761073	0.20202398	0.9390866	0.77049565	0.36034492	0.22835334	0.001948912	0.40401134	0.09552054	0.18058862	0.779659	0.19122343	0.8159597	0.7380383	0.39620742	0.87153494	0.3415848	0.6953568	0.823221	0.6039216	0.62376976	0.1957188	0.24315614	0.520875	0.047078457	0.59249943	0.3042383	0.17286436	0.5492452	0.5982279	0.055363417	0.07818982	0.97224975	0.71646374	0.3228178	0.009316032	0.24749126	0.93300414	0.6606999	0.4894559	0.64867544	0.43334386	0.88112855	0.74602276	0.44531894	0.8527712	0.4622631	0.020371098	0.19177161	0.2908577	0.21554473	0.16954689	0.57741165	0.11027071	0.12618637	0.92059225	0.21639249	0.79428846	0.4267286	0.9426855	0.8993863	0.13833722	0.14289147	0.96281713	0.29373172	0.2754965	0.23835607	0.71243083	0.6251121	0.50218743	0.20136604	0.90252954	0.7483757	0.14292905	0.8303706	0.35279363	0.17666304	0.86918056	0.05527037	0.042896952	0.7173806	0.5846493	0.12757616	0.37300506	0.8198024	0.14848578	0.1721672

data/ids.csv

@@ -0,0 +1,2 @@
+0
+1

data/keys.csv

@@ -0,0 +1,2 @@
+a1
+a2

faiss-server.yaml.sample

@@ -0,0 +1,24 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: faiss-server-deployment
+  labels:
+    app: faiss-server
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: faiss-server
+  template:
+    metadata:
+      labels:
+        app: faiss-server
+    spec:
+      containers:
+        - name: faiss-server
+          image: YOUR_CONTAINER_REGISTRY/faiss-server:v1
+          imagePullPolicy: Always
+          command: ["server.py"]
+              args: ["--dim", "200", "--log", "/app/log/faiss.log", "--save_path", "/app/data/index", "--nprobe", "1"]
+          ports:
+            - containerPort: 50051

faiss.proto

@@ -41,6 +41,7 @@ message RestoreRequest {
 message ImportRequest {
     string embs_path = 1;
     string ids_path = 2;
+    string keys_path = 3;
 }
 
 message EmptyRequest {

faiss_server.py

@@ -82,6 +82,7 @@ def _load_keys(self, keys_path):
         _, keys_path = self.down_if_remote_path(keys_path)
         keys = pd.read_csv(keys_path, header=None, squeeze=True, dtype=("str"))
         key_index = pd.Index(keys)
+        logging.debug("keys: keys=%s, keys_index=%s", keys.values[:10], key_index[:10])
         return keys.values, key_index
 
     def Total(self, request, context):
@@ -110,6 +111,8 @@ def Search(self, request, context):
             if not self._key_index.contains(request.key):
                 return pb2.SearchResponse()
             request.id = self._key_index.get_loc(request.key)
+            logging.debug("Search: key=%s, found id=%s", request.key, request.id)
+
         D, I = self._index.search_by_id(request.id, request.count)
         K = None
         if request.key:
@@ -132,9 +135,10 @@ def Restore(self, request, context):
         return pb2.SimpleResponse(message="Restored, %s!" % request.save_path)
 
     def Import(self, request, context):
-        logging.debug("importing - %s, %s", request.embs_path, request.ids_path)
+        logging.debug("importing - %s, %s, %s", request.embs_path, request.ids_path, request.keys_path)
         _, embs_path = self.down_if_remote_path(request.embs_path)
         _, ids_path = self.down_if_remote_path(request.ids_path)
+        _, keys_path = self.down_if_remote_path(request.keys_path)
         df = pd.read_csv(embs_path, delimiter="\t", header=None)
         X = df.values
         # logging.debug("X = %s", X)
@@ -146,7 +150,10 @@ def Import(self, request, context):
         ids = np.ascontiguousarray(ids, dtype=np.int64)
 
         self._index.replace(X, ids)
-        return pb2.SimpleResponse(message="Imported, %s, %s!" % (request.embs_path, request.ids_path))
+
+        self._keys, self._key_index = self._load_keys(keys_path)
+
+        return pb2.SimpleResponse(message="Imported, %s, %s, %s!" % (request.embs_path, request.ids_path, request.keys_path))
 
     def save(self):
         logging.debug("saving index to %s", self._save_path)