diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml index 4f345e7ea..565e57d63 100644 --- a/.github/workflows/benchmarks.yml +++ b/.github/workflows/benchmarks.yml @@ -54,6 +54,7 @@ jobs: - sklearn - sptag - tinyknn + - usearch - vald - vearch - vespa diff --git a/algos.yaml b/algos.yaml index 7de437df2..5d8515ca6 100644 --- a/algos.yaml +++ b/algos.yaml @@ -142,6 +142,90 @@ float: arg-groups: - {"M": 96, "efConstruction": 500} query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + usearch-f32: + docker-tag: ann-benchmarks-usearch + module: ann_benchmarks.algorithms.usearch + constructor: USearch + base-args: ["@metric", "f32"] + run-groups: + M-4: + arg-groups: + - {"M": 4, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-8: + arg-groups: + - {"M": 8, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-12: + arg-groups: + - {"M": 12, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-16: + arg-groups: + - {"M": 16, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-24: + arg-groups: + - {"M": 24, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-36: + arg-groups: + - {"M": 36, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-48: + arg-groups: + - {"M": 48, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-64: + arg-groups: + - {"M": 64, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-96: + arg-groups: + - {"M": 96, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + usearch-f8: + docker-tag: ann-benchmarks-usearch + module: ann_benchmarks.algorithms.usearch + constructor: USearch + base-args: ["@metric", "f8"] + run-groups: + M-4: + arg-groups: + - {"M": 4, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-8: + arg-groups: + - {"M": 8, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-12: + arg-groups: + - {"M": 12, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-16: + arg-groups: + - {"M": 16, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-24: + arg-groups: + - {"M": 24, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-36: + arg-groups: + - {"M": 36, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-48: + arg-groups: + - {"M": 48, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-64: + arg-groups: + - {"M": 64, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] + M-96: + arg-groups: + - {"M": 96, "efConstruction": 500} + query-args: [[10, 20, 40, 80, 120, 200, 400, 600, 800]] hnsw(faiss): docker-tag: ann-benchmarks-faiss diff --git a/ann_benchmarks/algorithms/usearch.py b/ann_benchmarks/algorithms/usearch.py new file mode 100644 index 000000000..036b0b1d6 --- /dev/null +++ b/ann_benchmarks/algorithms/usearch.py @@ -0,0 +1,63 @@ +from usearch.index import Index, MetricKind, ScalarKind +from usearch.numba import jit +import numpy as np + +from .base import BaseANN + +class USearch(BaseANN): + + def __init__(self, metric: str, accuracy: str, method_param: dict): + assert accuracy in ['f64', 'f32', 'f16', 'f8'] + assert metric in ['angular', 'euclidean'] + assert 'M' in method_param + assert 'efConstruction' in method_param + + self._method_param = method_param + self._accuracy = {'f64': ScalarKind.F64, 'f32': ScalarKind.F32, 'f8': ScalarKind.F8}[accuracy] + self._metric = {'angular': MetricKind.Cos, 'euclidean': MetricKind.L2sq}[metric] + + def __str__(self): + connectivity = self._method_param['M'] + expansion_add = self._method_param['efConstruction'] + return f'USearch(connecitivity={connectivity}, expansion_add={expansion_add})' + + def fit(self, X): + connectivity = self._method_param['M'] + expansion_add = self._method_param['efConstruction'] + metric = jit( + X.shape[1], + self._metric, + self._accuracy + ) + + self._index = Index( + ndim=len(X[0]), + metric=metric, + dtype=self._accuracy, + connectivity=connectivity, + expansion_add=expansion_add + ) + + labels = np.arange(len(X), dtype=np.longlong) + self._index.add(labels, np.asarray(X)) + + def get_memory_usage(self) -> int: + if not hasattr(self, '_index'): + return 0 + + return self._index.memory_usage / 1024 + + def set_query_arguments(self, ef: int): + self._index.expansion_search = ef + + def freeIndex(self): + del self._index + + def query(self, v, n): + return self._index.search(np.expand_dims(v, axis=0), k=n)[0][0] + + def batch_query(self, X, n): + self._batch_results = self._index.search(np.asarray(X), n) + + def get_batch_results(self): + return self._batch_results diff --git a/install/Dockerfile.usearch b/install/Dockerfile.usearch new file mode 100644 index 000000000..298730723 --- /dev/null +++ b/install/Dockerfile.usearch @@ -0,0 +1,8 @@ +FROM ann-benchmarks + +RUN apt-get install -y python-setuptools python-pip +RUN pip3 install pybind11 numpy setuptools numba +RUN git clone https://github.com/unum-cloud/usearch.git && cd usearch && git submodule update --init --recursive +RUN cd usearch && python3 setup.py install + +RUN python3 -c 'import usearch'