usearch 0.4.0 → 0.8.1

package/python/scripts/bench.py

@@ -0,0 +1,188 @@
+import time
+import numpy as np
+import math
+
+import fire
+
+from usearch.index import Index
+from usearch.io import load_matrix
+
+
+def measure(f) -> float:
+    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.array, queries: np.array, neighbors: np.array):
+    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.array, queries: np.array, neighbors: np.array):
+    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,
+):
+
+    vectors_mat = load_matrix(vectors)
+    queries_mat = load_matrix(queries)
+    neighbors_mat = load_matrix(neighbors)
+    dim = vectors_mat.shape[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.size_t, types.size_t)
+
+        @cfunc(signature_f32)
+        def inner_product_f32(a, b, n, m):
+            a_array = carray(a, n)
+            b_array = carray(b, n)
+            c = types.float32(0)
+            for i in range(n):
+                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.size_t, types.size_t)
+
+        @cfunc(signature_f8)
+        def cos_f8(a, b, n, m):
+            a_array = carray(a, n)
+            b_array = carray(b, n)
+            c = types.int32(0)
+            a_sq = types.int32(0)
+            b_sq = types.int32(0)
+            for i in range(n):
+                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,
+            expansion_add=expansion_add,
+            expansion_search=expansion_search,
+            connectivity=connectivity,
+            metric_pointer=inner_product_f32.address,
+        )
+        bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
+
+        print('USearch: HNSW + Quarter Precision + Numba JIT')
+        index = Index(
+            ndim=dim,
+            dtype='f8',
+            expansion_add=expansion_add,
+            expansion_search=expansion_search,
+            connectivity=connectivity,
+            metric_pointer=cos_f8.address,
+        )
+        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)

package/python/scripts/bench_params.py

@@ -0,0 +1,113 @@
+import os
+import time
+import numpy as np
+
+import fire
+import plotly.express as px
+import numpy as np
+import pandas as pd
+
+
+from usearch.index import Index
+from usearch.io import load_matrix
+
+
+def measure(f) -> float:
+    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_usearch(index, vectors: np.ndarray) -> float:
+    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]
+    insertions_per_second = vectors.shape[0]/dt
+    print(f'- Performance: {insertions_per_second:.2f} insertions/s')
+    return insertions_per_second
+
+
+class Main:
+
+    def dimensions(
+        self,
+        count: int = 1_000_000,
+        connectivity: int = 16,
+        expansion_add: int = 128,
+        expansion_search: int = 64,
+    ):
+
+        dimensions = [
+            2, 3, 4, 8, 16, 32, 96, 100,
+            256, 384, 512, 768, 1024, 1536]
+        results = []
+
+        try:
+            for ndim in dimensions:
+                vectors_mat = np.random.rand(count, ndim).astype(np.float32)
+
+                print(f'USearch for {ndim}-dimensional f32 vectors')
+                index = Index(
+                    ndim=ndim,
+                    expansion_add=expansion_add,
+                    expansion_search=expansion_search,
+                    connectivity=connectivity,
+                )
+                speed = bench_usearch(index, vectors_mat)
+                results.append({'ndim': ndim, 'speed': speed})
+                del index
+        except KeyboardInterrupt:
+            pass
+
+        df = pd.DataFrame(results)
+        fig = px.line(
+            df, x='ndim', y='speed', log_x=True,
+            title='HNSW Indexing Speed vs Vector Dimensions')
+        fig.write_image(os.path.join(os.path.dirname(
+            __file__), 'bench_dimensions.png'))
+        fig.show()
+
+    def connectivity(
+        self,
+        count: int = 1_000_000,
+        ndim: int = 256,
+        expansion_add: int = 128,
+        expansion_search: int = 64,
+    ):
+
+        connectivities = range(10, 20)
+        results = []
+
+        try:
+            for connectivity in connectivities:
+                vectors_mat = np.random.rand(count, ndim).astype(np.float32)
+
+                print(f'USearch for {connectivity} connectivity')
+                index = Index(
+                    ndim=ndim,
+                    expansion_add=expansion_add,
+                    expansion_search=expansion_search,
+                    connectivity=connectivity,
+                )
+                speed = bench_usearch(index, vectors_mat)
+                results.append({'connectivity': connectivity, 'speed': speed})
+                del index
+        except KeyboardInterrupt:
+            pass
+
+        df = pd.DataFrame(results)
+        fig = px.line(
+            df, x='connectivity', y='speed', log_x=True,
+            title='HNSW Indexing Speed vs Connectivity')
+        fig.write_image(os.path.join(os.path.dirname(
+            __file__), 'bench_connectivity.png'))
+        fig.show()
+
+
+if __name__ == '__main__':
+    fire.Fire(Main)

package/python/{test.py → scripts/test.py}

@@ -1,37 +1,57 @@
 import pytest
 import numpy as np
 
-
+from usearch.io import load_matrix
 from usearch.index import Index, SetsIndex
 
 dimensions = [13, 99, 100, 256]
 sizes = [1, 2, 100]
-types = ['f32', 'f16', 'f8', 'f64']
+index_types = ['f64', 'f32', 'f16', 'f8']
+numpy_types = [np.float64, np.float32, np.float16, np.byte]
+connectivities = [3, 13, 50]
 
 
 @pytest.mark.parametrize('ndim', dimensions)
-@pytest.mark.parametrize('dtype', types)
-def test_cos(ndim: int, dtype: str):
-    index = Index(ndim=ndim, metric='cos', dtype=dtype)
-    vector = np.random.uniform(0, 0.7, (index.ndim)).astype(np.float32)
+@pytest.mark.parametrize('index_type', index_types)
+@pytest.mark.parametrize('numpy_type', numpy_types)
+@pytest.mark.parametrize('connectivity', connectivities)
+def test_l2sq(ndim: int, index_type: str, numpy_type: str, connectivity: int):
+    index = Index(
+        metric='l2sq',
+        ndim=ndim,
+        dtype=index_type,
+        connectivity=connectivity,
+    )
+    assert index.ndim == ndim
+    assert index.connectivity == connectivity
+
+    vector = np.random.uniform(0, 0.7, (index.ndim)).astype(numpy_type)
     index.add(42, vector)
     matches, distances, count = index.search(vector, 10)
 
+    assert len(index) == 1
     assert len(matches) == count
     assert len(distances) == count
     assert count == 1
     assert matches[0] == 42
     assert distances[0] == pytest.approx(0, abs=1e-3)
 
+    index.save('tmp.usearch')
+    index.clear()
+    assert len(index) == 0
+
+    index.load('tmp.usearch')
+    assert len(index) == 1
+
 
 @pytest.mark.parametrize('ndim', dimensions)
 @pytest.mark.parametrize('size', sizes)
-@pytest.mark.parametrize('dtype', types)
-def test_cos_batch(ndim: int, size: int, dtype: str):
-    index = Index(ndim=ndim, metric='cos', dtype=dtype)
+@pytest.mark.parametrize('index_type', index_types)
+@pytest.mark.parametrize('numpy_type', numpy_types)
+def test_l2sq_batch(ndim: int, size: int, index_type: str, numpy_type: str):
+    index = Index(ndim=ndim, metric='l2sq', dtype=index_type)
 
-    vectors = np.random.uniform(
-        0, 0.3, (size, index.ndim)).astype(np.float32)
+    vectors = np.random.uniform(0, 0.7, (size, index.ndim)).astype(numpy_type)
     labels = np.array(range(size), dtype=np.longlong)
 
     index.add(labels, vectors)
@@ -84,8 +104,7 @@ def test_sets():
     index = SetsIndex()
     index.add(10, np.array([10, 12, 15], dtype=np.uint32))
     index.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
-    results = index.search(np.array([12, 15], dtype=np.uint32), 20)
-
+    # results = index.search(np.array([12, 15], dtype=np.uint32), 10)
     # assert list(results) == [10, 11]
 
 

package/python/usearch/client.py

@@ -1,34 +1,120 @@
+from typing import Union, Optional
+
 import numpy as np
 from ucall.client import Client
 
 
+def _vector_to_ascii(vector: np.ndarray) -> Optional[str]:
+    if vector.dtype != np.int8 and vector.dtype != np.uint8 and vector.dtype != np.byte:
+        return None
+    if not np.all((vector >= 0) | (vector <= 100)):
+        return None
+
+    # Let's map [0, 100] to the range from [23, 123],
+    # poking 60 and replacing with the 124.
+    vector += 23
+    vector[vector == 60] = 124
+    ascii = str(vector)
+    return ascii
+
+
 class IndexClient:
 
     def __init__(self, uri: str = '127.0.0.1', port: int = 8545, use_http: bool = True) -> None:
         self.client = Client(uri=uri, port=port, use_http=use_http)
 
-    def add(self, labels: np.array, vectors: np.array):
-        if isinstance(labels, int):
-            return self.client.add_one(label=labels, vectors=vectors)
+    def add_one(self, label: int, vector: np.ndarray):
+        assert isinstance(label, int)
+        assert isinstance(vector, np.ndarray)
+        vector = vector.flatten()
+        ascii = _vector_to_ascii(vector)
+        if ascii:
+            self.client.add_ascii(label=label, string=ascii)
         else:
-            return self.client.add_many(labels=labels, vectors=vectors)
+            self.client.add_one(label=label, vectors=vector)
+
+    def add_many(self, labels: np.ndarray, vectors: np.ndarray):
+        assert isinstance(labels, int)
+        assert isinstance(vectors, np.ndarray)
+        assert labels.ndim == 1 and vectors.ndim == 2
+        assert labels.shape[0] == vectors.shape[0]
+        self.client.add_many(labels=labels, vectors=vectors)
+
+    def add(self, labels: Union[np.ndarray, int], vectors: np.ndarray):
+        if isinstance(labels, int) or len(labels) == 1:
+            return self.add_one(labels, vectors)
+        else:
+            return self.add_many(labels, vectors)
+
+    def search_one(self, vector: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        matches: list[dict] = []
+        vector = vector.flatten()
+        ascii = _vector_to_ascii(vector)
+        if ascii:
+            matches = self.client.search_ascii(string=ascii, count=count)
+        else:
+            matches = self.client.search_one(vector=vector, count=count)
+
+        print(matches.data)
+        matches = matches.json
+
+        labels = np.array((1, count), dtype=np.uint32)
+        distances = np.array((1, count), dtype=np.float32)
+        counts = np.array((1), dtype=np.uint32)
+        for col, result in enumerate(matches):
+            labels[0, col] = result['label']
+            distances[0, col] = result['distance']
+        counts[0] = len(matches)
 
-    def search(self, vectors: np.array, count: int) -> tuple[np.array, np.array, np.array]:
-        matches = []
-        distances = []
-        counts = []
-        # return self.client.search_one(vectors=vectors, count=count)
-        return matches, distances, counts
+        return labels, distances, counts
+
+    def search_many(self, vectors: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        batch_size: int = vectors.shape[0]
+        list_of_matches: list[list[dict]] = self.client.search_many(
+            vectors=vectors, count=count)
+
+        labels = np.array((batch_size, count), dtype=np.uint32)
+        distances = np.array((batch_size, count), dtype=np.float32)
+        counts = np.array((batch_size), dtype=np.uint32)
+        for row, matches in enumerate(list_of_matches):
+            for col, result in enumerate(matches):
+                labels[row, col] = result['label']
+                distances[row, col] = result['distance']
+            counts[row] = len(results)
+
+        return labels, distances, counts
+
+    def search(self, vectors: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        if vectors.ndim == 1 or (vectors.ndim == 2 and vectors.shape[0] == 1):
+            return self.search_one(vectors, count)
+        else:
+            return self.search_many(vectors, count)
 
     def __len__(self):
-        return self.client.size()
+        return self.client.size().json()
 
     @property
     def ndim(self):
-        return self.client.ndim()
+        return self.client.ndim().json()
 
     def capacity(self):
-        return self.client.capacity()
+        return self.client.capacity().json()
 
     def connectivity(self):
-        return self.client.connectivity()
+        return self.client.connectivity().json()
+
+    def load(self, path: str):
+        raise NotImplementedError()
+
+    def view(self, path: str):
+        raise NotImplementedError()
+
+    def save(self, path: str):
+        raise NotImplementedError()
+
+
+if __name__ == '__main__':
+    index = IndexClient()
+    index.add(42, np.array([0.4] * 256, dtype=np.float32))
+    results = index.search(np.array([0.4] * 256, dtype=np.float32), 10)
+    print(results)

package/python/usearch/server.py

@@ -1,32 +1,49 @@
-import ucall.rich_posix as ucall
-import usearch
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
 
-import fire
+import os
+import argparse
 import numpy as np
-from PIL import Image
 
+from ucall.rich_posix import Server
+from usearch.index import Index
 
-def serve(ndim: int, metric: str = 'ip'):
 
-    server = ucall.Server()
-    index = usearch.Index(ndim=ndim, metric=metric)
+def _results_to_json(results: tuple[np.ndarray, np.ndarray, np.ndarray], row: int) -> list[dict]:
+    count = results[2][row]
+    labels = results[0][row, :count]
+    distances = results[1][row, :count]
+    return [{'label': int(l), 'distance': float(d)} for l, d in zip(labels, distances)]
 
-    @server
-    def add_one(label: int, vector: np.array):
-        labels = np.array([label], dtype=np.longlong)
-        vectors = vector.reshape(vector.shape[0], 1)
-        index.add(labels, vectors, copy=True)
 
-    @server
-    def add_many(labels: np.array, vectors: np.array):
-        labels = labels.astype(np.longlong)
-        index.add(labels, vectors, copy=True)
+def _ascii_to_vector(string: str) -> np.ndarray:
+    """
+    WARNING: A dirty performance hack!
+    Assuming the `f8` vectors in our implementations are just integers,
+    and generally contain scalars in the [0, 100] range, we can transmit
+    them as JSON-embedded strings. The only symbols we must avoid are
+    the double-quote '"' (code 22) and backslash '\' (code 60).
+    Printable ASCII characters are in [20, 126].
+    """
+    vector = np.array(string, dtype=np.int8)
+    vector[vector == 124] = 60
+    vector -= 23
+    return vector
 
-    @server
-    def search_one(vector: np.array, count: int) -> np.ndarray:
-        vectors = vector.reshape(vector.shape[0], 1)
-        results = index.search(vectors, 3)
-        return results[0][:results[2][0]]
+
+def serve(
+        ndim_: int, metric: str = 'ip',
+        port: int = 8545, threads: int = 1,
+        path: str = 'index.usearch', immutable: bool = False):
+
+    server = Server(port=port)
+    index = Index(ndim=ndim_, metric=metric)
+
+    if os.path.exists(path):
+        if immutable:
+            index.view(path)
+        else:
+            index.load(path)
 
     @server
     def size() -> int:
@@ -44,8 +61,72 @@ def serve(ndim: int, metric: str = 'ip'):
     def connectivity() -> int:
         return index.connectivity()
 
-    server.run()
+    @server
+    def add_one(label: int, vector: np.ndarray):
+        print('adding', label, vector)
+        labels = np.array([label], dtype=np.longlong)
+        vectors = vector.flatten().reshape(vector.shape[0], 1)
+        index.add(labels, vectors)
+
+    @server
+    def add_many(labels: np.ndarray, vectors: np.ndarray):
+        labels = labels.astype(np.longlong)
+        index.add(labels, vectors, threads=threads)
+
+    @server
+    def search_one(vector: np.ndarray, count: int) -> list[dict]:
+        print('search', vector, count)
+        vectors = vector.reshape(vector.shape[0], 1)
+        results = index.search(vectors, count)
+        return _results_to_json(results, 0)
+
+    @server
+    def search_many(vectors: np.ndarray, count: int) -> list[list[dict]]:
+        results = index.search(vectors, count)
+        return [_results_to_json(results, i) for i in range(vectors.shape[0])]
+
+    @server
+    def add_ascii(label: int, string: str):
+        return add_one(label, _ascii_to_vector(string))
+
+    @server
+    def search_ascii(string: str, count: int):
+        return search_one(_ascii_to_vector(string), count)
+
+    try:
+        server.run()
+    except KeyboardInterrupt:
+        if not immutable:
+            index.save(path)
 
 
 if __name__ == '__main__':
-    fire(serve)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-v', '--verbose', help='log server activity')
+    parser.add_argument(
+        '--ndim', type=int,
+        help='dimensionality of the vectors')
+    parser.add_argument(
+        '--immutable', type=bool, default=False,
+        help='the index can not be updated')
+
+    parser.add_argument(
+        '--metric', type=str, default='ip', choices=['ip', 'cos', 'l2sq', 'haversine'],
+        help='distance function to compare vectors')
+    parser.add_argument(
+        '-p', '--port', type=int, default=8545,
+        help='port to open for client connections')
+    parser.add_argument(
+        '-j', '--threads', type=int, default=1,
+        help='number of CPU threads to use')
+    parser.add_argument(
+        '--path', type=str, default='index.usearch',
+        help='where to store the index')
+
+    args = parser.parse_args()
+    assert args.ndim is not None, 'Define the number of dimensions!'
+    serve(
+        ndim_=args.ndim, metric=args.metric,
+        threads=args.threads, port=args.port,
+        path=args.path, immutable=args.immutable)