redisbench-admin 0.11.63__py3-none-any.whl → 0.11.65__py3-none-any.whl

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/diskann.py

@@ -0,0 +1,190 @@
+import sys
+import os
+import vamanapy as vp
+import numpy as np
+import struct
+import time
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class Vamana(BaseANN):
+    def __init__(self, metric, param):
+        self.metric = {'angular': 'cosine', 'euclidean': 'l2'}[metric]
+        self.l_build = int(param["l_build"])
+        self.max_outdegree = int(param["max_outdegree"])
+        self.alpha = float(param["alpha"])
+        print("Vamana: L_Build = " + str(self.l_build))
+        print("Vamana: R = " + str(self.max_outdegree))
+        print("Vamana: Alpha = " + str(self.alpha))
+        self.params = vp.Parameters()
+        self.params.set("L", self.l_build)
+        self.params.set("R", self.max_outdegree)
+        self.params.set("C", 750)
+        self.params.set("alpha", self.alpha)
+        self.params.set("saturate_graph", False)
+        self.params.set("num_threads", 1)
+
+    def fit(self, X):
+
+        def bin_to_float(binary):
+            return struct.unpack('!f',struct.pack('!I', int(binary, 2)))[0]
+
+        print("Vamana: Starting Fit...")
+        index_dir = 'indices'
+
+        if not os.path.exists(index_dir):
+            os.makedirs(index_dir)
+
+        data_path = os.path.join(index_dir, 'base.bin')
+        self.name = 'Vamana-{}-{}-{}'.format(self.l_build,
+                                             self.max_outdegree, self.alpha)
+        save_path = os.path.join(index_dir, self.name)
+        print('Vamana: Index Stored At: ' + save_path)
+        shape = [np.float32(bin_to_float('{:032b}'.format(X.shape[0]))),
+                 np.float32(bin_to_float('{:032b}'.format(X.shape[1])))]
+        X = X.flatten()
+        X = np.insert(X, 0, shape)
+        X.tofile(data_path)
+
+        if not os.path.exists(save_path):
+            print('Vamana: Creating Index')
+            s = time.time()
+            if self.metric == 'l2':
+                index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, data_path)
+            elif self.metric == 'cosine':
+                index = vp.SinglePrecisionIndex(vp.Metric.INNER_PRODUCT, data_path)
+            else:
+                print('Vamana: Unknown Metric Error!')
+            index.build(self.params, [])
+            t = time.time()
+            print('Vamana: Index Build Time (sec) = ' + str(t - s))
+            index.save(save_path)
+        if os.path.exists(save_path):
+            print('Vamana: Loading Index: ' + str(save_path))
+            s = time.time()
+            if self.metric == 'l2':
+                self.index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, data_path)
+            elif self.metric == 'cosine':
+                self.index = vp.SinglePrecisionIndex(vp.Metric.INNER_PRODUCT, data_path)
+            else:
+                print('Vamana: Unknown Metric Error!')
+            self.index.load(file_name = save_path)
+            print("Vamana: Index Loaded")
+            self.index.optimize_graph()
+            print("Vamana: Graph Optimization Completed")
+            t = time.time()
+            print('Vamana: Index Load Time (sec) = ' + str(t - s))
+        else:
+            print("Vamana: Unexpected Index Build Time Error")
+
+        print('Vamana: End of Fit')
+
+    def set_query_arguments(self, l_search):
+        print("Vamana: L_Search = " + str(l_search))
+        self.l_search = l_search
+
+    def query(self, v, n):
+        return self.index.single_numpy_query(v, n, self.l_search)
+
+    def batch_query(self, X, n):
+        self.num_queries = X.shape[0]
+        self.result = self.index.batch_numpy_query(X, n, self.num_queries, self.l_search)
+
+    def get_batch_results(self):
+        return self.result.reshape((self.num_queries, self.result.shape[0] // self.num_queries))
+
+
+class VamanaPQ(BaseANN):
+    def __init__(self, metric, param):
+        self.metric = {'angular': 'cosine', 'euclidean': 'l2'}[metric]
+        self.l_build = int(param["l_build"])
+        self.max_outdegree = int(param["max_outdegree"])
+        self.alpha = float(param["alpha"])
+        self.chunks = int(param["chunks"])
+        print("Vamana PQ: L_Build = " + str(self.l_build))
+        print("Vamana PQ: R = " + str(self.max_outdegree))
+        print("Vamana PQ: Alpha = " + str(self.alpha))
+        print("Vamana PQ: Chunks = " + str(self.chunks))
+        self.params = vp.Parameters()
+        self.params.set("L", self.l_build)
+        self.params.set("R", self.max_outdegree)
+        self.params.set("C", 750)
+        self.params.set("alpha", self.alpha)
+        self.params.set("saturate_graph", False)
+        self.params.set("num_chunks", self.chunks)
+        self.params.set("num_threads", 1)
+
+    def fit(self, X):
+
+        def bin_to_float(binary):
+            return struct.unpack('!f',struct.pack('!I', int(binary, 2)))[0]
+
+        print("Vamana PQ: Starting Fit...")
+        index_dir = 'indices'
+
+        if self.chunks > X.shape[1]:
+            raise ValueError
+
+        if not os.path.exists(index_dir):
+            os.makedirs(index_dir)
+
+        data_path = os.path.join(index_dir, 'base.bin')
+        pq_path = os.path.join(index_dir, 'pq_memory_index')
+        self.name = 'VamanaPQ-{}-{}-{}'.format(self.l_build,
+                                             self.max_outdegree, self.alpha)
+        save_path = os.path.join(index_dir, self.name)
+        print('Vamana PQ: Index Stored At: ' + save_path)
+        shape = [np.float32(bin_to_float('{:032b}'.format(X.shape[0]))),
+                 np.float32(bin_to_float('{:032b}'.format(X.shape[1])))]
+        X = X.flatten()
+        X = np.insert(X, 0, shape)
+        X.tofile(data_path)
+
+        if not os.path.exists(save_path):
+            print('Vamana PQ: Creating Index')
+            s = time.time()
+            if self.metric == 'l2':
+                index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, data_path)
+            elif self.metric == 'cosine':
+                index = vp.SinglePrecisionIndex(vp.Metric.INNER_PRODUCT, data_path)
+            else:
+                print('Vamana PQ: Unknown Metric Error!')
+            index.pq_build(data_path, pq_path, self.params)
+            t = time.time()
+            print('Vamana PQ: Index Build Time (sec) = ' + str(t - s))
+            index.save(save_path)
+        if os.path.exists(save_path):
+            print('Vamana PQ: Loading Index: ' + str(save_path))
+            s = time.time()
+            if self.metric == 'l2':
+                self.index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, data_path)
+            elif self.metric == 'cosine':
+                self.index = vp.SinglePrecisionIndex(vp.Metric.INNER_PRODUCT, data_path)
+            else:
+                print('Vamana PQ: Unknown Metric Error!')
+            self.index.load(file_name = save_path)
+            print("Vamana PQ: Index Loaded")
+            self.index.pq_load(pq_prefix_path = pq_path)
+            print("Vamana PQ: PQ Data Loaded")
+            self.index.optimize_graph()
+            print("Vamana PQ: Graph Optimization Completed")
+            t = time.time()
+            print('Vamana PQ: Index Load Time (sec) = ' + str(t - s))
+        else:
+            print("Vamana PQ: Unexpected Index Build Time Error")
+
+        print('Vamana PQ: End of Fit')
+
+    def set_query_arguments(self, l_search):
+        print("Vamana PQ: L_Search = " + str(l_search))
+        self.l_search = l_search
+
+    def query(self, v, n):
+        return self.index.pq_single_numpy_query(v, n, self.l_search)
+
+    def batch_query(self, X, n):
+        self.num_queries = X.shape[0]
+        self.result = self.index.pq_batch_numpy_query(X, n, self.num_queries, self.l_search)
+
+    def get_batch_results(self):
+        return self.result.reshape((self.num_queries, self.result.shape[0] // self.num_queries))

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/dolphinnpy.py

@@ -0,0 +1,31 @@
+from __future__ import absolute_import
+import sys
+sys.path.append("install/lib-dolphinnpy")  # noqa
+import numpy
+import ctypes
+from dolphinn import Dolphinn
+from utils import findmean, isotropize
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class DolphinnPy(BaseANN):
+    def __init__(self, num_probes):
+        self.name = 'Dolphinn(num_probes={} )'.format(num_probes)
+        self.num_probes = num_probes
+        self.m = 1
+        self._index = None
+
+    def fit(self, X):
+        if X.dtype != numpy.float32:
+            X = numpy.array(X, dtype=numpy.float32)
+        d = X.shape[1]
+        self.m = findmean(X, d, 10)
+        X = isotropize(X, d, self.m)
+        hypercube_dim = int(numpy.log2(len(X))) - 2
+        self._index = Dolphinn(X, d, hypercube_dim)
+
+    def query(self, v, n):
+        q = numpy.array([v])
+        q = isotropize(q, len(v), self.m)
+        res = self._index.queries(q, n, self.num_probes)
+        return res[0]

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/dummy_algo.py

@@ -0,0 +1,25 @@
+from __future__ import absolute_import
+import numpy as np
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class DummyAlgoMt(BaseANN):
+    def __init__(self, metric):
+        self.name = 'DummyAlgoMultiThread'
+
+    def fit(self, X):
+        self.len = len(X) - 1
+
+    def query(self, v, n):
+        return np.random.randint(self.len, size=n)
+
+
+class DummyAlgoSt(BaseANN):
+    def __init__(self, metric):
+        self.name = 'DummyAlgoSingleThread'
+
+    def fit(self, X):
+        self.len = len(X) - 1
+
+    def query(self, v, n):
+        return np.random.randint(self.len, size=n)

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/elasticsearch.py

@@ -0,0 +1,107 @@
+"""
+ann-benchmarks interfaces for Elasticsearch.
+Note that this requires X-Pack, which is not included in the OSS version of Elasticsearch.
+"""
+import logging
+from time import sleep
+from os import environ
+from urllib.error import URLError
+
+from elasticsearch import Elasticsearch, BadRequestError
+from elasticsearch.helpers import bulk
+from elastic_transport.client_utils import DEFAULT
+
+from ann_benchmarks.algorithms.base import BaseANN
+
+# Configure the elasticsearch logger.
+# By default, it writes an INFO statement for every request.
+logging.getLogger("elasticsearch").setLevel(logging.WARN)
+
+# Uncomment these lines if you want to see timing for every HTTP request and its duration.
+# logging.basicConfig(level=logging.INFO)
+# logging.getLogger("elasticsearch").setLevel(logging.INFO)
+
+def es_wait(es):
+    print("Waiting for elasticsearch health endpoint...")
+    for i in range(30):
+        try:
+            res = es.cluster.health(wait_for_status='yellow', timeout='1s')
+            if not res['timed_out']: # then status is OK
+                print("Elasticsearch is ready")
+                return
+        except URLError:
+            pass
+        sleep(1)
+    raise RuntimeError("Failed to connect to elasticsearch server")
+
+
+class ElasticsearchScriptScoreQuery(BaseANN):
+    """
+    KNN using the Elasticsearch dense_vector datatype and script score functions.
+    - Dense vector field type: https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html
+    - Dense vector queries: https://www.elastic.co/guide/en/elasticsearch/reference/master/query-dsl-script-score-query.html
+    """
+
+    def __init__(self, metric: str, dimension: int, conn_params, method_param):
+        self.name = f"elasticsearch-script-score-query_metric={metric}_dimension={dimension}_params{method_param}"
+        self.metric = {"euclidean": 'l2_norm', "angular": 'cosine'}[metric]
+        self.method_param = method_param
+        self.dimension = dimension
+        self.timeout = 60 * 60
+        h = conn_params['host'] if conn_params['host'] is not None else 'localhost'
+        p = conn_params['port'] if conn_params['port'] is not None else '9200'
+        u = conn_params['user'] if conn_params['user'] is not None else 'elastic'
+        a = conn_params['auth'] if conn_params['auth'] is not None else ''
+        self.index = "ann_benchmark"
+        self.shards = conn_params['shards']
+        try:
+            self.es = Elasticsearch(f"http://{h}:{p}",  request_timeout=self.timeout, basic_auth=(u, a), refresh_interval=-1)
+            self.es.info()
+        except Exception:
+            self.es = Elasticsearch(f"https://{h}:{p}", request_timeout=self.timeout, basic_auth=(u, a), ca_certs=environ.get('ELASTIC_CA', DEFAULT))
+        self.batch_res = []
+        es_wait(self.es)
+
+    def fit(self, X):
+        mappings = dict(
+            properties=dict(
+                id=dict(type="keyword", store=True),
+                vec=dict(
+                    type="dense_vector",
+                    dims=self.dimension,
+                    similarity=self.metric,
+                    index=True,
+                    index_options=self.method_param
+                )
+            )
+        )
+        try:
+            self.es.indices.create(index=self.index, mappings=mappings, settings=dict(number_of_shards=self.shards, number_of_replicas=0))
+        except BadRequestError as e:
+            if 'resource_already_exists_exception' not in e.message: raise e
+
+        def gen():
+            for i, vec in enumerate(X):
+                yield { "_op_type": "index", "_index": self.index, "vec": vec.tolist(), 'id': str(i) }
+
+        (_, errors) = bulk(self.es, gen(), chunk_size=500, max_retries=9)
+        assert len(errors) == 0, errors
+
+        self.es.indices.refresh(index=self.index)
+        self.es.indices.forcemerge(index=self.index, max_num_segments=1)
+
+    def set_query_arguments(self, ef):
+        self.ef = ef
+
+    def query(self, q, n):
+        knn = dict(field='vec', query_vector=q.tolist(), k=n, num_candidates=self.ef)
+        res = self.es.knn_search(index=self.index, knn=knn, source=False, docvalue_fields=['id'],
+                                 stored_fields="_none_", filter_path=["hits.hits.fields.id"])
+        return [int(h['fields']['id'][0]) for h in res['hits']['hits']]
+
+    def batch_query(self, X, n):
+        self.batch_res = [self.query(q, n) for q in X]
+
+    def get_batch_results(self):
+        return self.batch_res
+

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/elastiknn.py

@@ -0,0 +1,124 @@
+"""
+ann-benchmarks interfaces for elastiknn: https://github.com/alexklibisz/elastiknn
+Uses the elastiknn python client
+To install a local copy of the client, run `pip install --upgrade -e /path/to/elastiknn/client-python/`
+To monitor the Elasticsearch JVM using Visualvm, add `ports={ "8097": 8097 }` to the `containers.run` call in runner.py.
+"""
+from sys import stderr
+from urllib.error import URLError
+
+import numpy as np
+from elastiknn.api import Vec
+from elastiknn.models import ElastiknnModel
+from elastiknn.utils import dealias_metric
+
+from ann_benchmarks.algorithms.base import BaseANN
+
+from urllib.request import Request, urlopen
+from time import sleep, perf_counter
+
+import logging
+
+# Mute the elasticsearch logger.
+# By default, it writes an INFO statement for every request.
+logging.getLogger("elasticsearch").setLevel(logging.WARN)
+
+
+def es_wait():
+    print("Waiting for elasticsearch health endpoint...")
+    req = Request("http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s")
+    for i in range(30):
+        try:
+            res = urlopen(req)
+            if res.getcode() == 200:
+                print("Elasticsearch is ready")
+                return
+        except URLError:
+            pass
+        sleep(1)
+    raise RuntimeError("Failed to connect to local elasticsearch")
+
+
+class Exact(BaseANN):
+
+    def __init__(self, metric: str, dimension: int):
+        self.name = f"eknn-exact-metric={metric}_dimension={dimension}"
+        self.metric = metric
+        self.dimension = dimension
+        self.model = ElastiknnModel("exact", dealias_metric(metric))
+        self.batch_res = None
+        es_wait()
+
+    def _handle_sparse(self, X):
+        # convert list of lists of indices to sparse vectors.
+        return [Vec.SparseBool(x, self.dimension) for x in X]
+
+    def fit(self, X):
+        if self.metric in {'jaccard', 'hamming'}:
+            return self.model.fit(self._handle_sparse(X), shards=1)[0]
+        else:
+            return self.model.fit(X, shards=1)
+
+    def query(self, q, n):
+        if self.metric in {'jaccard', 'hamming'}:
+            return self.model.kneighbors(self._handle_sparse([q]), n)[0]
+        else:
+            return self.model.kneighbors(np.expand_dims(q, 0), n)[0]
+
+    def batch_query(self, X, n):
+        if self.metric in {'jaccard', 'hamming'}:
+            self.batch_res = self.model.kneighbors(self._handle_sparse(X), n)
+        else:
+            self.batch_res = self.model.kneighbors(X, n)
+
+    def get_batch_results(self):
+        return self.batch_res
+
+
+class L2Lsh(BaseANN):
+
+    def __init__(self, L: int, k: int, w: int):
+        self.name_prefix = f"eknn-l2lsh-L={L}-k={k}-w={w}"
+        self.name = None  # set based on query args.
+        self.model = ElastiknnModel("lsh", "l2", mapping_params=dict(L=L, k=k, w=w))
+        self.X_max = 1.0
+        self.query_params = dict()
+        self.batch_res = None
+        self.sum_query_dur = 0
+        self.num_queries = 0
+        es_wait()
+
+    def fit(self, X):
+        print(f"{self.name_prefix}: indexing {len(X)} vectors")
+
+        # I found it's best to scale the vectors into [0, 1], i.e. divide by the max.
+        self.X_max = X.max()
+        return self.model.fit(X / self.X_max, shards=1)
+
+    def set_query_arguments(self, candidates: int, probes: int):
+        # This gets called when starting a new batch of queries.
+        # Update the name and model's query parameters based on the given params.
+        self.name = f"{self.name_prefix}_candidates={candidates}_probes={probes}"
+        self.model.set_query_params(dict(candidates=candidates, probes=probes))
+        # Reset the counters.
+        self.num_queries = 0
+        self.sum_query_dur = 0
+
+    def query(self, q, n):
+        # If QPS after 100 queries is < 10, this setting is bad and won't complete within the default timeout.
+        if self.num_queries > 100 and self.num_queries / self.sum_query_dur < 10:
+            print("Throughput after 100 queries is less than 10 q/s. Terminating to avoid wasteful computation.", flush=True)
+            exit(0)
+        else:
+            t0 = perf_counter()
+            res = self.model.kneighbors(np.expand_dims(q, 0) / self.X_max, n)[0]
+            dur = (perf_counter() - t0)
+            self.sum_query_dur += dur
+            self.num_queries += 1
+            return res
+
+    def batch_query(self, X, n):
+        self.batch_res = self.model.kneighbors(X, n)
+
+    def get_batch_results(self):
+        return self.batch_res

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/faiss.py

@@ -0,0 +1,124 @@
+from __future__ import absolute_import
+import sys
+sys.path.append("install/lib-faiss")  # noqa
+import numpy
+import sklearn.preprocessing
+import ctypes
+import faiss
+from ann_benchmarks.algorithms.base import BaseANN
+
+
+class Faiss(BaseANN):
+    def query(self, v, n):
+        if self._metric == 'angular':
+            v /= numpy.linalg.norm(v)
+        D, I = self.index.search(numpy.expand_dims(
+            v, axis=0).astype(numpy.float32), n)
+        return I[0]
+
+    def batch_query(self, X, n):
+        if self._metric == 'angular':
+            X /= numpy.linalg.norm(X)
+        self.res = self.index.search(X.astype(numpy.float32), n)
+
+    def get_batch_results(self):
+        D, L = self.res
+        res = []
+        for i in range(len(D)):
+            r = []
+            for l, d in zip(L[i], D[i]):
+                if l != -1:
+                    r.append(l)
+            res.append(r)
+        return res
+
+
+class FaissLSH(Faiss):
+    def __init__(self, metric, n_bits):
+        self._n_bits = n_bits
+        self.index = None
+        self._metric = metric
+        self.name = 'FaissLSH(n_bits={})'.format(self._n_bits)
+
+    def fit(self, X):
+        if X.dtype != numpy.float32:
+            X = X.astype(numpy.float32)
+        f = X.shape[1]
+        self.index = faiss.IndexLSH(f, self._n_bits)
+        self.index.train(X)
+        self.index.add(X)
+
+
+class FaissIVF(Faiss):
+    def __init__(self, metric, n_list):
+        self._n_list = n_list
+        self._metric = metric
+
+    def fit(self, X):
+        if self._metric == 'angular':
+            X = sklearn.preprocessing.normalize(X, axis=1, norm='l2')
+
+        if X.dtype != numpy.float32:
+            X = X.astype(numpy.float32)
+
+        self.quantizer = faiss.IndexFlatL2(X.shape[1])
+        index = faiss.IndexIVFFlat(
+            self.quantizer, X.shape[1], self._n_list, faiss.METRIC_L2)
+        index.train(X)
+        index.add(X)
+        self.index = index
+
+    def set_query_arguments(self, n_probe):
+        faiss.cvar.indexIVF_stats.reset()
+        self._n_probe = n_probe
+        self.index.nprobe = self._n_probe
+
+    def get_additional(self):
+        return {"dist_comps": faiss.cvar.indexIVF_stats.ndis +      # noqa
+                faiss.cvar.indexIVF_stats.nq * self._n_list}
+
+    def __str__(self):
+        return 'FaissIVF(n_list=%d, n_probe=%d)' % (self._n_list,
+                                                    self._n_probe)
+
+
+class FaissIVFPQfs(Faiss):
+    def __init__(self, metric, n_list):
+        self._n_list = n_list
+        self._metric = metric
+
+    def fit(self, X):
+        if X.dtype != numpy.float32:
+            X = X.astype(numpy.float32)
+        if self._metric == 'angular':
+            faiss.normalize_L2(X)
+
+        d = X.shape[1]
+        faiss_metric = faiss.METRIC_INNER_PRODUCT if self._metric == 'angular' else faiss.METRIC_L2
+        factory_string = f"IVF{self._n_list},PQ{d//2}x4fs"
+        index = faiss.index_factory(d, factory_string, faiss_metric)
+        index.train(X)
+        index.add(X)
+        index_refine = faiss.IndexRefineFlat(index, faiss.swig_ptr(X))
+        self.base_index = index
+        self.refine_index = index_refine
+
+    def set_query_arguments(self, n_probe, k_reorder):
+        faiss.cvar.indexIVF_stats.reset()
+        self._n_probe = n_probe
+        self._k_reorder = k_reorder
+        self.base_index.nprobe = self._n_probe
+        self.refine_index.k_factor = self._k_reorder
+        if self._k_reorder == 0:
+            self.index = self.base_index
+        else:
+            self.index = self.refine_index
+
+    def get_additional(self):
+        return {"dist_comps": faiss.cvar.indexIVF_stats.ndis +      # noqa
+                faiss.cvar.indexIVF_stats.nq * self._n_list}
+
+    def __str__(self):
+        return 'FaissIVFPQfs(n_list=%d, n_probe=%d, k_reorder=%d)' % (self._n_list,
+                                                                      self._n_probe,
+                                                                      self._k_reorder)

redisbench_admin/run/ann/pkg/ann_benchmarks/algorithms/faiss_gpu.py

@@ -0,0 +1,61 @@
+from __future__ import absolute_import
+import sys
+# Assumes local installation of FAISS
+sys.path.append("faiss")  # noqa
+import numpy
+import ctypes
+import faiss
+from ann_benchmarks.algorithms.base import BaseANN
+
+# Implementation based on
+# https://github.com/facebookresearch/faiss/blob/master/benchs/bench_gpu_sift1m.py  # noqa
+
+
+class FaissGPU(BaseANN):
+    def __init__(self, n_bits, n_probes):
+        self.name = 'FaissGPU(n_bits={}, n_probes={})'.format(
+            n_bits, n_probes)
+        self._n_bits = n_bits
+        self._n_probes = n_probes
+        self._res = faiss.StandardGpuResources()
+        self._index = None
+
+    def fit(self, X):
+        X = X.astype(numpy.float32)
+        self._index = faiss.GpuIndexIVFFlat(self._res, len(X[0]), self._n_bits,
+                                            faiss.METRIC_L2)
+        # self._index = faiss.index_factory(len(X[0]),
+        #                                   "IVF%d,Flat" % self._n_bits)
+        # co = faiss.GpuClonerOptions()
+        # co.useFloat16 = True
+        # self._index = faiss.index_cpu_to_gpu(self._res, 0,
+        #                                      self._index, co)
+        self._index.train(X)
+        self._index.add(X)
+        self._index.setNumProbes(self._n_probes)
+
+    def query(self, v, n):
+        return [label for label, _ in self.query_with_distances(v, n)]
+
+    def query_with_distances(self, v, n):
+        v = v.astype(numpy.float32).reshape(1, -1)
+        distances, labels = self._index.search(v, n)
+        r = []
+        for l, d in zip(labels[0], distances[0]):
+            if l != -1:
+                r.append((l, d))
+        return r
+
+    def batch_query(self, X, n):
+        self.res = self._index.search(X.astype(numpy.float32), n)
+
+    def get_batch_results(self):
+        D, L = self.res
+        res = []
+        for i in range(len(D)):
+            r = []
+            for l, d in zip(L[i], D[i]):
+                if l != -1:
+                    r.append(l)
+            res.append(r)
+        return res