usearch 2.8.7 → 2.8.9

package/python/scripts/bench.ipynb

@@ -1,193 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install faiss-cpu numba pandas plotly nbformat jinja2 kaleido tqdm scikit-learn scipy matplotlib"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from bench import *\n",
-    "import plotly.graph_objects as go"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "('auto', 'auto', 'auto')"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from usearch.index import Index\n",
-    "\n",
-    "Index(ndim=256, metric=\"cos\", dtype=\"f32\").hardware_acceleration, Index(ndim=256, metric=\"cos\", dtype=\"f16\").hardware_acceleration, Index(ndim=256, metric=\"cos\", dtype=\"i8\").hardware_acceleration"
-   ]
-  },
-  {
-   "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": [
-    "from usearch.eval import Dataset\n",
-    "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, jit=True)\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, jit=True)\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.5"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

package/python/scripts/bench.py

@@ -1,137 +0,0 @@
-import itertools
-from typing import List
-from dataclasses import asdict
-
-import numpy as np
-import pandas as pd
-
-from usearch.index import Index, Key, MetricKind, ScalarKind
-from usearch.numba import jit as njit
-from usearch.eval import Evaluation, 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,
-    jit: bool = False,
-    train: bool = False,
-) -> pd.DataFrame:
-    # Build various indexes:
-    indexes = []
-    jit_options = [False, True] if jit else [False]
-    dtype_options = [ScalarKind.F32, ScalarKind.F16, ScalarKind.I8]
-    for jit, dtype in itertools.product(jit_options, dtype_options):
-        metric = MetricKind.IP
-        if jit:
-            metric = njit(eval.ndim, metric, dtype)
-        index = Index(
-            ndim=eval.ndim,
-            metric=metric,
-            dtype=dtype,
-            expansion_add=expansion_add,
-            expansion_search=expansion_search,
-            connectivity=connectivity,
-            path="USearch" + ["", "+JIT"][jit] + ":" + str(dtype),
-        )
-
-        # Skip the cases, where JIT-ing is impossible
-        if jit and not index.jit:
-            continue
-        indexes.append(index)
-
-    # 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,
-                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(
-            keys=np.arange(count, dtype=Key),
-            vectors=np.random.rand(count, ndim).astype(np.float32),
-        )
-        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/bench_exact.py

@@ -1,44 +0,0 @@
-from time import time
-
-from usearch.index import search, MetricKind, ScalarKind
-from usearch.compiled import hardware_acceleration
-from faiss import knn, METRIC_L2
-import numpy as np
-import fire
-
-
-def run(
-    n: int = 10**5,
-    q: int = 10,
-    k: int = 100,
-    ndim: int = 256,
-    half: bool = False,
-):
-    usearch_dtype = ScalarKind.F16 if half else ScalarKind.F32
-    acceleration = hardware_acceleration(
-        dtype=usearch_dtype,
-        ndim=ndim,
-        metric_kind=MetricKind.L2sq,
-    )
-    print("Hardware acceleration in USearch: ", acceleration)
-
-    x = np.random.random((n, ndim))
-    x = x.astype(np.float16) if half else x.astype(np.float32)
-
-    start = time()
-    _ = search(
-        x,
-        x[:q],
-        k,
-        MetricKind.L2sq,
-        exact=True,
-    ).keys
-    print("USearch: ", time() - start)
-
-    start = time()
-    _ = knn(x[:q], x, k, metric=METRIC_L2)[1]
-    print("FAISS:   ", time() - start)
-
-
-if __name__ == "__main__":
-    fire.Fire(run)

package/python/scripts/cluster.ipynb

@@ -1,160 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from usearch.io import load_matrix\n",
-    "import numpy as np"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vectors = load_matrix(\"datasets/arxiv_2M/abstract.e5-base-v2.fbin\", view=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vectors.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from usearch.index import Index, USES_SIMSIMD, USES_NATIVE_F16\n",
-    "\n",
-    "index = Index(ndim=vectors.shape[1], metric=\"cos\", dtype=\"i8\")\n",
-    "index.hardware_acceleration, USES_SIMSIMD, USES_NATIVE_F16"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "import os\n",
-    "\n",
-    "if os.path.exists(\"cluster.i8.usearch\"):\n",
-    "    index.load(\"cluster.i8.usearch\")\n",
-    "\n",
-    "if len(index) == 0:\n",
-    "    index.add(None, vectors, log=True)\n",
-    "    index.save(\"cluster.i8.usearch\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "index"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clustering = index.cluster(min_count=10, max_count=15, threads=60)\n",
-    "clustering"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(f\"{len(clustering.centroids_popularity[0])} unique clusters for {len(index)} members\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clustering.plot_centroids_popularity()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "nx.draw_networkx(g)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import networkx as nx\n",
-    "\n",
-    "G = g\n",
-    "\n",
-    "node_sizes = nx.get_node_attributes(G, \"size\").values()\n",
-    "max_node_size = max(node_sizes)\n",
-    "node_sizes = [size * 10 / max_node_size for size in node_sizes]\n",
-    "\n",
-    "edge_labels = nx.get_edge_attributes(G, \"distance\")\n",
-    "edge_labels = {edge: f\"{label:.2}\" for edge, label in edge_labels.items()}\n",
-    "\n",
-    "pos = nx.spring_layout(G, seed=7)\n",
-    "nx.draw_networkx_nodes(G, pos, node_size=node_sizes)\n",
-    "nx.draw_networkx_edges(G, pos, edgelist=G.edges(data=False))\n",
-    "nx.draw_networkx_labels(G, pos, font_size=10, font_family=\"sans-serif\")\n",
-    "nx.draw_networkx_edge_labels(G, pos, edge_labels, font_size=5)\n",
-    "\n",
-    "ax = plt.gca()\n",
-    "ax.margins(0.08)\n",
-    "plt.axis(\"off\")\n",
-    "plt.tight_layout()\n",
-    "plt.show()"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "base",
-   "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.10.11"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

package/python/scripts/index_faiss.py

@@ -1,100 +0,0 @@
-import os
-from typing import Optional
-
-import numpy as np
-from faiss import IndexHNSWFlat, IndexIVFPQ, read_index
-
-from usearch.index import Matches, BatchMatches
-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, keys, vectors):
-        # Adding keys isn't supported for most index types
-        # self._faiss.add_with_ids(vectors, keys)
-        self._faiss.add(vectors)
-
-    def search(self, queries, k: int) -> BatchMatches:
-        distances, keys = self._faiss.search(queries, k)
-        return BatchMatches(keys, 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"