usearch 0.22.3 → 1.1.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/prerelease.yml +3 -3
- package/.releaserc +10 -2
- package/.vscode/launch.json +10 -10
- package/.vscode/settings.json +9 -1
- package/.vscode/tasks.json +4 -4
- package/CMakeLists.txt +6 -1
- package/Cargo.toml +2 -2
- package/README.md +35 -41
- package/build.rs +2 -2
- package/c/CMakeLists.txt +48 -0
- package/c/README.md +34 -0
- package/c/lib.cpp +83 -72
- package/c/test.c +20 -20
- package/c/usearch.h +132 -32
- package/conanfile.py +1 -1
- package/cpp/CMakeLists.txt +3 -1
- package/cpp/README.md +16 -16
- package/cpp/bench.cpp +50 -126
- package/cpp/test.cpp +164 -80
- package/docs/benchmarks.md +26 -5
- package/docs/c/reference.rst +3 -3
- package/docs/compilation.md +31 -49
- package/docs/cpp/index.rst +1 -1
- package/docs/cpp/reference.rst +4 -4
- package/golang/lib.go +275 -141
- package/golang/lib_test.go +105 -0
- package/include/usearch/index.hpp +1681 -1656
- package/include/usearch/index_dense.hpp +1400 -0
- package/include/usearch/index_plugins.hpp +1263 -0
- package/java/README.md +1 -1
- package/java/cloud/unum/usearch/Index.java +10 -10
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +34 -37
- package/java/test/IndexTest.java +1 -1
- package/javascript/README.md +6 -6
- package/javascript/lib.cpp +156 -55
- package/javascript/test.js +26 -11
- package/javascript/usearch.d.ts +102 -0
- package/javascript/usearch.js +2 -118
- package/objc/USearchObjective.mm +57 -31
- package/objc/include/USearchObjective.h +32 -22
- package/package.json +5 -3
- package/python/README.md +23 -14
- package/python/lib.cpp +468 -444
- package/python/scripts/bench.ipynb +3 -2
- package/python/scripts/bench.py +4 -4
- package/python/scripts/index_faiss.py +5 -5
- package/python/scripts/join.py +4 -4
- package/python/scripts/test.py +152 -61
- package/python/scripts/test_jit.py +10 -12
- package/python/scripts/test_sparse.py +32 -0
- package/python/usearch/client.py +26 -28
- package/python/usearch/eval.py +73 -23
- package/python/usearch/index.py +520 -228
- package/python/usearch/io.py +57 -25
- package/python/usearch/numba.py +4 -4
- package/python/usearch/server.py +10 -10
- package/rust/lib.cpp +34 -30
- package/rust/lib.hpp +14 -12
- package/rust/lib.rs +8 -7
- package/setup.py +5 -2
- package/simsimd/.github/workflows/prerelease.yml +73 -0
- package/simsimd/.github/workflows/release.yml +110 -0
- package/simsimd/.github/workflows/update_version.sh +3 -0
- package/simsimd/.releaserc +77 -0
- package/simsimd/.vscode/launch.json +27 -1
- package/simsimd/.vscode/settings.json +22 -1
- package/simsimd/.vscode/tasks.json +6 -0
- package/simsimd/CMakeLists.txt +3 -2
- package/simsimd/README.md +82 -46
- package/simsimd/bench.cxx +63 -24
- package/simsimd/include/simsimd/simsimd.h +88 -66
- package/simsimd/include/simsimd/simsimd_chem.h +126 -0
- package/simsimd/package.json +8 -0
- package/simsimd/pyproject.toml +42 -0
- package/simsimd/python/python.c +81 -0
- package/simsimd/python/run.sh +1 -0
- package/simsimd/python/test.c +109 -0
- package/simsimd/python/test.py +23 -0
- package/simsimd/setup.py +44 -0
- package/swift/Index+Sugar.swift +75 -23
- package/swift/README.md +26 -2
- package/swift/Test.swift +19 -7
- package/swift/USearch.swift +2 -2
- package/tsconfig.json +8 -0
- package/wasm/CMakeLists.txt +1 -0
- package/wasmer.toml +1 -1
- package/wolfram/CMakeLists.txt +1 -0
- package/wolfram/lib.cpp +11 -13
- package/golang/usearch.h +0 -9
- package/golang/usearch_par_test.go +0 -40
- package/golang/usearch_test.go +0 -82
- package/include/usearch/index_punned_dense.hpp +0 -1001
- package/include/usearch/index_punned_helpers.hpp +0 -676
package/python/usearch/index.py
CHANGED
|
@@ -7,13 +7,17 @@ from __future__ import annotations
|
|
|
7
7
|
import os
|
|
8
8
|
import math
|
|
9
9
|
from typing import Optional, Union, NamedTuple, List, Iterable
|
|
10
|
+
from dataclasses import dataclass
|
|
10
11
|
|
|
11
12
|
import numpy as np
|
|
12
13
|
from tqdm import tqdm
|
|
13
14
|
|
|
14
|
-
from usearch.compiled import Index as _CompiledIndex
|
|
15
|
-
from usearch.compiled import
|
|
15
|
+
from usearch.compiled import Index as _CompiledIndex
|
|
16
|
+
from usearch.compiled import Indexes as _CompiledIndexes
|
|
17
|
+
from usearch.compiled import IndexStats as _CompiledIndexStats
|
|
16
18
|
|
|
19
|
+
from usearch.compiled import index_dense_metadata as _index_dense_metadata
|
|
20
|
+
from usearch.compiled import exact_search as _exact_search
|
|
17
21
|
from usearch.compiled import MetricKind, ScalarKind, MetricSignature
|
|
18
22
|
from usearch.compiled import (
|
|
19
23
|
DEFAULT_CONNECTIVITY,
|
|
@@ -30,12 +34,10 @@ MetricKindBitwise = (
|
|
|
30
34
|
MetricKind.Sorensen,
|
|
31
35
|
)
|
|
32
36
|
|
|
33
|
-
|
|
37
|
+
Key = np.uint64
|
|
34
38
|
|
|
35
|
-
Label = np.uint64
|
|
36
39
|
|
|
37
|
-
|
|
38
|
-
def _normalize_dtype(dtype, metric: MetricKind = MetricKind.IP) -> ScalarKind:
|
|
40
|
+
def _normalize_dtype(dtype, metric: MetricKind = MetricKind.Cos) -> ScalarKind:
|
|
39
41
|
if dtype is None or dtype == "":
|
|
40
42
|
return ScalarKind.B1 if metric in MetricKindBitwise else ScalarKind.F32
|
|
41
43
|
|
|
@@ -49,14 +51,16 @@ def _normalize_dtype(dtype, metric: MetricKind = MetricKind.IP) -> ScalarKind:
|
|
|
49
51
|
"f64": ScalarKind.F64,
|
|
50
52
|
"f32": ScalarKind.F32,
|
|
51
53
|
"f16": ScalarKind.F16,
|
|
52
|
-
"
|
|
54
|
+
"i8": ScalarKind.I8,
|
|
53
55
|
"b1": ScalarKind.B1,
|
|
54
56
|
"b1x8": ScalarKind.B1,
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
57
|
+
"float64": ScalarKind.F64,
|
|
58
|
+
"float32": ScalarKind.F32,
|
|
59
|
+
"float16": ScalarKind.F16,
|
|
60
|
+
"int8": ScalarKind.I8,
|
|
59
61
|
}
|
|
62
|
+
if isinstance(dtype, np.dtype):
|
|
63
|
+
dtype = dtype.name
|
|
60
64
|
return _normalize[dtype]
|
|
61
65
|
|
|
62
66
|
|
|
@@ -65,7 +69,7 @@ def _to_numpy_compatible_dtype(dtype: ScalarKind) -> ScalarKind:
|
|
|
65
69
|
ScalarKind.F64: ScalarKind.F64,
|
|
66
70
|
ScalarKind.F32: ScalarKind.F32,
|
|
67
71
|
ScalarKind.F16: ScalarKind.F16,
|
|
68
|
-
ScalarKind.
|
|
72
|
+
ScalarKind.I8: ScalarKind.F16,
|
|
69
73
|
ScalarKind.B1: ScalarKind.B1,
|
|
70
74
|
}
|
|
71
75
|
return _normalize[dtype]
|
|
@@ -76,7 +80,7 @@ def _to_numpy_dtype(dtype: ScalarKind):
|
|
|
76
80
|
ScalarKind.F64: np.float64,
|
|
77
81
|
ScalarKind.F32: np.float32,
|
|
78
82
|
ScalarKind.F16: np.float16,
|
|
79
|
-
ScalarKind.
|
|
83
|
+
ScalarKind.I8: np.float16,
|
|
80
84
|
ScalarKind.B1: np.uint8,
|
|
81
85
|
}
|
|
82
86
|
return _normalize[dtype]
|
|
@@ -84,15 +88,19 @@ def _to_numpy_dtype(dtype: ScalarKind):
|
|
|
84
88
|
|
|
85
89
|
def _normalize_metric(metric):
|
|
86
90
|
if metric is None:
|
|
87
|
-
return MetricKind.
|
|
91
|
+
return MetricKind.Cos
|
|
88
92
|
|
|
89
93
|
if isinstance(metric, str):
|
|
90
94
|
_normalize = {
|
|
91
95
|
"cos": MetricKind.Cos,
|
|
96
|
+
"cosine": MetricKind.Cos,
|
|
92
97
|
"ip": MetricKind.IP,
|
|
98
|
+
"dot": MetricKind.IP,
|
|
99
|
+
"inner_product": MetricKind.IP,
|
|
100
|
+
"l2sq": MetricKind.L2sq,
|
|
93
101
|
"l2_sq": MetricKind.L2sq,
|
|
94
102
|
"haversine": MetricKind.Haversine,
|
|
95
|
-
"
|
|
103
|
+
"pearson": MetricKind.Pearson,
|
|
96
104
|
"hamming": MetricKind.Hamming,
|
|
97
105
|
"tanimoto": MetricKind.Tanimoto,
|
|
98
106
|
"sorensen": MetricKind.Sorensen,
|
|
@@ -102,59 +110,234 @@ def _normalize_metric(metric):
|
|
|
102
110
|
return metric
|
|
103
111
|
|
|
104
112
|
|
|
105
|
-
|
|
106
|
-
|
|
113
|
+
def _search_in_compiled(
|
|
114
|
+
*,
|
|
115
|
+
compiled: Union[_CompiledIndex, _CompiledIndexes],
|
|
116
|
+
vectors: np.ndarray,
|
|
117
|
+
k: int,
|
|
118
|
+
threads: int,
|
|
119
|
+
exact: bool,
|
|
120
|
+
log: Union[str, bool],
|
|
121
|
+
batch_size: int,
|
|
122
|
+
) -> Union[Matches, BatchMatches]:
|
|
123
|
+
#
|
|
124
|
+
assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
|
|
125
|
+
assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
|
|
126
|
+
if vectors.ndim == 1:
|
|
127
|
+
vectors = vectors.reshape(1, len(vectors))
|
|
128
|
+
count_vectors = vectors.shape[0]
|
|
129
|
+
|
|
130
|
+
def distil_batch(
|
|
131
|
+
batch_matches: BatchMatches,
|
|
132
|
+
) -> Union[BatchMatches, Matches]:
|
|
133
|
+
return batch_matches[0] if count_vectors == 1 else batch_matches
|
|
134
|
+
|
|
135
|
+
if log and batch_size == 0:
|
|
136
|
+
batch_size = int(math.ceil(count_vectors / 100))
|
|
137
|
+
|
|
138
|
+
if batch_size:
|
|
139
|
+
tasks = [
|
|
140
|
+
vectors[start_row : start_row + batch_size, :]
|
|
141
|
+
for start_row in range(0, count_vectors, batch_size)
|
|
142
|
+
]
|
|
143
|
+
tasks_matches = []
|
|
144
|
+
name = log if isinstance(log, str) else "Search"
|
|
145
|
+
pbar = tqdm(
|
|
146
|
+
tasks,
|
|
147
|
+
desc=name,
|
|
148
|
+
total=count_vectors,
|
|
149
|
+
unit="vector",
|
|
150
|
+
disable=log is False,
|
|
151
|
+
)
|
|
152
|
+
for vectors in tasks:
|
|
153
|
+
tuple_ = compiled.search(
|
|
154
|
+
vectors,
|
|
155
|
+
k,
|
|
156
|
+
exact=exact,
|
|
157
|
+
threads=threads,
|
|
158
|
+
)
|
|
159
|
+
tasks_matches.append(BatchMatches(*tuple_))
|
|
160
|
+
pbar.update(vectors.shape[0])
|
|
161
|
+
|
|
162
|
+
pbar.close()
|
|
163
|
+
return distil_batch(
|
|
164
|
+
BatchMatches(
|
|
165
|
+
keys=np.vstack([m.keys for m in tasks_matches]),
|
|
166
|
+
distances=np.vstack([m.distances for m in tasks_matches]),
|
|
167
|
+
counts=np.concatenate([m.counts for m in tasks_matches], axis=None),
|
|
168
|
+
)
|
|
169
|
+
)
|
|
170
|
+
|
|
171
|
+
else:
|
|
172
|
+
tuple_ = compiled.search(
|
|
173
|
+
vectors,
|
|
174
|
+
k,
|
|
175
|
+
exact=exact,
|
|
176
|
+
threads=threads,
|
|
177
|
+
)
|
|
178
|
+
return distil_batch(BatchMatches(*tuple_))
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
def _add_to_compiled(
|
|
182
|
+
*,
|
|
183
|
+
compiled,
|
|
184
|
+
keys,
|
|
185
|
+
vectors,
|
|
186
|
+
copy: bool,
|
|
187
|
+
threads: int,
|
|
188
|
+
log: Union[str, bool],
|
|
189
|
+
batch_size: int,
|
|
190
|
+
) -> Union[int, np.ndarray]:
|
|
191
|
+
assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
|
|
192
|
+
assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
|
|
193
|
+
if vectors.ndim == 1:
|
|
194
|
+
vectors = vectors.reshape(1, len(vectors))
|
|
195
|
+
|
|
196
|
+
# Validate or generate the keys
|
|
197
|
+
count_vectors = vectors.shape[0]
|
|
198
|
+
generate_labels = keys is None
|
|
199
|
+
if generate_labels:
|
|
200
|
+
start_id = len(compiled)
|
|
201
|
+
keys = np.arange(start_id, start_id + count_vectors, dtype=Key)
|
|
202
|
+
else:
|
|
203
|
+
if not isinstance(keys, Iterable):
|
|
204
|
+
assert count_vectors == 1, "Each vector must have a key"
|
|
205
|
+
keys = [keys]
|
|
206
|
+
keys = np.array(keys).astype(Key)
|
|
207
|
+
|
|
208
|
+
assert len(keys) == count_vectors
|
|
209
|
+
|
|
210
|
+
# If logging is requested, and batch size is undefined, set it to grow 1% at a time:
|
|
211
|
+
if log and batch_size == 0:
|
|
212
|
+
batch_size = int(math.ceil(count_vectors / 100))
|
|
213
|
+
|
|
214
|
+
# Split into batches and log progress, if needed
|
|
215
|
+
if batch_size:
|
|
216
|
+
keys = [
|
|
217
|
+
keys[start_row : start_row + batch_size]
|
|
218
|
+
for start_row in range(0, count_vectors, batch_size)
|
|
219
|
+
]
|
|
220
|
+
vectors = [
|
|
221
|
+
vectors[start_row : start_row + batch_size, :]
|
|
222
|
+
for start_row in range(0, count_vectors, batch_size)
|
|
223
|
+
]
|
|
224
|
+
tasks = zip(keys, vectors)
|
|
225
|
+
name = log if isinstance(log, str) else "Add"
|
|
226
|
+
pbar = tqdm(
|
|
227
|
+
tasks,
|
|
228
|
+
desc=name,
|
|
229
|
+
total=count_vectors,
|
|
230
|
+
unit="vector",
|
|
231
|
+
disable=log is False,
|
|
232
|
+
)
|
|
233
|
+
for keys, vectors in tasks:
|
|
234
|
+
compiled.add(keys, vectors, copy=copy, threads=threads)
|
|
235
|
+
pbar.update(len(keys))
|
|
236
|
+
|
|
237
|
+
pbar.close()
|
|
238
|
+
|
|
239
|
+
else:
|
|
240
|
+
compiled.add(keys, vectors, copy=copy, threads=threads)
|
|
241
|
+
|
|
242
|
+
return keys
|
|
243
|
+
|
|
244
|
+
|
|
245
|
+
@dataclass
|
|
246
|
+
class Match:
|
|
247
|
+
"""This class contains information about retrieved vector."""
|
|
248
|
+
|
|
249
|
+
key: int
|
|
250
|
+
distance: float
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
@dataclass
|
|
254
|
+
class Matches:
|
|
255
|
+
"""This class contains information about multiple retrieved vectors for single query,
|
|
256
|
+
i.e it is a set of `Match` instances."""
|
|
257
|
+
|
|
258
|
+
keys: np.ndarray
|
|
107
259
|
distances: np.ndarray
|
|
108
|
-
counts: Union[np.ndarray, int]
|
|
109
260
|
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
return isinstance(self.counts, np.ndarray)
|
|
261
|
+
visited_members: int
|
|
262
|
+
computed_distances: int
|
|
113
263
|
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
return len(self.counts) if isinstance(self.counts, np.ndarray) else 1
|
|
264
|
+
def __len__(self) -> int:
|
|
265
|
+
return len(self.keys)
|
|
117
266
|
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
267
|
+
def __getitem__(self, index: int) -> Match:
|
|
268
|
+
if isinstance(index, int) and index < len(self):
|
|
269
|
+
return Match(
|
|
270
|
+
key=self.keys[index],
|
|
271
|
+
distance=self.distances[index],
|
|
272
|
+
)
|
|
273
|
+
else:
|
|
274
|
+
raise IndexError(f"`index` must be an integer under {len(self)}")
|
|
275
|
+
|
|
276
|
+
def to_list(self) -> List[tuple]:
|
|
277
|
+
"""
|
|
278
|
+
Convert matches to the list of tuples which contain matches' indices and distances to them.
|
|
279
|
+
"""
|
|
280
|
+
|
|
281
|
+
return [(int(l), float(d)) for l, d in zip(self.keys, self.distances)]
|
|
282
|
+
|
|
283
|
+
def __repr__(self) -> str:
|
|
284
|
+
return f"usearch.Matches({len(self)})"
|
|
121
285
|
|
|
122
|
-
def to_list(self, row: Optional[int] = None) -> Union[List[dict], List[List[dict]]]:
|
|
123
|
-
if not self.is_multiple:
|
|
124
|
-
assert row is None, "Exporting a single sequence is only for batch requests"
|
|
125
|
-
labels = self.labels
|
|
126
|
-
distances = self.distances
|
|
127
286
|
|
|
128
|
-
|
|
129
|
-
|
|
287
|
+
@dataclass
|
|
288
|
+
class BatchMatches:
|
|
289
|
+
"""This class contains information about multiple retrieved vectors for multiple queries,
|
|
290
|
+
i.e it is a set of `Matches` instances."""
|
|
130
291
|
|
|
292
|
+
keys: np.ndarray
|
|
293
|
+
distances: np.ndarray
|
|
294
|
+
counts: np.ndarray
|
|
295
|
+
|
|
296
|
+
visited_members: int
|
|
297
|
+
computed_distances: int
|
|
298
|
+
|
|
299
|
+
def __len__(self) -> int:
|
|
300
|
+
return len(self.counts)
|
|
301
|
+
|
|
302
|
+
def __getitem__(self, index: int) -> Matches:
|
|
303
|
+
if isinstance(index, int) and index < len(self):
|
|
304
|
+
return Matches(
|
|
305
|
+
keys=self.keys[index, : self.counts[index]],
|
|
306
|
+
distances=self.distances[index, : self.counts[index]],
|
|
307
|
+
visited_members=self.visited_members // len(self),
|
|
308
|
+
computed_distances=self.computed_distances // len(self),
|
|
309
|
+
)
|
|
131
310
|
else:
|
|
132
|
-
|
|
133
|
-
labels = self.labels[row, :count]
|
|
134
|
-
distances = self.distances[row, :count]
|
|
311
|
+
raise IndexError(f"`index` must be an integer under {len(self)}")
|
|
135
312
|
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
]
|
|
313
|
+
def to_list(self) -> List[List[tuple]]:
|
|
314
|
+
"""Convert the result for each query to the list of tuples with information about its matches."""
|
|
315
|
+
list_of_matches = [self.__getitem__(row) for row in range(self.__len__())]
|
|
316
|
+
return [match.to_list() for matches in list_of_matches for match in matches]
|
|
140
317
|
|
|
141
|
-
def
|
|
142
|
-
|
|
143
|
-
|
|
318
|
+
def mean_recall(self, expected: np.ndarray, k: Optional[int] = None) -> float:
|
|
319
|
+
"""Measures recall [0, 1] as of `Matches` that contain the corresponding
|
|
320
|
+
`expected` entry anywhere among results."""
|
|
321
|
+
return self.count_matches(expected, k=k) / len(expected)
|
|
144
322
|
|
|
145
|
-
def
|
|
146
|
-
|
|
323
|
+
def count_matches(self, expected: np.ndarray, k: Optional[int] = None) -> int:
|
|
324
|
+
"""Measures recall [0, len(expected)] as of `Matches` that contain the corresponding
|
|
325
|
+
`expected` entry anywhere among results.
|
|
326
|
+
"""
|
|
327
|
+
assert len(expected) == len(self)
|
|
147
328
|
recall = 0
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
329
|
+
if k is None:
|
|
330
|
+
k = self.keys.shape[1]
|
|
331
|
+
|
|
332
|
+
if k == 1:
|
|
333
|
+
recall = np.sum(self.keys[:, 0] == expected)
|
|
334
|
+
else:
|
|
335
|
+
for i in range(len(self)):
|
|
336
|
+
recall += expected[i] in self.keys[i, :k]
|
|
337
|
+
return recall
|
|
151
338
|
|
|
152
339
|
def __repr__(self) -> str:
|
|
153
|
-
return (
|
|
154
|
-
f"usearch.Matches({self.total_matches})"
|
|
155
|
-
if self.is_multiple
|
|
156
|
-
else f"usearch.Matches({self.total_matches} across {self.batch_size} queries)"
|
|
157
|
-
)
|
|
340
|
+
return f"usearch.BatchMatches({np.sum(self.counts)} across {len(self)} queries)"
|
|
158
341
|
|
|
159
342
|
|
|
160
343
|
class CompiledMetric(NamedTuple):
|
|
@@ -166,7 +349,7 @@ class CompiledMetric(NamedTuple):
|
|
|
166
349
|
class Index:
|
|
167
350
|
"""Fast JIT-compiled vector-search index for dense equi-dimensional embeddings.
|
|
168
351
|
|
|
169
|
-
Vector
|
|
352
|
+
Vector keys must be integers.
|
|
170
353
|
Vectors must have the same number of dimensions within the index.
|
|
171
354
|
Supports Inner Product, Cosine Distance, Ln measures
|
|
172
355
|
like the Euclidean metric, as well as automatic downcasting
|
|
@@ -175,13 +358,13 @@ class Index:
|
|
|
175
358
|
|
|
176
359
|
def __init__(
|
|
177
360
|
self,
|
|
361
|
+
*,
|
|
178
362
|
ndim: int = 0,
|
|
179
|
-
metric: Union[str, MetricKind, CompiledMetric] = MetricKind.
|
|
363
|
+
metric: Union[str, MetricKind, CompiledMetric] = MetricKind.Cos,
|
|
180
364
|
dtype: Optional[Union[str, ScalarKind]] = None,
|
|
181
|
-
connectivity: int =
|
|
182
|
-
expansion_add: int =
|
|
183
|
-
expansion_search: int =
|
|
184
|
-
tune: bool = False,
|
|
365
|
+
connectivity: Optional[int] = None,
|
|
366
|
+
expansion_add: Optional[int] = None,
|
|
367
|
+
expansion_search: Optional[int] = None,
|
|
185
368
|
path: Optional[os.PathLike] = None,
|
|
186
369
|
view: bool = False,
|
|
187
370
|
) -> None:
|
|
@@ -189,21 +372,20 @@ class Index:
|
|
|
189
372
|
|
|
190
373
|
:param ndim: Number of vector dimensions
|
|
191
374
|
:type ndim: int
|
|
192
|
-
Required for some metrics,
|
|
375
|
+
Required for some metrics, pre-set for others.
|
|
193
376
|
Haversine, for example, only applies to 2-dimensional latitude/longitude
|
|
194
377
|
coordinates. Angular (Cos) and Euclidean (L2sq), obviously, apply to
|
|
195
378
|
vectors with arbitrary number of dimensions.
|
|
196
379
|
|
|
197
|
-
:param metric: Distance function, defaults to MetricKind.
|
|
380
|
+
:param metric: Distance function, defaults to MetricKind.Cos
|
|
198
381
|
:type metric: Union[MetricKind, Callable, str], optional
|
|
199
382
|
Kind of the distance function, or the Numba `cfunc` JIT-compiled object.
|
|
200
383
|
Possible `MetricKind` values: IP, Cos, L2sq, Haversine, Pearson,
|
|
201
384
|
Hamming, Tanimoto, Sorensen.
|
|
202
|
-
Not every kind is JIT-able. For Jaccard distance, use `SparseIndex`.
|
|
203
385
|
|
|
204
386
|
:param dtype: Scalar type for internal vector storage, defaults to None
|
|
205
387
|
:type dtype: Optional[Union[str, ScalarKind]], optional
|
|
206
|
-
For continuous metrics can be: f16, f32, f64, or
|
|
388
|
+
For continuous metrics can be: f16, f32, f64, or i8.
|
|
207
389
|
For bitwise metrics it's implementation-defined, and can't change.
|
|
208
390
|
Example: you can use the `f16` index with `f32` vectors in Euclidean space,
|
|
209
391
|
which will be automatically downcasted.
|
|
@@ -226,15 +408,23 @@ class Index:
|
|
|
226
408
|
nearest neighbors. The original paper calls it "ef".
|
|
227
409
|
Can be changed afterwards, as the `.expansion_search`.
|
|
228
410
|
|
|
229
|
-
:param tune: Automatically adjusts hyper-parameters, defaults to False
|
|
230
|
-
:type tune: bool, optional
|
|
231
|
-
|
|
232
411
|
:param path: Where to store the index, defaults to None
|
|
233
412
|
:type path: Optional[os.PathLike], optional
|
|
234
413
|
:param view: Are we simply viewing an immutable index, defaults to False
|
|
235
414
|
:type view: bool, optional
|
|
236
415
|
"""
|
|
237
416
|
|
|
417
|
+
if connectivity is None:
|
|
418
|
+
connectivity = DEFAULT_CONNECTIVITY
|
|
419
|
+
if expansion_add is None:
|
|
420
|
+
expansion_add = DEFAULT_EXPANSION_ADD
|
|
421
|
+
if expansion_search is None:
|
|
422
|
+
expansion_search = DEFAULT_EXPANSION_SEARCH
|
|
423
|
+
|
|
424
|
+
assert isinstance(connectivity, int), "Expects integer `connectivity`"
|
|
425
|
+
assert isinstance(expansion_add, int), "Expects integer `expansion_add`"
|
|
426
|
+
assert isinstance(expansion_search, int), "Expects integer `expansion_search`"
|
|
427
|
+
|
|
238
428
|
metric = _normalize_metric(metric)
|
|
239
429
|
if isinstance(metric, MetricKind):
|
|
240
430
|
self._metric_kind = metric
|
|
@@ -256,13 +446,12 @@ class Index:
|
|
|
256
446
|
self._compiled = _CompiledIndex(
|
|
257
447
|
ndim=ndim,
|
|
258
448
|
dtype=dtype,
|
|
259
|
-
metric=self._metric_kind,
|
|
260
|
-
metric_pointer=self._metric_pointer,
|
|
261
|
-
metric_signature=self._metric_signature,
|
|
262
449
|
connectivity=connectivity,
|
|
263
450
|
expansion_add=expansion_add,
|
|
264
451
|
expansion_search=expansion_search,
|
|
265
|
-
|
|
452
|
+
metric_kind=self._metric_kind,
|
|
453
|
+
metric_pointer=self._metric_pointer,
|
|
454
|
+
metric_signature=self._metric_signature,
|
|
266
455
|
)
|
|
267
456
|
|
|
268
457
|
self.path = path
|
|
@@ -273,31 +462,30 @@ class Index:
|
|
|
273
462
|
self._compiled.load(path)
|
|
274
463
|
|
|
275
464
|
@staticmethod
|
|
276
|
-
def metadata(path: os.PathLike) ->
|
|
277
|
-
|
|
465
|
+
def metadata(path: os.PathLike) -> Optional[dict]:
|
|
466
|
+
if not os.path.exists(path):
|
|
467
|
+
return None
|
|
468
|
+
try:
|
|
469
|
+
return _index_dense_metadata(path)
|
|
470
|
+
except Exception:
|
|
471
|
+
return None
|
|
278
472
|
|
|
279
473
|
@staticmethod
|
|
280
|
-
def restore(path: os.PathLike, view: bool = False) -> Index:
|
|
474
|
+
def restore(path: os.PathLike, view: bool = False) -> Optional[Index]:
|
|
281
475
|
meta = Index.metadata(path)
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
ScalarKind.F16: 16,
|
|
285
|
-
ScalarKind.F32: 32,
|
|
286
|
-
ScalarKind.F64: 64,
|
|
287
|
-
ScalarKind.B1: 1,
|
|
288
|
-
}[meta.scalar_kind]
|
|
289
|
-
ndim = meta.bytes_for_vectors * 8 // meta.size // bits_per_scalar
|
|
476
|
+
if not meta:
|
|
477
|
+
return None
|
|
290
478
|
return Index(
|
|
291
|
-
ndim=
|
|
292
|
-
|
|
293
|
-
metric=meta
|
|
479
|
+
ndim=meta["dimensions"],
|
|
480
|
+
dtype=meta["kind_scalar"],
|
|
481
|
+
metric=meta["kind_metric"],
|
|
294
482
|
path=path,
|
|
295
483
|
view=view,
|
|
296
484
|
)
|
|
297
485
|
|
|
298
486
|
def add(
|
|
299
487
|
self,
|
|
300
|
-
|
|
488
|
+
keys,
|
|
301
489
|
vectors,
|
|
302
490
|
*,
|
|
303
491
|
copy: bool = True,
|
|
@@ -307,22 +495,22 @@ class Index:
|
|
|
307
495
|
) -> Union[int, np.ndarray]:
|
|
308
496
|
"""Inserts one or move vectors into the index.
|
|
309
497
|
|
|
310
|
-
For maximal performance the `
|
|
498
|
+
For maximal performance the `keys` and `vectors`
|
|
311
499
|
should conform to the Python's "buffer protocol" spec.
|
|
312
500
|
|
|
313
501
|
To index a single entry:
|
|
314
|
-
|
|
502
|
+
keys: int, vectors: np.ndarray.
|
|
315
503
|
To index many entries:
|
|
316
|
-
|
|
504
|
+
keys: np.ndarray, vectors: np.ndarray.
|
|
317
505
|
|
|
318
506
|
When working with extremely large indexes, you may want to
|
|
319
507
|
pass `copy=False`, if you can guarantee the lifetime of the
|
|
320
508
|
primary vectors store during the process of construction.
|
|
321
509
|
|
|
322
|
-
:param
|
|
323
|
-
:type
|
|
510
|
+
:param keys: Unique identifier(s) for passed vectors, optional
|
|
511
|
+
:type keys: np.ndarray
|
|
324
512
|
:param vectors: Vector or a row-major matrix
|
|
325
|
-
:type vectors:
|
|
513
|
+
:type vectors: np.ndarray
|
|
326
514
|
:param copy: Should the index store a copy of vectors, defaults to True
|
|
327
515
|
:type copy: bool, optional
|
|
328
516
|
:param threads: Optimal number of cores to use, defaults to 0
|
|
@@ -331,66 +519,18 @@ class Index:
|
|
|
331
519
|
:type log: Union[str, bool], optional
|
|
332
520
|
:param batch_size: Number of vectors to process at once, defaults to 0
|
|
333
521
|
:type batch_size: int, optional
|
|
334
|
-
:return: Inserted
|
|
522
|
+
:return: Inserted key or keys
|
|
335
523
|
:type: Union[int, np.ndarray]
|
|
336
524
|
"""
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
if generate_labels:
|
|
347
|
-
start_id = len(self._compiled)
|
|
348
|
-
if is_multiple:
|
|
349
|
-
labels = np.arange(start_id, start_id + count_vectors, dtype=Label)
|
|
350
|
-
else:
|
|
351
|
-
labels = start_id
|
|
352
|
-
else:
|
|
353
|
-
if isinstance(labels, np.ndarray):
|
|
354
|
-
if not is_multiple:
|
|
355
|
-
labels = int(labels[0])
|
|
356
|
-
else:
|
|
357
|
-
labels = labels.astype(Label)
|
|
358
|
-
count_labels = len(labels) if isinstance(labels, Iterable) else 1
|
|
359
|
-
assert count_labels == count_vectors
|
|
360
|
-
|
|
361
|
-
# If logging is requested, and batch size is undefined, set it to grow 1% at a time:
|
|
362
|
-
if log and batch_size == 0:
|
|
363
|
-
batch_size = int(math.ceil(count_vectors / 100))
|
|
364
|
-
|
|
365
|
-
# Split into batches and log progress, if needed
|
|
366
|
-
if batch_size:
|
|
367
|
-
labels = [
|
|
368
|
-
labels[start_row : start_row + batch_size]
|
|
369
|
-
for start_row in range(0, count_vectors, batch_size)
|
|
370
|
-
]
|
|
371
|
-
vectors = [
|
|
372
|
-
vectors[start_row : start_row + batch_size, :]
|
|
373
|
-
for start_row in range(0, count_vectors, batch_size)
|
|
374
|
-
]
|
|
375
|
-
tasks = zip(labels, vectors)
|
|
376
|
-
name = log if isinstance(log, str) else "Add"
|
|
377
|
-
pbar = tqdm(
|
|
378
|
-
tasks,
|
|
379
|
-
desc=name,
|
|
380
|
-
total=count_vectors,
|
|
381
|
-
unit="Vector",
|
|
382
|
-
disable=log is False,
|
|
383
|
-
)
|
|
384
|
-
for labels, vectors in tasks:
|
|
385
|
-
self._compiled.add(labels, vectors, copy=copy, threads=threads)
|
|
386
|
-
pbar.update(len(labels))
|
|
387
|
-
|
|
388
|
-
pbar.close()
|
|
389
|
-
|
|
390
|
-
else:
|
|
391
|
-
self._compiled.add(labels, vectors, copy=copy, threads=threads)
|
|
392
|
-
|
|
393
|
-
return labels
|
|
525
|
+
return _add_to_compiled(
|
|
526
|
+
compiled=self._compiled,
|
|
527
|
+
keys=keys,
|
|
528
|
+
vectors=vectors,
|
|
529
|
+
copy=copy,
|
|
530
|
+
threads=threads,
|
|
531
|
+
log=log,
|
|
532
|
+
batch_size=batch_size,
|
|
533
|
+
)
|
|
394
534
|
|
|
395
535
|
def search(
|
|
396
536
|
self,
|
|
@@ -401,12 +541,12 @@ class Index:
|
|
|
401
541
|
exact: bool = False,
|
|
402
542
|
log: Union[str, bool] = False,
|
|
403
543
|
batch_size: int = 0,
|
|
404
|
-
) -> Matches:
|
|
544
|
+
) -> Union[Matches, BatchMatches]:
|
|
405
545
|
"""
|
|
406
546
|
Performs approximate nearest neighbors search for one or more queries.
|
|
407
547
|
|
|
408
548
|
:param vectors: Query vector or vectors.
|
|
409
|
-
:type vectors:
|
|
549
|
+
:type vectors: np.ndarray
|
|
410
550
|
:param k: Upper limit on the number of matches to find, defaults to 10
|
|
411
551
|
:type k: int, optional
|
|
412
552
|
:param threads: Optimal number of cores to use, defaults to 0
|
|
@@ -417,60 +557,23 @@ class Index:
|
|
|
417
557
|
:type log: Union[str, bool], optional
|
|
418
558
|
:param batch_size: Number of vectors to process at once, defaults to 0
|
|
419
559
|
:type batch_size: int, optional
|
|
420
|
-
:return:
|
|
421
|
-
:rtype: Matches
|
|
560
|
+
:return: Matches for one or more queries
|
|
561
|
+
:rtype: Union[Matches, BatchMatches]
|
|
422
562
|
"""
|
|
423
563
|
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
vectors[start_row : start_row + batch_size, :]
|
|
434
|
-
for start_row in range(0, count_vectors, batch_size)
|
|
435
|
-
]
|
|
436
|
-
tasks_matches = []
|
|
437
|
-
name = log if isinstance(log, str) else "Search"
|
|
438
|
-
pbar = tqdm(
|
|
439
|
-
tasks,
|
|
440
|
-
desc=name,
|
|
441
|
-
total=count_vectors,
|
|
442
|
-
unit="Vector",
|
|
443
|
-
disable=log is False,
|
|
444
|
-
)
|
|
445
|
-
for vectors in tasks:
|
|
446
|
-
tuple_ = self._compiled.search(
|
|
447
|
-
vectors,
|
|
448
|
-
k,
|
|
449
|
-
exact=exact,
|
|
450
|
-
threads=threads,
|
|
451
|
-
)
|
|
452
|
-
tasks_matches.append(Matches(*tuple_))
|
|
453
|
-
pbar.update(vectors.shape[0])
|
|
454
|
-
|
|
455
|
-
pbar.close()
|
|
456
|
-
return Matches(
|
|
457
|
-
labels=np.vstack([m.labels for m in tasks_matches]),
|
|
458
|
-
distances=np.vstack([m.distances for m in tasks_matches]),
|
|
459
|
-
counts=np.concatenate([m.counts for m in tasks_matches], axis=None),
|
|
460
|
-
)
|
|
461
|
-
|
|
462
|
-
else:
|
|
463
|
-
tuple_ = self._compiled.search(
|
|
464
|
-
vectors,
|
|
465
|
-
k,
|
|
466
|
-
exact=exact,
|
|
467
|
-
threads=threads,
|
|
468
|
-
)
|
|
469
|
-
return Matches(*tuple_)
|
|
564
|
+
return _search_in_compiled(
|
|
565
|
+
compiled=self._compiled,
|
|
566
|
+
vectors=vectors,
|
|
567
|
+
k=k,
|
|
568
|
+
exact=exact,
|
|
569
|
+
threads=threads,
|
|
570
|
+
log=log,
|
|
571
|
+
batch_size=batch_size,
|
|
572
|
+
)
|
|
470
573
|
|
|
471
574
|
def remove(
|
|
472
575
|
self,
|
|
473
|
-
|
|
576
|
+
keys: Union[int, Iterable[int]],
|
|
474
577
|
*,
|
|
475
578
|
compact: bool = False,
|
|
476
579
|
threads: int = 0,
|
|
@@ -481,8 +584,8 @@ class Index:
|
|
|
481
584
|
mark some entries deleted, instead of rebuilding a filtered index.
|
|
482
585
|
In other cases, rebuilding - is the recommended approach.
|
|
483
586
|
|
|
484
|
-
:param
|
|
485
|
-
:type
|
|
587
|
+
:param keys: Unique identifier for passed vectors, optional
|
|
588
|
+
:type keys: np.ndarray
|
|
486
589
|
:param compact: Removes links to removed nodes (expensive), defaults to False
|
|
487
590
|
:type compact: bool, optional
|
|
488
591
|
:param threads: Optimal number of cores to use, defaults to 0
|
|
@@ -490,7 +593,7 @@ class Index:
|
|
|
490
593
|
:return: Number of removed entries
|
|
491
594
|
:type: Union[bool, int]
|
|
492
595
|
"""
|
|
493
|
-
return self._compiled.remove(
|
|
596
|
+
return self._compiled.remove(keys, compact=compact, threads=threads)
|
|
494
597
|
|
|
495
598
|
def rename(self, label_from: int, label_to: int) -> bool:
|
|
496
599
|
"""Relabel existing entry"""
|
|
@@ -516,16 +619,16 @@ class Index:
|
|
|
516
619
|
def __len__(self) -> int:
|
|
517
620
|
return self._compiled.__len__()
|
|
518
621
|
|
|
519
|
-
def __delitem__(self,
|
|
520
|
-
raise self.remove(
|
|
622
|
+
def __delitem__(self, key: int) -> bool:
|
|
623
|
+
raise self.remove(key)
|
|
521
624
|
|
|
522
|
-
def __contains__(self,
|
|
523
|
-
return self._compiled.__contains__(
|
|
625
|
+
def __contains__(self, key: int) -> bool:
|
|
626
|
+
return self._compiled.__contains__(key)
|
|
524
627
|
|
|
525
|
-
def __getitem__(self,
|
|
628
|
+
def __getitem__(self, key: int) -> np.ndarray:
|
|
526
629
|
dtype = self.dtype
|
|
527
630
|
get_dtype = _to_numpy_compatible_dtype(dtype)
|
|
528
|
-
vector = self._compiled.__getitem__(
|
|
631
|
+
vector = self._compiled.__getitem__(key, get_dtype)
|
|
529
632
|
view_dtype = _to_numpy_dtype(dtype)
|
|
530
633
|
return None if vector is None else vector.view(view_dtype)
|
|
531
634
|
|
|
@@ -545,6 +648,28 @@ class Index:
|
|
|
545
648
|
def metric(self) -> Union[MetricKind, CompiledMetric]:
|
|
546
649
|
return self._metric_jit if self._metric_jit else self._metric_kind
|
|
547
650
|
|
|
651
|
+
@metric.setter
|
|
652
|
+
def metric(self, metric: Union[str, MetricKind, CompiledMetric]):
|
|
653
|
+
metric = _normalize_metric(metric)
|
|
654
|
+
if isinstance(metric, MetricKind):
|
|
655
|
+
metric_kind = metric
|
|
656
|
+
metric_pointer = 0
|
|
657
|
+
metric_signature = MetricSignature.ArrayArraySize
|
|
658
|
+
elif isinstance(metric, CompiledMetric):
|
|
659
|
+
metric_kind = metric.kind
|
|
660
|
+
metric_pointer = metric.pointer
|
|
661
|
+
metric_signature = metric.signature
|
|
662
|
+
else:
|
|
663
|
+
raise ValueError(
|
|
664
|
+
"The `metric` must be a `CompiledMetric` or a `MetricKind`"
|
|
665
|
+
)
|
|
666
|
+
|
|
667
|
+
return self._compiled.change_metric(
|
|
668
|
+
metric_kind=metric_kind,
|
|
669
|
+
metric_pointer=metric_pointer,
|
|
670
|
+
metric_signature=metric_signature,
|
|
671
|
+
)
|
|
672
|
+
|
|
548
673
|
@property
|
|
549
674
|
def dtype(self) -> ScalarKind:
|
|
550
675
|
return self._compiled.dtype
|
|
@@ -581,25 +706,27 @@ class Index:
|
|
|
581
706
|
path = path if path else self.path
|
|
582
707
|
if path is None:
|
|
583
708
|
raise Exception("Define `path` argument")
|
|
584
|
-
self._compiled.save(path)
|
|
709
|
+
self._compiled.save(os.fspath(path))
|
|
585
710
|
|
|
586
711
|
def load(self, path: Optional[os.PathLike] = None):
|
|
587
712
|
path = path if path else self.path
|
|
588
713
|
if path is None:
|
|
589
714
|
raise Exception("Define `path` argument")
|
|
590
|
-
self._compiled.load(path)
|
|
715
|
+
self._compiled.load(os.fspath(path))
|
|
591
716
|
|
|
592
717
|
def view(self, path: Optional[os.PathLike] = None):
|
|
593
718
|
path = path if path else self.path
|
|
594
719
|
if path is None:
|
|
595
720
|
raise Exception("Define `path` argument")
|
|
596
|
-
self._compiled.view(path)
|
|
721
|
+
self._compiled.view(os.fspath(path))
|
|
597
722
|
|
|
598
723
|
def clear(self):
|
|
724
|
+
"""Erases all the vectors from the index, preserving the space for future insertions."""
|
|
599
725
|
self._compiled.clear()
|
|
600
726
|
|
|
601
|
-
def
|
|
602
|
-
|
|
727
|
+
def reset(self):
|
|
728
|
+
"""Erases all members from index, closing files, and returning RAM to OS."""
|
|
729
|
+
self._compiled.reset()
|
|
603
730
|
|
|
604
731
|
def copy(self) -> Index:
|
|
605
732
|
result = Index(
|
|
@@ -626,7 +753,7 @@ class Index:
|
|
|
626
753
|
:type max_proposals: int, optional
|
|
627
754
|
:param exact: Controls if underlying `search` should be exact, defaults to False
|
|
628
755
|
:type exact: bool, optional
|
|
629
|
-
:return: Mapping from
|
|
756
|
+
:return: Mapping from keys of `self` to keys of `other`
|
|
630
757
|
:rtype: dict
|
|
631
758
|
"""
|
|
632
759
|
return self._compiled.join(
|
|
@@ -635,49 +762,71 @@ class Index:
|
|
|
635
762
|
exact=exact,
|
|
636
763
|
)
|
|
637
764
|
|
|
638
|
-
def
|
|
765
|
+
def get_keys(self, offset: int = 0, limit: int = 0) -> np.ndarray:
|
|
639
766
|
if limit == 0:
|
|
640
767
|
limit = 2**63 - 1
|
|
641
|
-
return self._compiled.
|
|
768
|
+
return self._compiled.get_keys(offset, limit)
|
|
642
769
|
|
|
643
770
|
@property
|
|
644
|
-
def
|
|
645
|
-
"""Retrieves the
|
|
771
|
+
def keys(self) -> np.ndarray:
|
|
772
|
+
"""Retrieves the keys of all vectors present in `self`
|
|
646
773
|
|
|
647
|
-
:return: Array of
|
|
774
|
+
:return: Array of keys
|
|
648
775
|
:rtype: np.ndarray
|
|
649
776
|
"""
|
|
650
|
-
return self._compiled.
|
|
777
|
+
return self._compiled.keys
|
|
651
778
|
|
|
652
779
|
def get_vectors(
|
|
653
780
|
self,
|
|
654
|
-
|
|
781
|
+
keys: np.ndarray,
|
|
655
782
|
dtype: ScalarKind = ScalarKind.F32,
|
|
656
783
|
) -> np.ndarray:
|
|
657
|
-
"""Retrieves vectors associated with given `
|
|
784
|
+
"""Retrieves vectors associated with given `keys`
|
|
658
785
|
|
|
659
786
|
:return: Matrix of vectors (row-major)
|
|
660
787
|
:rtype: np.ndarray
|
|
661
788
|
"""
|
|
662
789
|
dtype = _normalize_dtype(dtype, self._metric_kind)
|
|
663
790
|
get_dtype = _to_numpy_compatible_dtype(dtype)
|
|
664
|
-
vectors = np.vstack([self._compiled.__getitem__(l, get_dtype) for l in
|
|
791
|
+
vectors = np.vstack([self._compiled.__getitem__(l, get_dtype) for l in keys])
|
|
665
792
|
view_dtype = _to_numpy_dtype(dtype)
|
|
666
793
|
return vectors.view(view_dtype)
|
|
667
794
|
|
|
668
795
|
@property
|
|
669
796
|
def vectors(self) -> np.ndarray:
|
|
670
|
-
return self.get_vectors(self.
|
|
797
|
+
return self.get_vectors(self.keys, self.dtype)
|
|
671
798
|
|
|
672
799
|
@property
|
|
673
800
|
def max_level(self) -> int:
|
|
674
801
|
return self._compiled.max_level
|
|
675
802
|
|
|
676
803
|
@property
|
|
677
|
-
def levels_stats(self) ->
|
|
804
|
+
def levels_stats(self) -> _CompiledIndexStats:
|
|
805
|
+
"""Get the accumulated statistics for the entire multi-level graph.
|
|
806
|
+
|
|
807
|
+
:return: Statistics for the entire multi-level graph.
|
|
808
|
+
:rtype: _CompiledIndexStats
|
|
809
|
+
|
|
810
|
+
Statistics:
|
|
811
|
+
- ``nodes`` (int): The number of nodes in that level.
|
|
812
|
+
- ``edges`` (int): The number of edges in that level.
|
|
813
|
+
- ``max_edges`` (int): The maximum possible number of edges in that level.
|
|
814
|
+
- ``allocated_bytes`` (int): The amount of allocated memory for that level.
|
|
815
|
+
"""
|
|
678
816
|
return self._compiled.levels_stats
|
|
679
817
|
|
|
680
|
-
def level_stats(self, level: int) ->
|
|
818
|
+
def level_stats(self, level: int) -> _CompiledIndexStats:
|
|
819
|
+
"""Get statistics for one level of the index - one graph.
|
|
820
|
+
|
|
821
|
+
:return: Statistics for one level of the index - one graph.
|
|
822
|
+
:rtype: _CompiledIndexStats
|
|
823
|
+
|
|
824
|
+
Statistics:
|
|
825
|
+
- ``nodes`` (int): The number of nodes in that level.
|
|
826
|
+
- ``edges`` (int): The number of edges in that level.
|
|
827
|
+
- ``max_edges`` (int): The maximum possible number of edges in that level.
|
|
828
|
+
- ``allocated_bytes`` (int): The amount of allocated memory for that level.
|
|
829
|
+
"""
|
|
681
830
|
return self._compiled.level_stats(level)
|
|
682
831
|
|
|
683
832
|
def __repr__(self) -> str:
|
|
@@ -711,3 +860,146 @@ class Index:
|
|
|
711
860
|
*level_stats,
|
|
712
861
|
]
|
|
713
862
|
)
|
|
863
|
+
|
|
864
|
+
|
|
865
|
+
class Indexes:
|
|
866
|
+
def __init__(
|
|
867
|
+
self,
|
|
868
|
+
indexes: Iterable[Index] = [],
|
|
869
|
+
paths: Iterable[os.PathLike] = [],
|
|
870
|
+
view: bool = False,
|
|
871
|
+
threads: int = 0,
|
|
872
|
+
) -> None:
|
|
873
|
+
self._compiled = _CompiledIndexes()
|
|
874
|
+
for index in indexes:
|
|
875
|
+
self._compiled.add(index._compiled)
|
|
876
|
+
self._compiled.add_paths(paths, view=view, threads=threads)
|
|
877
|
+
|
|
878
|
+
def add(self, index: Index):
|
|
879
|
+
self._compiled.add(index._compiled)
|
|
880
|
+
|
|
881
|
+
def __len__(self) -> int:
|
|
882
|
+
return self._compiled.__len__()
|
|
883
|
+
|
|
884
|
+
def search(
|
|
885
|
+
self,
|
|
886
|
+
vectors,
|
|
887
|
+
k: int = 10,
|
|
888
|
+
*,
|
|
889
|
+
threads: int = 0,
|
|
890
|
+
exact: bool = False,
|
|
891
|
+
):
|
|
892
|
+
return _search_in_compiled(
|
|
893
|
+
compiled=self._compiled,
|
|
894
|
+
vectors=vectors,
|
|
895
|
+
k=k,
|
|
896
|
+
exact=exact,
|
|
897
|
+
threads=threads,
|
|
898
|
+
log=False,
|
|
899
|
+
batch_size=None,
|
|
900
|
+
)
|
|
901
|
+
|
|
902
|
+
|
|
903
|
+
def search(
|
|
904
|
+
dataset: np.ndarray,
|
|
905
|
+
query: np.ndarray,
|
|
906
|
+
k: int = 10,
|
|
907
|
+
metric: Union[str, MetricKind, CompiledMetric] = MetricKind.Cos,
|
|
908
|
+
*,
|
|
909
|
+
exact: bool = False,
|
|
910
|
+
threads: int = 0,
|
|
911
|
+
log: Union[str, bool] = False,
|
|
912
|
+
batch_size: int = 0,
|
|
913
|
+
) -> Union[Matches, BatchMatches]:
|
|
914
|
+
"""Shortcut for search, that can avoid index construction. Particularly useful for
|
|
915
|
+
tiny datasets, where brute-force exact search works fast enough.
|
|
916
|
+
|
|
917
|
+
:param dataset: Row-major matrix.
|
|
918
|
+
:type dataset: np.ndarray
|
|
919
|
+
:param query: Query vector or vectors (also row-major), to find in `dataset`.
|
|
920
|
+
:type query: np.ndarray
|
|
921
|
+
|
|
922
|
+
:param k: Upper limit on the number of matches to find, defaults to 10
|
|
923
|
+
:type k: int, optional
|
|
924
|
+
|
|
925
|
+
:param metric: Distance function, defaults to MetricKind.Cos
|
|
926
|
+
:type metric: Union[MetricKind, Callable, str], optional
|
|
927
|
+
Kind of the distance function, or the Numba `cfunc` JIT-compiled object.
|
|
928
|
+
Possible `MetricKind` values: IP, Cos, L2sq, Haversine, Pearson,
|
|
929
|
+
Hamming, Tanimoto, Sorensen.
|
|
930
|
+
|
|
931
|
+
:param threads: Optimal number of cores to use, defaults to 0
|
|
932
|
+
:type threads: int, optional
|
|
933
|
+
:param exact: Perform exhaustive linear-time exact search, defaults to False
|
|
934
|
+
:type exact: bool, optional
|
|
935
|
+
:param log: Whether to print the progress bar, default to False
|
|
936
|
+
:type log: Union[str, bool], optional
|
|
937
|
+
:param batch_size: Number of vectors to process at once, defaults to 0
|
|
938
|
+
:type batch_size: int, optional
|
|
939
|
+
|
|
940
|
+
:return: Matches for one or more queries
|
|
941
|
+
:rtype: Union[Matches, BatchMatches]
|
|
942
|
+
"""
|
|
943
|
+
assert dataset.ndim == 2, "Dataset must be a matrix, with a vector in each row"
|
|
944
|
+
|
|
945
|
+
if not exact:
|
|
946
|
+
index = Index(
|
|
947
|
+
dataset.shape[1],
|
|
948
|
+
metric=metric,
|
|
949
|
+
dtype=dataset.dtype,
|
|
950
|
+
)
|
|
951
|
+
index.add(
|
|
952
|
+
None,
|
|
953
|
+
dataset,
|
|
954
|
+
threads=threads,
|
|
955
|
+
log=log,
|
|
956
|
+
batch_size=batch_size,
|
|
957
|
+
)
|
|
958
|
+
return index.search(
|
|
959
|
+
query,
|
|
960
|
+
k,
|
|
961
|
+
threads=threads,
|
|
962
|
+
log=log,
|
|
963
|
+
batch_size=batch_size,
|
|
964
|
+
)
|
|
965
|
+
|
|
966
|
+
metric = _normalize_metric(metric)
|
|
967
|
+
if isinstance(metric, MetricKind):
|
|
968
|
+
metric_kind = metric
|
|
969
|
+
metric_pointer = 0
|
|
970
|
+
metric_signature = MetricSignature.ArrayArraySize
|
|
971
|
+
elif isinstance(metric, CompiledMetric):
|
|
972
|
+
metric_kind = metric.kind
|
|
973
|
+
metric_pointer = metric.pointer
|
|
974
|
+
metric_signature = metric.signature
|
|
975
|
+
else:
|
|
976
|
+
raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`")
|
|
977
|
+
|
|
978
|
+
class WrappedDataset:
|
|
979
|
+
def __init__(self) -> None:
|
|
980
|
+
pass
|
|
981
|
+
|
|
982
|
+
def search(self, query, k, **kwargs):
|
|
983
|
+
kwargs.pop("exact")
|
|
984
|
+
kwargs.update(
|
|
985
|
+
dict(
|
|
986
|
+
metric_kind=metric_kind,
|
|
987
|
+
metric_pointer=metric_pointer,
|
|
988
|
+
metric_signature=metric_signature,
|
|
989
|
+
)
|
|
990
|
+
)
|
|
991
|
+
assert dataset.shape[1] == query.shape[1], "Number of dimensions differs"
|
|
992
|
+
if dataset.dtype != query.dtype:
|
|
993
|
+
query = query.astype(dataset.dtype)
|
|
994
|
+
|
|
995
|
+
return _exact_search(dataset, query, k, **kwargs)
|
|
996
|
+
|
|
997
|
+
return _search_in_compiled(
|
|
998
|
+
compiled=WrappedDataset(),
|
|
999
|
+
vectors=query,
|
|
1000
|
+
k=k,
|
|
1001
|
+
threads=threads,
|
|
1002
|
+
exact=True,
|
|
1003
|
+
log=log,
|
|
1004
|
+
batch_size=batch_size,
|
|
1005
|
+
)
|