usearch 0.13.0 → 0.15.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.github/workflows/release.yml +62 -7
- package/.github/workflows/update_version.sh +6 -1
- package/.vscode/launch.json +21 -22
- package/.vscode/settings.json +7 -1
- package/CITATION.cff +4 -4
- package/CMakeLists.txt +1 -1
- package/Cargo.toml +1 -1
- package/Package.swift +0 -1
- package/README.md +100 -27
- package/binding.gyp +0 -1
- package/build.gradle +1 -1
- package/build.rs +3 -3
- package/c/lib.cpp +32 -9
- package/c/usearch.h +67 -36
- package/conanfile.py +1 -1
- package/{src → cpp}/bench.cpp +52 -36
- package/{src → cpp}/test.cpp +17 -7
- package/docs/benchmarks.md +33 -10
- package/docs/c/index.rst +11 -0
- package/docs/c/reference.rst +7 -0
- package/docs/compilation.md +24 -8
- package/docs/conf.dox +6 -6
- package/docs/conf.py +3 -1
- package/docs/cpp/index.rst +11 -0
- package/docs/cpp/reference.rst +7 -0
- package/docs/index.rst +10 -0
- package/docs/java/index.rst +11 -0
- package/docs/java/reference.rst +7 -0
- package/docs/javascript/index.rst +11 -0
- package/docs/javascript/reference.rst +9 -0
- package/docs/python/index.rst +11 -0
- package/docs/python/reference.rst +35 -0
- package/golang/lib.go +1 -1
- package/include/usearch/{usearch.hpp → index.hpp} +152 -148
- package/include/usearch/index_punned_dense.hpp +633 -0
- package/include/usearch/index_punned_helpers.hpp +528 -0
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +8 -7
- package/javascript/lib.cpp +26 -50
- package/javascript/usearch.js +117 -2
- package/objc/USearchObjective.mm +145 -49
- package/objc/include/USearchObjective.h +106 -30
- package/package.json +2 -2
- package/python/lib.cpp +178 -320
- package/python/scripts/bench.ipynb +170 -0
- package/python/scripts/bench.py +102 -185
- package/python/scripts/index_faiss.py +100 -0
- package/python/scripts/test.py +57 -34
- package/python/usearch/eval.py +353 -0
- package/python/usearch/index.py +138 -117
- package/python/usearch/io.py +7 -2
- package/python/usearch/join.py +100 -0
- package/python/usearch/numba.py +101 -0
- package/rust/lib.cpp +27 -14
- package/rust/lib.hpp +4 -10
- package/rust/lib.rs +50 -17
- package/setup.py +12 -2
- package/simsimd/include/simsimd/simsimd.h +32 -3
- package/swift/Index+Sugar.swift +55 -29
- package/swift/Test.swift +1 -1
- package/wolfram/lib.cpp +4 -4
- package/python/scripts/bench_params.py +0 -161
- package/python/usearch/synthetic.py +0 -32
- package/src/punned.hpp +0 -887
- /package/{src → cpp}/CMakeLists.txt +0 -0
package/python/usearch/index.py
CHANGED
|
@@ -4,28 +4,60 @@
|
|
|
4
4
|
# into the primary `Index` class, connecting USearch with Numba.
|
|
5
5
|
import os
|
|
6
6
|
from math import sqrt
|
|
7
|
-
from typing import Optional, Callable, Union, NamedTuple, List
|
|
7
|
+
from typing import Optional, Callable, Union, NamedTuple, List, Iterable
|
|
8
8
|
|
|
9
9
|
import numpy as np
|
|
10
10
|
|
|
11
11
|
from usearch.compiled import Index as _CompiledIndex
|
|
12
|
-
from usearch.compiled import
|
|
13
|
-
from usearch.compiled import BitsIndex as _CompiledBitsIndex
|
|
12
|
+
from usearch.compiled import SparseIndex as _CompiledSetsIndex
|
|
14
13
|
|
|
15
|
-
from usearch.compiled import MetricKind
|
|
14
|
+
from usearch.compiled import MetricKind, ScalarKind
|
|
16
15
|
from usearch.compiled import (
|
|
17
16
|
DEFAULT_CONNECTIVITY,
|
|
18
17
|
DEFAULT_EXPANSION_ADD,
|
|
19
18
|
DEFAULT_EXPANSION_SEARCH,
|
|
19
|
+
|
|
20
|
+
USES_OPENMP,
|
|
21
|
+
USES_SIMSIMD,
|
|
20
22
|
)
|
|
21
23
|
|
|
22
|
-
|
|
23
|
-
MetricKind.
|
|
24
|
-
MetricKind.
|
|
25
|
-
MetricKind.
|
|
24
|
+
MetricKindBitwise = (
|
|
25
|
+
MetricKind.Hamming,
|
|
26
|
+
MetricKind.Tanimoto,
|
|
27
|
+
MetricKind.Sorensen,
|
|
26
28
|
)
|
|
27
29
|
|
|
28
|
-
|
|
30
|
+
SparseIndex = _CompiledSetsIndex
|
|
31
|
+
|
|
32
|
+
Label = np.longlong
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _normalize_dtype(dtype) -> ScalarKind:
|
|
36
|
+
if dtype is None:
|
|
37
|
+
return ScalarKind.F32
|
|
38
|
+
if isinstance(dtype, ScalarKind):
|
|
39
|
+
return dtype
|
|
40
|
+
|
|
41
|
+
if isinstance(dtype, str):
|
|
42
|
+
_normalize = {
|
|
43
|
+
'f64': ScalarKind.F64,
|
|
44
|
+
'f32': ScalarKind.F32,
|
|
45
|
+
'f16': ScalarKind.F16,
|
|
46
|
+
'f8': ScalarKind.F8,
|
|
47
|
+
'b1': ScalarKind.B1,
|
|
48
|
+
}
|
|
49
|
+
return _normalize[dtype.lower()]
|
|
50
|
+
|
|
51
|
+
if not isinstance(dtype, ScalarKind):
|
|
52
|
+
_normalize = {
|
|
53
|
+
np.float64: ScalarKind.F64,
|
|
54
|
+
np.float32: ScalarKind.F32,
|
|
55
|
+
np.float16: ScalarKind.F16,
|
|
56
|
+
np.int8: ScalarKind.F8,
|
|
57
|
+
}
|
|
58
|
+
return _normalize[dtype]
|
|
59
|
+
|
|
60
|
+
return dtype
|
|
29
61
|
|
|
30
62
|
|
|
31
63
|
class Matches(NamedTuple):
|
|
@@ -65,83 +97,15 @@ class Matches(NamedTuple):
|
|
|
65
97
|
for label, distance in zip(labels, distances)
|
|
66
98
|
]
|
|
67
99
|
|
|
100
|
+
def recall_first(self, expected: np.ndarray) -> float:
|
|
101
|
+
best_matches = self.labels if not self.is_batch else self.labels[:, 0]
|
|
102
|
+
return np.sum(best_matches == expected) / len(expected)
|
|
103
|
+
|
|
68
104
|
def __repr__(self) -> str:
|
|
69
105
|
return f'usearch.Matches({self.total_matches})' if self.is_batch else \
|
|
70
106
|
f'usearch.Matches({self.total_matches} across {self.batch_size} queries)'
|
|
71
107
|
|
|
72
108
|
|
|
73
|
-
def jit_metric(ndim: int, metric: MetricKind, dtype: str = 'f32') -> Callable:
|
|
74
|
-
|
|
75
|
-
try:
|
|
76
|
-
from numba import cfunc, types, carray
|
|
77
|
-
except ImportError:
|
|
78
|
-
raise ModuleNotFoundError(
|
|
79
|
-
'To use JIT install Numba with `pip install numba`.'
|
|
80
|
-
'Alternatively, reinstall usearch with `pip install usearch[jit]`')
|
|
81
|
-
|
|
82
|
-
# Showcases how to use Numba to JIT-compile similarity measures for USearch.
|
|
83
|
-
# https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
|
|
84
|
-
|
|
85
|
-
if dtype == 'f32':
|
|
86
|
-
signature = types.float32(
|
|
87
|
-
types.CPointer(types.float32),
|
|
88
|
-
types.CPointer(types.float32),
|
|
89
|
-
types.uint64, types.uint64)
|
|
90
|
-
|
|
91
|
-
if metric == MetricKind.IP:
|
|
92
|
-
|
|
93
|
-
@cfunc(signature)
|
|
94
|
-
def numba_ip(a, b, _n, _m):
|
|
95
|
-
a_array = carray(a, ndim)
|
|
96
|
-
b_array = carray(b, ndim)
|
|
97
|
-
ab = 0.0
|
|
98
|
-
for i in range(ndim):
|
|
99
|
-
ab += a_array[i] * b_array[i]
|
|
100
|
-
return 1 - ab
|
|
101
|
-
|
|
102
|
-
return numba_ip
|
|
103
|
-
|
|
104
|
-
if metric == MetricKind.Cos:
|
|
105
|
-
|
|
106
|
-
@cfunc(signature)
|
|
107
|
-
def numba_cos(a, b, _n, _m):
|
|
108
|
-
a_array = carray(a, ndim)
|
|
109
|
-
b_array = carray(b, ndim)
|
|
110
|
-
ab = 0.0
|
|
111
|
-
a_sq = 0.0
|
|
112
|
-
b_sq = 0.0
|
|
113
|
-
for i in range(ndim):
|
|
114
|
-
ab += a_array[i] * b_array[i]
|
|
115
|
-
a_sq += a_array[i] * a_array[i]
|
|
116
|
-
b_sq += b_array[i] * b_array[i]
|
|
117
|
-
a_norm = sqrt(a_sq)
|
|
118
|
-
b_norm = sqrt(b_sq)
|
|
119
|
-
if a_norm == 0 and b_norm == 0:
|
|
120
|
-
return 0
|
|
121
|
-
elif a_norm == 0 or b_norm == 0 or ab == 0:
|
|
122
|
-
return 1
|
|
123
|
-
else:
|
|
124
|
-
return 1 - ab / (a_norm * b_norm)
|
|
125
|
-
|
|
126
|
-
return numba_cos
|
|
127
|
-
|
|
128
|
-
if metric == MetricKind.L2sq:
|
|
129
|
-
|
|
130
|
-
@cfunc(signature)
|
|
131
|
-
def numba_l2sq(a, b, _n, _m):
|
|
132
|
-
a_array = carray(a, ndim)
|
|
133
|
-
b_array = carray(b, ndim)
|
|
134
|
-
ab_delta_sq = types.float32(0.0)
|
|
135
|
-
for i in range(ndim):
|
|
136
|
-
ab_delta_sq += (a_array[i] - b_array[i]) * \
|
|
137
|
-
(a_array[i] - b_array[i])
|
|
138
|
-
return ab_delta_sq
|
|
139
|
-
|
|
140
|
-
return numba_l2sq
|
|
141
|
-
|
|
142
|
-
return None
|
|
143
|
-
|
|
144
|
-
|
|
145
109
|
class Index:
|
|
146
110
|
"""Fast JIT-compiled vector-search index for dense equi-dimensional embeddings.
|
|
147
111
|
|
|
@@ -155,7 +119,7 @@ class Index:
|
|
|
155
119
|
def __init__(
|
|
156
120
|
self,
|
|
157
121
|
ndim: int,
|
|
158
|
-
metric: Union[MetricKind, Callable] = MetricKind.IP,
|
|
122
|
+
metric: Union[MetricKind, Callable, str] = MetricKind.IP,
|
|
159
123
|
dtype: Optional[str] = None,
|
|
160
124
|
jit: bool = False,
|
|
161
125
|
|
|
@@ -173,11 +137,11 @@ class Index:
|
|
|
173
137
|
:type ndim: int
|
|
174
138
|
|
|
175
139
|
:param metric: Distance function, defaults to MetricKind.IP
|
|
176
|
-
:type metric: Union[MetricKind, Callable], optional
|
|
140
|
+
:type metric: Union[MetricKind, Callable, str], optional
|
|
177
141
|
Kind of the distance function, or the Numba `cfunc` JIT-compiled object.
|
|
178
142
|
Possible `MetricKind` values: IP, Cos, L2sq, Haversine, Pearson,
|
|
179
|
-
|
|
180
|
-
Not every kind is JIT-able. For Jaccard distance, use `
|
|
143
|
+
Hamming, Tanimoto, Sorensen.
|
|
144
|
+
Not every kind is JIT-able. For Jaccard distance, use `SparseIndex`.
|
|
181
145
|
|
|
182
146
|
:param dtype: Scalar type for internal vector storage, defaults to None
|
|
183
147
|
:type dtype: str, optional
|
|
@@ -221,6 +185,18 @@ class Index:
|
|
|
221
185
|
|
|
222
186
|
if metric is None:
|
|
223
187
|
metric = MetricKind.IP
|
|
188
|
+
elif isinstance(metric, str):
|
|
189
|
+
_normalize = {
|
|
190
|
+
'cos': MetricKind.Cos,
|
|
191
|
+
'ip': MetricKind.IP,
|
|
192
|
+
'l2_sq': MetricKind.L2sq,
|
|
193
|
+
'haversine': MetricKind.Haversine,
|
|
194
|
+
'perason': MetricKind.Pearson,
|
|
195
|
+
'hamming': MetricKind.Hamming,
|
|
196
|
+
'tanimoto': MetricKind.Tanimoto,
|
|
197
|
+
'sorensen': MetricKind.Sorensen,
|
|
198
|
+
}
|
|
199
|
+
metric = _normalize[metric.lower()]
|
|
224
200
|
|
|
225
201
|
if isinstance(metric, Callable):
|
|
226
202
|
self._metric_kind = MetricKind.Unknown
|
|
@@ -229,10 +205,18 @@ class Index:
|
|
|
229
205
|
|
|
230
206
|
elif isinstance(metric, MetricKind):
|
|
231
207
|
if jit:
|
|
208
|
+
|
|
209
|
+
try:
|
|
210
|
+
from usearch.numba import jit
|
|
211
|
+
except ImportError:
|
|
212
|
+
raise ModuleNotFoundError(
|
|
213
|
+
'To use JIT install Numba with `pip install numba`.'
|
|
214
|
+
'Alternatively, reinstall usearch with `pip install usearch[jit]`')
|
|
215
|
+
|
|
232
216
|
self._metric_kind = metric
|
|
233
|
-
self._metric_jit =
|
|
217
|
+
self._metric_jit = jit(
|
|
234
218
|
ndim=ndim,
|
|
235
|
-
|
|
219
|
+
metric=metric,
|
|
236
220
|
dtype=dtype,
|
|
237
221
|
)
|
|
238
222
|
self._metric_pointer = self._metric_jit.address if \
|
|
@@ -245,27 +229,22 @@ class Index:
|
|
|
245
229
|
raise ValueError(
|
|
246
230
|
'The `metric` must be Numba callback or a `MetricKind`')
|
|
247
231
|
|
|
248
|
-
if metric in
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
metric=self._metric_kind,
|
|
252
|
-
connectivity=connectivity,
|
|
253
|
-
expansion_add=expansion_add,
|
|
254
|
-
expansion_search=expansion_search,
|
|
255
|
-
)
|
|
232
|
+
if metric in MetricKindBitwise:
|
|
233
|
+
assert dtype is None or dtype == ScalarKind.B1
|
|
234
|
+
dtype = ScalarKind.B1
|
|
256
235
|
else:
|
|
257
|
-
|
|
258
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
236
|
+
dtype = _normalize_dtype(dtype)
|
|
237
|
+
|
|
238
|
+
self._compiled = _CompiledIndex(
|
|
239
|
+
ndim=ndim,
|
|
240
|
+
metric=self._metric_kind,
|
|
241
|
+
metric_pointer=self._metric_pointer,
|
|
242
|
+
dtype=dtype,
|
|
243
|
+
connectivity=connectivity,
|
|
244
|
+
expansion_add=expansion_add,
|
|
245
|
+
expansion_search=expansion_search,
|
|
246
|
+
tune=tune,
|
|
247
|
+
)
|
|
269
248
|
|
|
270
249
|
self.path = path
|
|
271
250
|
if path and os.path.exists(path):
|
|
@@ -276,7 +255,7 @@ class Index:
|
|
|
276
255
|
|
|
277
256
|
def add(
|
|
278
257
|
self, labels, vectors, *,
|
|
279
|
-
copy: bool = True, threads: int = 0):
|
|
258
|
+
copy: bool = True, threads: int = 0) -> Union[int, np.ndarray]:
|
|
280
259
|
"""Inserts one or move vectors into the index.
|
|
281
260
|
|
|
282
261
|
For maximal performance the `labels` and `vectors`
|
|
@@ -291,7 +270,7 @@ class Index:
|
|
|
291
270
|
pass `copy=False`, if you can guarantee the lifetime of the
|
|
292
271
|
primary vectors store during the process of construction.
|
|
293
272
|
|
|
294
|
-
:param labels: Unique identifier for passed vectors
|
|
273
|
+
:param labels: Unique identifier for passed vectors, optional
|
|
295
274
|
:type labels: Buffer
|
|
296
275
|
:param vectors: Collection of vectors.
|
|
297
276
|
:type vectors: Buffer
|
|
@@ -299,6 +278,9 @@ class Index:
|
|
|
299
278
|
:type copy: bool, optional
|
|
300
279
|
:param threads: Optimal number of cores to use, defaults to 0
|
|
301
280
|
:type threads: int, optional
|
|
281
|
+
|
|
282
|
+
:return: Inserted label or labels
|
|
283
|
+
:type: Union[int, np.ndarray]
|
|
302
284
|
"""
|
|
303
285
|
assert isinstance(vectors, np.ndarray), 'Expects a NumPy array'
|
|
304
286
|
assert vectors.ndim == 1 or vectors.ndim == 2, 'Expects a matrix or vector'
|
|
@@ -314,11 +296,12 @@ class Index:
|
|
|
314
296
|
labels = start_id
|
|
315
297
|
if isinstance(labels, np.ndarray):
|
|
316
298
|
labels = labels.astype(np.longlong)
|
|
299
|
+
elif isinstance(labels, Iterable):
|
|
300
|
+
labels = np.array(labels, dtype=np.longlong)
|
|
317
301
|
|
|
318
302
|
self._compiled.add(labels, vectors, copy=copy, threads=threads)
|
|
319
303
|
|
|
320
|
-
|
|
321
|
-
return labels
|
|
304
|
+
return labels
|
|
322
305
|
|
|
323
306
|
def search(
|
|
324
307
|
self, vectors, k: int = 10, *,
|
|
@@ -343,8 +326,24 @@ class Index:
|
|
|
343
326
|
)
|
|
344
327
|
return Matches(*tuple_)
|
|
345
328
|
|
|
329
|
+
@property
|
|
330
|
+
def specs(self) -> dict:
|
|
331
|
+
return {
|
|
332
|
+
'Class': 'usearch.Index',
|
|
333
|
+
'Connectivity': self.connectivity,
|
|
334
|
+
'Size': self.size,
|
|
335
|
+
'Dimensions': self.ndim,
|
|
336
|
+
'Expansion@Add': self.expansion_add,
|
|
337
|
+
'Expansion@Search': self.expansion_search,
|
|
338
|
+
'OpenMP': USES_OPENMP,
|
|
339
|
+
'SimSIMD': USES_SIMSIMD,
|
|
340
|
+
'JIT': self.jit,
|
|
341
|
+
'DType': self.dtype,
|
|
342
|
+
'Path': self.path,
|
|
343
|
+
}
|
|
344
|
+
|
|
346
345
|
def __len__(self) -> int:
|
|
347
|
-
return
|
|
346
|
+
return self._compiled.__len__()
|
|
348
347
|
|
|
349
348
|
@property
|
|
350
349
|
def jit(self) -> bool:
|
|
@@ -363,7 +362,7 @@ class Index:
|
|
|
363
362
|
return self._metric_kind
|
|
364
363
|
|
|
365
364
|
@property
|
|
366
|
-
def dtype(self) ->
|
|
365
|
+
def dtype(self) -> ScalarKind:
|
|
367
366
|
return self._compiled.dtype
|
|
368
367
|
|
|
369
368
|
@property
|
|
@@ -387,11 +386,11 @@ class Index:
|
|
|
387
386
|
return self._compiled.expansion_search
|
|
388
387
|
|
|
389
388
|
@expansion_add.setter
|
|
390
|
-
def
|
|
389
|
+
def expansion_add(self, v: int):
|
|
391
390
|
self._compiled.expansion_add = v
|
|
392
391
|
|
|
393
392
|
@expansion_search.setter
|
|
394
|
-
def
|
|
393
|
+
def expansion_search(self, v: int):
|
|
395
394
|
self._compiled.expansion_search = v
|
|
396
395
|
|
|
397
396
|
def save(self, path: os.PathLike):
|
|
@@ -406,5 +405,27 @@ class Index:
|
|
|
406
405
|
def clear(self):
|
|
407
406
|
self._compiled.clear()
|
|
408
407
|
|
|
409
|
-
|
|
410
|
-
|
|
408
|
+
@property
|
|
409
|
+
def labels(self) -> np.ndarray:
|
|
410
|
+
return self._compiled.labels
|
|
411
|
+
|
|
412
|
+
def get_vectors(
|
|
413
|
+
self,
|
|
414
|
+
labels: np.ndarray,
|
|
415
|
+
dtype: ScalarKind = ScalarKind.F32) -> np.ndarray:
|
|
416
|
+
|
|
417
|
+
dtype = _normalize_dtype(dtype)
|
|
418
|
+
return np.vstack([self._compiled.__getitem__(l, dtype) for l in labels])
|
|
419
|
+
|
|
420
|
+
@property
|
|
421
|
+
def vectors(self) -> np.ndarray:
|
|
422
|
+
return self.get_vectors(self.labels)
|
|
423
|
+
|
|
424
|
+
def __delitem__(self, label: int):
|
|
425
|
+
raise NotImplementedError()
|
|
426
|
+
|
|
427
|
+
def __contains__(self, label: int) -> bool:
|
|
428
|
+
return self._compiled.__contains__(label)
|
|
429
|
+
|
|
430
|
+
def __getitem__(self, label: int) -> np.ndarray:
|
|
431
|
+
return self._compiled.__getitem__(label)
|
package/python/usearch/io.py
CHANGED
|
@@ -7,7 +7,7 @@ def load_matrix(
|
|
|
7
7
|
filename: str,
|
|
8
8
|
start_row: int = 0, count_rows: int = None,
|
|
9
9
|
view: bool = False) -> np.ndarray:
|
|
10
|
-
"""Read *.ibin, *.hbin, *.fbin, *.dbin files with
|
|
10
|
+
"""Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
|
|
11
11
|
|
|
12
12
|
:param filename: path to the matrix file
|
|
13
13
|
:param start_row: start reading vectors from this index
|
|
@@ -31,6 +31,9 @@ def load_matrix(
|
|
|
31
31
|
elif filename.endswith('.ibin'):
|
|
32
32
|
dtype = np.int32
|
|
33
33
|
scalar_size = 4
|
|
34
|
+
elif filename.endswith('.bbin'):
|
|
35
|
+
dtype = np.int8
|
|
36
|
+
scalar_size = 1
|
|
34
37
|
else:
|
|
35
38
|
raise Exception('Unknown file type')
|
|
36
39
|
|
|
@@ -49,7 +52,7 @@ def load_matrix(
|
|
|
49
52
|
|
|
50
53
|
|
|
51
54
|
def save_matrix(vectors: np.ndarray, filename: str):
|
|
52
|
-
"""Write *.ibin, *.hbin, *.fbin, *.dbin files with
|
|
55
|
+
"""Write *.ibin, *.bbib, *.hbin, *.fbin, *.dbin files with matrices.
|
|
53
56
|
|
|
54
57
|
:param vectors: the matrix to serialize
|
|
55
58
|
:type vectors: numpy.ndarray
|
|
@@ -65,6 +68,8 @@ def save_matrix(vectors: np.ndarray, filename: str):
|
|
|
65
68
|
dtype = np.float16
|
|
66
69
|
elif filename.endswith('.ibin'):
|
|
67
70
|
dtype = np.int32
|
|
71
|
+
elif filename.endswith('.bbin'):
|
|
72
|
+
dtype = np.int8
|
|
68
73
|
else:
|
|
69
74
|
raise Exception('Unknown file type')
|
|
70
75
|
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
from collections import deque
|
|
2
|
+
|
|
3
|
+
import numpy as np
|
|
4
|
+
from numba import njit
|
|
5
|
+
|
|
6
|
+
from index import Index, Matches, Label
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class bidict(dict):
|
|
10
|
+
|
|
11
|
+
def __init__(self, *args, **kwargs):
|
|
12
|
+
super(bidict, self).__init__(*args, **kwargs)
|
|
13
|
+
self.inverse = {}
|
|
14
|
+
for key, value in self.items():
|
|
15
|
+
self.inverse.setdefault(value, []).append(key)
|
|
16
|
+
|
|
17
|
+
def __setitem__(self, key, value):
|
|
18
|
+
if key in self:
|
|
19
|
+
self.inverse[self[key]].remove(key)
|
|
20
|
+
super(bidict, self).__setitem__(key, value)
|
|
21
|
+
self.inverse.setdefault(value, []).append(key)
|
|
22
|
+
|
|
23
|
+
def __delitem__(self, key):
|
|
24
|
+
self.inverse.setdefault(self[key], []).remove(key)
|
|
25
|
+
if self[key] in self.inverse and not self.inverse[self[key]]:
|
|
26
|
+
del self.inverse[self[key]]
|
|
27
|
+
super(bidict, self).__delitem__(key)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@njit
|
|
31
|
+
def index_in_array(array: np.ndarray, item: Label) -> int:
|
|
32
|
+
for idx, val in np.ndenumerate(array):
|
|
33
|
+
if val == item:
|
|
34
|
+
return idx
|
|
35
|
+
|
|
36
|
+
return len(array)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
|
|
40
|
+
"""Performs a Semantic Join, matching one entry from `a` with
|
|
41
|
+
one entry from `b`, converging towards a Stable Marriage.
|
|
42
|
+
Assuming the collections can be different in size, classical solution
|
|
43
|
+
doesn't provide
|
|
44
|
+
|
|
45
|
+
:param resolution: Approximate matches per member to consider, defaults to 10
|
|
46
|
+
:type resolution: int, optional
|
|
47
|
+
:return: Bidirectional mapping from men to women and `.inverse`
|
|
48
|
+
:rtype: bidict
|
|
49
|
+
"""
|
|
50
|
+
|
|
51
|
+
men, women = (a, b) if len(a) > len(b) else (b, a)
|
|
52
|
+
matches = bidict()
|
|
53
|
+
|
|
54
|
+
men_vectors = np.vstack(men[i] for i in range(len(men)))
|
|
55
|
+
women_vectors = np.vstack(women[i] for i in range(len(women)))
|
|
56
|
+
|
|
57
|
+
man_count_proposed = np.zeros(len(men))
|
|
58
|
+
man_to_women_preferences: Matches = men.search(women_vectors, resolution)
|
|
59
|
+
woman_to_men_preferences: Matches = women.search(men_vectors, resolution)
|
|
60
|
+
|
|
61
|
+
# A nice optimization is to prioritize free men by the quality of their
|
|
62
|
+
# best remaining variant. TODO: Replace with `heapq`
|
|
63
|
+
free_men = deque(range(len(men)))
|
|
64
|
+
|
|
65
|
+
idle_cycles: int = 0
|
|
66
|
+
hopeless_men_count = len(men) - len(women)
|
|
67
|
+
while len(free_men) > hopeless_men_count:
|
|
68
|
+
|
|
69
|
+
# In the worst case scenario we may need to match more candidates for every
|
|
70
|
+
# remaining man. This, however, may drastically increase the runtime.
|
|
71
|
+
if len(free_men) == idle_cycles:
|
|
72
|
+
break
|
|
73
|
+
|
|
74
|
+
man: Label = free_men.popleft()
|
|
75
|
+
count_proposals = man_count_proposed[man]
|
|
76
|
+
if count_proposals == man_to_women_preferences.counts[man]:
|
|
77
|
+
free_men.append(man)
|
|
78
|
+
idle_cycles += 1
|
|
79
|
+
continue
|
|
80
|
+
|
|
81
|
+
woman: Label = man_to_women_preferences.labels[man, count_proposals]
|
|
82
|
+
man_count_proposed[man] += 1
|
|
83
|
+
|
|
84
|
+
if woman not in matches.inverse:
|
|
85
|
+
matches[man] = woman
|
|
86
|
+
|
|
87
|
+
else:
|
|
88
|
+
her_preferences = woman_to_men_preferences.labels[woman, :]
|
|
89
|
+
husband: Label = matches.inverse[woman]
|
|
90
|
+
husband_rank = index_in_array(her_preferences, husband)
|
|
91
|
+
challenger_rank = index_in_array(her_preferences, man)
|
|
92
|
+
if challenger_rank < husband_rank:
|
|
93
|
+
del matches[husband]
|
|
94
|
+
free_men.append(husband)
|
|
95
|
+
matches[man] = woman
|
|
96
|
+
|
|
97
|
+
else:
|
|
98
|
+
free_men.append(man)
|
|
99
|
+
|
|
100
|
+
return matches
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
# The purpose of this file is to provide Pythonic wrapper on top
|
|
2
|
+
# the native precompiled CPython module. It improves compatibility
|
|
3
|
+
# Python tooling, linters, and static analyzers. It also embeds JIT
|
|
4
|
+
# into the primary `Index` class, connecting USearch with Numba.
|
|
5
|
+
from math import sqrt
|
|
6
|
+
from typing import Callable
|
|
7
|
+
|
|
8
|
+
import numpy as np
|
|
9
|
+
from numba import cfunc, types, carray
|
|
10
|
+
|
|
11
|
+
from usearch.compiled import MetricKind
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
signature_f16args = types.float32(types.CPointer(
|
|
15
|
+
types.float16), types.CPointer(types.float16))
|
|
16
|
+
|
|
17
|
+
signature_f32args = types.float32(types.CPointer(
|
|
18
|
+
types.float32), types.CPointer(types.float32))
|
|
19
|
+
|
|
20
|
+
signature_f64args = types.float32(types.CPointer(
|
|
21
|
+
types.float64), types.CPointer(types.float64))
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def jit(ndim: int,
|
|
25
|
+
metric: MetricKind = MetricKind.Cos,
|
|
26
|
+
dtype: np.generic = np.float32) -> Callable:
|
|
27
|
+
"""JIT-compiles a distance metric specifically tuned for the target hardware
|
|
28
|
+
and number of dimensions.
|
|
29
|
+
|
|
30
|
+
Uses Numba `cfunc` functionality, annotating it with Numba `types` instead
|
|
31
|
+
of `ctypes` to support half-precision.
|
|
32
|
+
https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
|
|
33
|
+
"""
|
|
34
|
+
normalize_dtype = {
|
|
35
|
+
# Half-precision is still unsupported
|
|
36
|
+
# https://github.com/numba/numba/issues/4402
|
|
37
|
+
# 'f16': np.float16,
|
|
38
|
+
'f32': np.float32,
|
|
39
|
+
'f64': np.float64,
|
|
40
|
+
}
|
|
41
|
+
if dtype in normalize_dtype.keys():
|
|
42
|
+
dtype = normalize_dtype[dtype]
|
|
43
|
+
if dtype not in normalize_dtype.values():
|
|
44
|
+
return None
|
|
45
|
+
|
|
46
|
+
scalar_kind_to_accumulator_type = {
|
|
47
|
+
np.float16: types.float32,
|
|
48
|
+
np.float32: types.float32,
|
|
49
|
+
np.float64: types.float64,
|
|
50
|
+
}
|
|
51
|
+
accumulator = scalar_kind_to_accumulator_type[dtype]
|
|
52
|
+
|
|
53
|
+
def numba_ip(a, b):
|
|
54
|
+
a_array = carray(a, ndim)
|
|
55
|
+
b_array = carray(b, ndim)
|
|
56
|
+
ab = accumulator(0)
|
|
57
|
+
for i in range(ndim):
|
|
58
|
+
ab += a_array[i] * b_array[i]
|
|
59
|
+
return types.float32(1 - ab)
|
|
60
|
+
|
|
61
|
+
def numba_cos(a, b):
|
|
62
|
+
a_array = carray(a, ndim)
|
|
63
|
+
b_array = carray(b, ndim)
|
|
64
|
+
ab = accumulator(0)
|
|
65
|
+
a_sq = accumulator(0)
|
|
66
|
+
b_sq = accumulator(0)
|
|
67
|
+
for i in range(ndim):
|
|
68
|
+
ab += a_array[i] * b_array[i]
|
|
69
|
+
a_sq += a_array[i] * a_array[i]
|
|
70
|
+
b_sq += b_array[i] * b_array[i]
|
|
71
|
+
a_norm = sqrt(a_sq)
|
|
72
|
+
b_norm = sqrt(b_sq)
|
|
73
|
+
if a_norm == 0 and b_norm == 0:
|
|
74
|
+
return types.float32(0)
|
|
75
|
+
elif a_norm == 0 or b_norm == 0 or ab == 0:
|
|
76
|
+
return types.float32(1)
|
|
77
|
+
else:
|
|
78
|
+
return types.float32(1 - ab / (a_norm * b_norm))
|
|
79
|
+
|
|
80
|
+
def numba_l2sq(a, b):
|
|
81
|
+
a_array = carray(a, ndim)
|
|
82
|
+
b_array = carray(b, ndim)
|
|
83
|
+
ab_delta_sq = accumulator(0)
|
|
84
|
+
for i in range(ndim):
|
|
85
|
+
ab_delta_sq += (a_array[i] - b_array[i]) * \
|
|
86
|
+
(a_array[i] - b_array[i])
|
|
87
|
+
return types.float32(ab_delta_sq)
|
|
88
|
+
|
|
89
|
+
scalar_kind_to_signature = {
|
|
90
|
+
np.float16: signature_f16args,
|
|
91
|
+
np.float32: signature_f32args,
|
|
92
|
+
np.float64: signature_f64args,
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
metric_kind_to_function = {
|
|
96
|
+
MetricKind.IP: numba_ip,
|
|
97
|
+
MetricKind.Cos: numba_cos,
|
|
98
|
+
MetricKind.L2sq: numba_l2sq,
|
|
99
|
+
}
|
|
100
|
+
|
|
101
|
+
return cfunc(scalar_kind_to_signature[dtype])(metric_kind_to_function[metric])
|
package/rust/lib.cpp
CHANGED
|
@@ -62,7 +62,7 @@ void Index::save(rust::Str path) const { index_->save(std::string(path).c_str())
|
|
|
62
62
|
void Index::load(rust::Str path) const { index_->load(std::string(path).c_str()); }
|
|
63
63
|
void Index::view(rust::Str path) const { index_->view(std::string(path).c_str()); }
|
|
64
64
|
|
|
65
|
-
|
|
65
|
+
scalar_kind_t accuracy(rust::Str quant) { return scalar_kind_from_name(quant.data(), quant.size()); }
|
|
66
66
|
|
|
67
67
|
std::unique_ptr<Index> wrap(punned_t&& index) {
|
|
68
68
|
std::unique_ptr<punned_t> punned_ptr;
|
|
@@ -72,26 +72,39 @@ std::unique_ptr<Index> wrap(punned_t&& index) {
|
|
|
72
72
|
return result;
|
|
73
73
|
}
|
|
74
74
|
|
|
75
|
-
index_config_t config(size_t connectivity
|
|
75
|
+
index_config_t config(size_t connectivity) {
|
|
76
76
|
index_config_t result;
|
|
77
77
|
result.connectivity = connectivity ? connectivity : default_connectivity();
|
|
78
|
-
result.expansion_add = exp_add ? exp_add : default_expansion_add();
|
|
79
|
-
result.expansion_search = exp_search ? exp_search : default_expansion_search();
|
|
80
78
|
return result;
|
|
81
79
|
}
|
|
82
80
|
|
|
83
|
-
|
|
84
|
-
|
|
81
|
+
metric_kind_t rust_to_cpp_metric(MetricKind value) {
|
|
82
|
+
switch (value) {
|
|
83
|
+
case MetricKind::IP: return metric_kind_t::ip_k;
|
|
84
|
+
case MetricKind::L2Sq: return metric_kind_t::l2sq_k;
|
|
85
|
+
case MetricKind::Cos: return metric_kind_t::cos_k;
|
|
86
|
+
case MetricKind::Pearson: return metric_kind_t::pearson_k;
|
|
87
|
+
case MetricKind::Haversine: return metric_kind_t::haversine_k;
|
|
88
|
+
case MetricKind::Hamming: return metric_kind_t::hamming_k;
|
|
89
|
+
case MetricKind::Tanimoto: return metric_kind_t::tanimoto_k;
|
|
90
|
+
case MetricKind::Sorensen: return metric_kind_t::sorensen_k;
|
|
91
|
+
default: return metric_kind_t::unknown_k;
|
|
92
|
+
}
|
|
85
93
|
}
|
|
86
94
|
|
|
87
|
-
|
|
88
|
-
|
|
95
|
+
scalar_kind_t rust_to_cpp_scalar(ScalarKind value) {
|
|
96
|
+
switch (value) {
|
|
97
|
+
case ScalarKind::F8: return scalar_kind_t::f8_k;
|
|
98
|
+
case ScalarKind::F16: return scalar_kind_t::f16_k;
|
|
99
|
+
case ScalarKind::F32: return scalar_kind_t::f32_k;
|
|
100
|
+
case ScalarKind::F64: return scalar_kind_t::f64_k;
|
|
101
|
+
case ScalarKind::B1: return scalar_kind_t::b1x8_k;
|
|
102
|
+
default: return scalar_kind_t::unknown_k;
|
|
103
|
+
}
|
|
89
104
|
}
|
|
90
105
|
|
|
91
|
-
std::unique_ptr<Index>
|
|
92
|
-
return wrap(punned_t::
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
std::unique_ptr<Index> new_haversine(rust::Str quant, size_t connectivity, size_t exp_add, size_t exp_search) {
|
|
96
|
-
return wrap(punned_t::haversine(accuracy(quant), config(connectivity, exp_add, exp_search)));
|
|
106
|
+
std::unique_ptr<Index> new_index(IndexOptions const& options) {
|
|
107
|
+
return wrap(punned_t::make(options.dimensions, rust_to_cpp_metric(options.metric), config(options.connectivity),
|
|
108
|
+
rust_to_cpp_scalar(options.quantization), options.expansion_add,
|
|
109
|
+
options.expansion_search));
|
|
97
110
|
}
|