usearch 2.23.0__cp314-cp314t-macosx_11_0_arm64.whl

usearch/io.py

@@ -0,0 +1,138 @@
+import os
+import struct
+import typing
+
+import numpy as np
+
+
+def numpy_scalar_size(dtype) -> int:
+    return {
+        np.float64: 8,
+        np.int64: 8,
+        np.uint64: 8,
+        np.float32: 4,
+        np.int32: 4,
+        np.uint32: 4,
+        np.float16: 2,
+        np.int16: 2,
+        np.uint16: 2,
+        np.int8: 1,
+        np.uint8: 1,
+    }[dtype]
+
+
+def guess_numpy_dtype_from_filename(filename) -> typing.Optional[type]:
+    if filename.endswith(".fbin"):
+        return np.float32
+    elif filename.endswith(".dbin"):
+        return np.float64
+    elif filename.endswith(".hbin"):
+        return np.float16
+    elif filename.endswith(".ibin"):
+        return np.int32
+    elif filename.endswith(".bbin"):
+        return np.uint8
+    elif filename.endswith(".i8bin"):
+        return np.int8
+    elif filename.endswith(".i32bin"):
+        return np.int32
+    elif filename.endswith(".f32bin"):
+        return np.float32
+    else:
+        return None
+
+
+def load_matrix(
+    filename: str,
+    start_row: int = 0,
+    count_rows: int = None,
+    view: bool = False,
+    dtype: typing.Optional[type] = None,
+) -> typing.Optional[np.ndarray]:
+    """Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin, *.i8bin, *.i32bin files with matrices.
+
+    :param filename: path to the matrix file
+    :param start_row: start reading vectors from this index
+    :param count_rows: number of vectors to read. If None, read all vectors
+    :param view: set to `True` to memory-map the file instead of loading to RAM
+
+    :return: parsed matrix
+    :rtype: numpy.ndarray
+    """
+    if dtype is None:
+        dtype = guess_numpy_dtype_from_filename(filename)
+        if dtype is None:
+            raise Exception("Unknown file type")
+    scalar_size = numpy_scalar_size(dtype)
+
+    if not os.path.exists(filename):
+        return None
+
+    with open(filename, "rb") as f:
+        rows, cols = np.fromfile(f, count=2, dtype=np.int32).astype(np.uint64)
+        
+        # Validate file size matches expected data size
+        f.seek(0, 2)  # Go to end
+        file_size = f.tell()
+        expected_size = 8 + (rows * cols * scalar_size)  # Header + data
+        
+        if file_size != expected_size:
+            if file_size < expected_size:
+                raise ValueError(f"File {filename} is truncated. Expected {expected_size:,} bytes, got {file_size:,} bytes")
+            else:
+                raise ValueError(f"File {filename} is larger than expected. Expected {expected_size:,} bytes, got {file_size:,} bytes")
+        
+        f.seek(8)  # Back to start of data
+        rows = (rows - start_row) if count_rows is None else count_rows
+        row_offset = start_row * scalar_size * cols
+
+        if view:
+            return np.memmap(
+                f,
+                dtype=dtype,
+                mode="r",
+                offset=8 + row_offset,
+                shape=(rows, cols),
+            )
+        else:
+            return np.fromfile(
+                f,
+                count=rows * cols,
+                dtype=dtype,
+                offset=row_offset,
+            ).reshape(rows, cols)
+
+
+def save_matrix(vectors: np.ndarray, filename: str):
+    """Write *.ibin, *.bbib, *.hbin, *.fbin, *.dbin, *.i8bin, *.i32bin, *.f32bin files with matrices.
+
+    :param vectors: the matrix to serialize
+    :type vectors: numpy.ndarray
+    :param filename: path to the matrix file
+    :type filename: str
+    """
+    if filename.endswith(".fbin"):
+        dtype = np.float32
+    elif filename.endswith(".dbin"):
+        dtype = np.float64
+    elif filename.endswith(".hbin"):
+        dtype = np.float16
+    elif filename.endswith(".ibin"):
+        dtype = np.int32
+    elif filename.endswith(".bbin"):
+        dtype = np.uint8
+    elif filename.endswith(".i8bin"):
+        dtype = np.int8
+    elif filename.endswith(".i32bin"):
+        dtype = np.int32
+    elif filename.endswith(".f32bin"):
+        dtype = np.float32
+    else:
+        dtype = vectors.dtype
+
+    assert len(vectors.shape) == 2, "Input array must have 2 dimensions"
+    with open(filename, "wb") as f:
+        count, dim = vectors.shape
+        f.write(struct.pack("<i", count))
+        f.write(struct.pack("<i", dim))
+        vectors.astype(dtype).flatten().tofile(f)

usearch/numba.py

@@ -0,0 +1,110 @@
+# The purpose of this file is to provide Pythonic wrapper on top
+# the native precompiled CPython module. It improves compatibility
+# Python tooling, linters, and static analyzers. It also embeds JIT
+# into the primary `Index` class, connecting USearch with Numba.
+from math import sqrt
+
+from usearch.index import MetricKind, ScalarKind, MetricSignature, CompiledMetric
+
+
+def jit(
+    ndim: int,
+    metric: MetricKind = MetricKind.Cos,
+    dtype: ScalarKind = ScalarKind.F32,
+) -> CompiledMetric:
+    """JIT-compiles the metric for target hardware and number of dimensions.
+
+    This can result in up-to 3x performance difference on very large vectors
+    and very recent hardware, as the Python module is compiled with high
+    compatibility in mind and avoids very fancy assembly instructions.
+
+    Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
+    of `ctypes` to support half-precision.
+    https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
+    """
+    assert isinstance(metric, MetricKind)
+    assert isinstance(dtype, ScalarKind)
+
+    from numba import cfunc, types, carray
+
+    signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8))
+    signature_f16args = types.float32(types.CPointer(types.float16), types.CPointer(types.float16))
+    signature_f32args = types.float32(types.CPointer(types.float32), types.CPointer(types.float32))
+    signature_f64args = types.float32(types.CPointer(types.float64), types.CPointer(types.float64))
+
+    numba_supported_types = (
+        ScalarKind.I8,
+        # Half-precision is still unsupported
+        # https://github.com/numba/numba/issues/4402
+        # ScalarKind.F16: np.float16,
+        ScalarKind.F32,
+        ScalarKind.F64,
+    )
+    if dtype not in numba_supported_types:
+        return metric
+
+    scalar_kind_to_accumulator_type = {
+        ScalarKind.I8: types.int32,
+        ScalarKind.F16: types.float16,
+        ScalarKind.F32: types.float32,
+        ScalarKind.F64: types.float64,
+    }
+    accumulator = scalar_kind_to_accumulator_type[dtype]
+
+    def numba_ip(a, b):
+        a_array = carray(a, ndim)
+        b_array = carray(b, ndim)
+        ab = accumulator(0)
+        for i in range(ndim):
+            ab += a_array[i] * b_array[i]
+        return types.float32(1 - ab)
+
+    def numba_cos(a, b):
+        a_array = carray(a, ndim)
+        b_array = carray(b, ndim)
+        ab = accumulator(0)
+        a_sq = accumulator(0)
+        b_sq = accumulator(0)
+        for i in range(ndim):
+            ab += a_array[i] * b_array[i]
+            a_sq += a_array[i] * a_array[i]
+            b_sq += b_array[i] * b_array[i]
+        a_norm = sqrt(a_sq)
+        b_norm = sqrt(b_sq)
+        if a_norm == 0 and b_norm == 0:
+            return types.float32(0)
+        elif a_norm == 0 or b_norm == 0 or ab == 0:
+            return types.float32(1)
+        else:
+            return types.float32(1 - ab / (a_norm * b_norm))
+
+    def numba_l2sq(a, b):
+        a_array = carray(a, ndim)
+        b_array = carray(b, ndim)
+        ab_delta_sq = accumulator(0)
+        for i in range(ndim):
+            ab_delta_sq += (a_array[i] - b_array[i]) * (a_array[i] - b_array[i])
+        return types.float32(ab_delta_sq)
+
+    scalar_kind_to_signature = {
+        ScalarKind.I8: signature_i8args,
+        ScalarKind.F16: signature_f16args,
+        ScalarKind.F32: signature_f32args,
+        ScalarKind.F64: signature_f64args,
+    }
+
+    metric_kind_to_function = {
+        MetricKind.IP: numba_ip,
+        MetricKind.Cos: numba_cos,
+        MetricKind.L2sq: numba_l2sq,
+    }
+
+    if dtype == ScalarKind.I8 and metric == MetricKind.IP:
+        metric = MetricKind.Cos
+
+    pointer = cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
+    return CompiledMetric(
+        pointer=pointer.address,
+        kind=metric,
+        signature=MetricSignature.ArrayArray,
+    )

usearch/server.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import os
+import argparse
+import numpy as np
+from typing import List
+
+from ucall.rich_posix import Server
+from usearch.index import Index, Matches, Key
+
+
+def _ascii_to_vector(string: str) -> np.ndarray:
+    """
+    WARNING: A dirty performance hack!
+    Assuming the `i8` 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
+
+
+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:
+        return len(index)
+
+    @server
+    def ndim() -> int:
+        return index.ndim
+
+    @server
+    def capacity() -> int:
+        return index.capacity()
+
+    @server
+    def connectivity() -> int:
+        return index.connectivity()
+
+    @server
+    def add_one(key: int, vector: np.ndarray):
+        print("adding", key, vector)
+        keys = np.array([key], dtype=Key)
+        vectors = vector.flatten().reshape(vector.shape[0], 1)
+        index.add(keys, vectors)
+
+    @server
+    def add_many(keys: np.ndarray, vectors: np.ndarray):
+        index.add(keys, 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: Matches = index.search(vectors, count)
+        return results.to_list()
+
+    @server
+    def search_many(vectors: np.ndarray, count: int) -> List[List[dict]]:
+        results: Matches = index.search(vectors, count)
+        return results.to_list()
+
+    @server
+    def add_ascii(key: int, string: str):
+        return add_one(key, _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__":
+    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,
+    )