usearch 0.13.0 → 0.15.0

package/python/scripts/bench.ipynb

@@ -0,0 +1,170 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install pandas nbformat jinja2 kaleido tqdm sklearn scikit-learn scipy matplotlib"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from bench import *\n",
+    "import plotly.graph_objects as go"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def speed_chart(df, title):\n",
+    "    subtitles = [\n",
+    "        f'{n}, recall@1 ~{r*100:.2f}%'\n",
+    "        for n, r in zip(df['names'], df['recall_at_one'])\n",
+    "    ]\n",
+    "    fig = go.Figure(\n",
+    "        data=[\n",
+    "            go.Bar(name='Add', x=subtitles, y=df['add_per_second']),\n",
+    "            go.Bar(name='Search', x=subtitles, y=df['search_per_second']),\n",
+    "        ])\n",
+    "    fig.update_layout(\n",
+    "        width=1000,\n",
+    "        height=1000,\n",
+    "        barmode='group',\n",
+    "        title=dict(\n",
+    "            text=title,\n",
+    "        ),\n",
+    "        legend=dict(\n",
+    "            orientation='h',\n",
+    "            entrywidth=70,\n",
+    "            yanchor='bottom',\n",
+    "            y=1.02,\n",
+    "            xanchor='right',\n",
+    "            x=1,\n",
+    "        ),\n",
+    "    )\n",
+    "    fig.write_image(title + '.png')\n",
+    "    return fig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = Dataset.build(\n",
+    "    vectors='../../datasets/wiki_1M/base.1M.fbin',\n",
+    "    queries='../../datasets/wiki_1M/query.public.100K.fbin',\n",
+    "    neighbors='../../datasets/wiki_1M/groundtruth.public.100K.ibin',\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eval_default = Evaluation.for_dataset(data)\n",
+    "results_default = bench_speed(eval_default)\n",
+    "results_default"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "speed_chart(results_default, 'USearch vs FAISS Speed')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eval_sliced = Evaluation.for_dataset(data, batch_size=256)\n",
+    "results_sliced = bench_speed(eval_sliced)\n",
+    "results_sliced"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "speed_chart(results_sliced, 'USearch vs FAISS Speed with 256 Batch Size')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eval_clustered = Evaluation.for_dataset(data, clusters=10_000)\n",
+    "len(eval_clustered.tasks)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results_clustered = bench_speed(eval_clustered)\n",
+    "results_clustered"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "speed_chart(results_clustered, 'USearch vs FAISS Speed on Clustered Data')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "usearch-test",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.3"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

package/python/scripts/bench.py

@@ -1,198 +1,115 @@
-import time
-import numpy as np
-import math
-from typing import Tuple, Any, Optional
+import os
+import logging
+import itertools
+from typing import List, Optional, Iterable
+from dataclasses import asdict
 
 import fire
+import numpy as np
+import pandas as pd
+
+from usearch.index import Index, Label
+from usearch.eval import Evaluation, Dataset, AddTask
+from usearch.index import (
+    DEFAULT_CONNECTIVITY,
+    DEFAULT_EXPANSION_ADD,
+    DEFAULT_EXPANSION_SEARCH,
+)
+
+
+def bench_speed(
+    eval: Evaluation,
+    connectivity: int = DEFAULT_CONNECTIVITY,
+    expansion_add: int = DEFAULT_EXPANSION_ADD,
+    expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+    train: bool = False,
+) -> pd.DataFrame:
+
+    # Build various indexes:
+    indexes = []
+    jit_options = [False]
+    dtype_options = ['f32', 'f16', 'f8']
+    for jit, dtype in itertools.product(jit_options, dtype_options):
 
-from usearch.index import Index
-from usearch.io import load_matrix
-
-
-def measure(f) -> Tuple[float, Any]:
-    a = time.time_ns()
-    result = f()
-    b = time.time_ns()
-    c = b - a
-    secs = c / (10 ** 9)
-    print(f'- Took: {secs:.2f} seconds')
-    return secs, result
-
-
-def bench_faiss(
-        index,
-        vectors: np.ndarray, queries: np.ndarray, neighbors: np.ndarray):
-
-    dt, _ = measure(lambda: index.add(vectors))
-    print(f'- Performance: {vectors.shape[0]/dt:.2f} insertions/s')
-
-    dt, results = measure(lambda: index.search(queries, neighbors.shape[1]))
-    print(f'- Performance: {queries.shape[0]/dt:.2f} queries/s')
-
-    recall_at_one = 0.0
-    matches = results[1]
-    for i in range(queries.shape[0]):
-        recall_at_one += neighbors[i, 0] in matches[i, :]
-    recall_at_one /= queries.shape[0]
-    print(f'- Recall@1: {recall_at_one * 100:.2f} %')
-
-
-def bench_usearch(
-        index,
-        vectors: np.ndarray, queries: np.ndarray, neighbors: np.ndarray):
-
-    labels = np.arange(vectors.shape[0], dtype=np.longlong)
-    assert len(index) == 0
-    dt, _ = measure(lambda: index.add(labels, vectors))
-    assert len(index) == vectors.shape[0]
-    print(f'- Performance: {vectors.shape[0]/dt:.2f} insertions/s')
-
-    dt, results = measure(lambda: index.search(queries, neighbors.shape[1]))
-    print(f'- Performance: {queries.shape[0]/dt:.2f} queries/s')
-
-    recall_at_one = 0.0
-    matches = results[0]
-    for i in range(queries.shape[0]):
-        recall_at_one += neighbors[i, 0] in matches[i, :]
-    recall_at_one /= queries.shape[0]
-    print(f'- Recall@1: {recall_at_one * 100:.2f} %')
-
-
-def main(
-    vectors: str,
-    queries: str,
-    neighbors: str,
-    connectivity: int = 16,
-    expansion_add: int = 128,
-    expansion_search: int = 64,
-    k: Optional[int] = None,
-):
-
-    vectors_mat = load_matrix(vectors)
-    queries_mat = load_matrix(queries)
-    neighbors_mat = load_matrix(neighbors)
-    dim = vectors_mat.shape[1]
-    if k:
-        neighbors_mat = neighbors_mat[:, :1]
-
-    print('USearch: HNSW')
-    index = Index(
-        ndim=dim,
-        expansion_add=expansion_add,
-        expansion_search=expansion_search,
-        connectivity=connectivity,
-    )
-    bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
-
-    print('USearch: HNSW + Half Precision')
-    index = Index(
-        ndim=dim,
-        dtype='f16',
-        expansion_add=expansion_add,
-        expansion_search=expansion_search,
-        connectivity=connectivity,
-    )
-    bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
-
-    print('USearch: HNSW + Quarter Precision')
-    index = Index(
-        ndim=dim,
-        dtype='f8',
-        expansion_add=expansion_add,
-        expansion_search=expansion_search,
-        connectivity=connectivity,
-    )
-    bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
-
-    # Showcases how to use Numba to JIT-compile similarity measures for USearch.
-    # https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
-    try:
-        from numba import cfunc, types, carray
-
-        signature_f32 = types.float32(
-            types.CPointer(types.float32),
-            types.CPointer(types.float32),
-            types.uint64, types.uint64)
-
-        @cfunc(signature_f32)
-        def inner_product_f32(a, b, n, m):
-            a_array = carray(a, dim)
-            b_array = carray(b, dim)
-            c = types.float32(0)
-            for i in range(dim):
-                c += a_array[i] * b_array[i]
-            return types.float32(1 - c)
-
-        signature_f8 = types.float32(
-            types.CPointer(types.int8),
-            types.CPointer(types.int8),
-            types.uint64, types.uint64)
-
-        @cfunc(signature_f8)
-        def cos_f8(a, b, n, m):
-            a_array = carray(a, dim)
-            b_array = carray(b, dim)
-            c = types.int32(0)
-            a_sq = types.int32(0)
-            b_sq = types.int32(0)
-            for i in range(dim):
-                ai = types.int16(a_array[i])
-                bi = types.int16(b_array[i])
-                c += ai * bi
-                a_sq += ai * ai
-                b_sq += bi * bi
-            a_norm = math.sqrt(types.float32(a_sq))
-            b_norm = math.sqrt(types.float32(b_sq))
-            c_norm = types.float32(c) / (a_norm * b_norm)
-            return types.float32(1 - c_norm)
-
-        print('USearch: HNSW + Numba JIT')
         index = Index(
-            ndim=dim,
+            ndim=eval.ndim,
+            jit=jit,
+            dtype=dtype,
             expansion_add=expansion_add,
             expansion_search=expansion_search,
             connectivity=connectivity,
-            metric=inner_product_f32,
+            path='USearch' + ['', '+JIT'][jit] + ':' + dtype,
         )
-        bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
+        indexes.append(index)
 
-        print('USearch: HNSW + Quarter Precision + Numba JIT')
-        index = Index(
-            ndim=dim,
-            dtype='f8',
+    # Add FAISS indexes to the mix:
+    try:
+        from index_faiss import IndexFAISS, IndexQuantizedFAISS
+        indexes.append(IndexFAISS(
+            ndim=eval.ndim,
             expansion_add=expansion_add,
             expansion_search=expansion_search,
             connectivity=connectivity,
-            metric=cos_f8,
+            path='FAISS:f32',
+        ))
+        if train:
+            indexes.append(IndexQuantizedFAISS(
+                train=eval.tasks[0].vectors,
+                expansion_add=expansion_add,
+                expansion_search=expansion_search,
+                connectivity=connectivity,
+                path='FAISS+IVFPQ:f32',
+            ))
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+    # Time to evaluate:
+    results = [eval(index) for index in indexes]
+    return pd.DataFrame({
+        'names': [i.path for i in indexes],
+        'add_per_second': [x['add_per_second'] for x in results],
+        'search_per_second': [x['search_per_second'] for x in results],
+        'recall_at_one': [x['recall_at_one'] for x in results],
+    })
+
+
+def bench_params(
+    count: int = 1_000_000,
+    connectivities: int = range(10, 20),
+    dimensions: List[int] = [
+        2, 3, 4, 8, 16, 32, 96, 100,
+        256, 384, 512, 768, 1024, 1536],
+
+    expansion_add: int = DEFAULT_EXPANSION_ADD,
+    expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+) -> pd.DataFrame:
+    """Measures indexing speed for different dimensionality vectors.
+
+    :param count: Number of vectors, defaults to 1_000_000
+    :type count: int, optional
+    """
+
+    results = []
+    for connectivity, ndim in itertools.product(connectivities, dimensions):
+
+        task = AddTask(
+            labels=np.arange(count, dtype=Label),
+            vectors=np.random.rand(count, ndim).astype(np.float32),
         )
-        bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
-    except ImportError:
-        print('Skipping Numba benchmarks')
-
-    # Don't depend on the FAISS installation for benchmarks
-    try:
-        from faiss import IndexHNSWFlat, IndexIVFPQ
-
-        print('FAISS: HNSW')
-        index = IndexHNSWFlat(dim, connectivity)
-        index.hnsw.efSearch = expansion_search
-        index.hnsw.efConstruction = expansion_add
-        bench_faiss(index, vectors_mat, queries_mat, neighbors_mat)
-
-        print('FAISS: HNSW + IVFPQ')
-        # Number of inverted lists (number of partitions or cells).
-        nlist = 10000
-        nsegment = 16  # Number of segments for PQ (number of subquantizers).
-        nbit = 8       # Number of bits to encode each segment.
-        coarse_quantizer = IndexHNSWFlat(dim, connectivity)
-        index = IndexIVFPQ(coarse_quantizer, dim, nlist, nsegment, nbit)
-        index.train(vectors_mat)
-        index.nprobe = 10
-        bench_faiss(index, vectors_mat, queries_mat, neighbors_mat)
-
-    except ImportError:
-        print('Skipping FAISS benchmarks')
-
-
-if __name__ == '__main__':
-    fire.Fire(main)
+        index = Index(
+            ndim=ndim,
+            connectivity=connectivity,
+            expansion_add=expansion_add,
+            expansion_search=expansion_search,
+        )
+        result = asdict(task(index))
+        result['ndim'] = dimensions
+        result['connectivity'] = connectivity
+        results.append(result)
+
+    # return self._execute_tasks(
+    #     tasks,
+    #     title='HNSW Indexing Speed vs Vector Dimensions',
+    #     x='ndim', y='add_per_second', log_x=True,
+    # )
+    return pd.DataFrame(results)

package/python/scripts/index_faiss.py

@@ -0,0 +1,100 @@
+import os
+from typing import Optional
+
+import numpy as np
+from faiss import IndexHNSWFlat, IndexIVFPQ, read_index
+
+from usearch.index import Matches
+from usearch.index import (
+    DEFAULT_CONNECTIVITY,
+    DEFAULT_EXPANSION_ADD,
+    DEFAULT_EXPANSION_SEARCH,
+)
+
+
+class IndexFAISS:
+
+    def __init__(
+        self,
+        index: IndexHNSWFlat = None,
+        ndim: int = 0,
+        connectivity: int = DEFAULT_CONNECTIVITY,
+        expansion_add: int = DEFAULT_EXPANSION_ADD,
+        expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+        path: Optional[os.PathLike] = None,
+        *args, **kwargs,
+    ):
+
+        if index is None:
+            index = IndexHNSWFlat(ndim, connectivity)
+            index.hnsw.efConstruction = expansion_add
+            index.hnsw.efSearch = expansion_search
+
+        self._faiss = index
+        self._specs = {
+            'Class': 'usearch.IndexFAISS',
+            'Dimensions': ndim,
+            'Connectivity': connectivity,
+            'Expansion@Add': expansion_add,
+            'Expansion@Search': expansion_search,
+        }
+
+        self.path = path
+
+    def add(self, labels, vectors):
+        # Adding labels isn't supported for most index types
+        # self._faiss.add_with_ids(vectors, labels)
+        self._faiss.add(vectors)
+
+    def search(self, queries, k: int) -> Matches:
+        distances, labels = self._faiss.search(queries, k)
+        return Matches(labels, distances, np.array([k] * queries.shape[0]))
+
+    def __len__(self) -> int:
+        return self._faiss.ntotal
+
+    def clear(self):
+        self._faiss.reset()
+
+    @property
+    def specs(self) -> dict:
+        self._specs.update({
+            'Size': len(self),
+        })
+        return self._specs
+
+    def load(self, path: os.PathLike):
+        self._faiss = read_index(path)
+
+
+class IndexQuantizedFAISS(IndexFAISS):
+
+    def __init__(
+        self,
+        train: np.ndarray,
+        connectivity: int = DEFAULT_CONNECTIVITY,
+        expansion_add: int = DEFAULT_EXPANSION_ADD,
+        expansion_search: int = DEFAULT_EXPANSION_SEARCH,
+        *args, **kwargs,
+    ):
+
+        ndim = train.shape[1]
+        super().__init__(
+            ndim=ndim,
+            connectivity=connectivity,
+            expansion_add=expansion_add,
+            expansion_search=expansion_search,
+            *args, **kwargs,
+        )
+
+        nlist = 10000  # Number of inverted lists (number of partitions).
+        nsegment = 16  # Number of segments for PQ (number of subquantizers).
+        nbit = 8       # Number of bits to encode each segment.
+
+        self._original_faiss = self._faiss
+        self._faiss = IndexIVFPQ(
+            self._faiss, ndim, nlist, nsegment, nbit)
+
+        self._faiss.train(train)
+        self._faiss.nprobe = 10
+        self._specs['Class'] = 'usearch.IndexQuantizedFAISS'