def Remove(self, request, context):
logging.debug("remove - id: %d", request.id)
ids = np.array([request.id], dtype=np.int64)
removed_count = self._index.remove(ids)
if removed_count < 1:
return pb2.SimpleResponse(message="Not existed, %s!" % request.id)
return pb2.SimpleResponse(message="Removed, %s!" % request.id)
def Add(self, request, context):
logging.debug("add - id: %d, %s", request.id, request.key)
if request.key:
# if self._key_index is None or not self._key_index.contains(request.key):
if self._key_index is None or request.key not in self._key_index:
if self._key_index is None:
self._key_index = pd.Index([request.key])
else:
self._key_index = self._key_index.append(
pd.Index([request.key]))
request.id = self._key_index.get_loc(request.key)
if self._keys is None:
self._keys = np.array([request.key])
else:
self._keys = np.append(self._keys, [request.key])
else:
request.id = self._key_index.get_loc(request.key)
# For debugging
# if self._keys is not None and self._key_index is not None:
# logging.debug("keys: keys=%s, keys_index=%s", self._keys, self._key_index)
xb = np.expand_dims(np.array(request.embedding, dtype=np.float32), 0)
ids = np.array([request.id], dtype=np.int64)
self._index.replace(xb, ids)
return pb2.SimpleResponse(message="Added, %d!" % request.id)
def Restore(self, request, context):
logging.debug("restore - %s", request.save_path)
remote_path, save_path = self.down_if_remote_path(request.save_path)
self._remote_path = remote_path
self._save_path = save_path
self._index.restore(request.save_path)
return pb2.SimpleResponse(message="Restored, %s!" % request.save_path)
def Import(self, request, context):
logging.info("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)
df = pd.read_csv(ids_path, header=None)
ids = df[0].values
logging.info("ids[size=%d] = %s", len(ids), ids)
X = np.ascontiguousarray(X, dtype=np.float32)
ids = np.ascontiguousarray(ids, dtype=np.int64)
# self._index.replace(X, ids)
self._index.rebuild(X, ids)
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 Reset(self, request, context):
logging.debug("reset")
self._index.reset()
self._keys = None
self._key_index = None
return pb2.SimpleResponse(message="Reset!")