usearch 0.1.10 → 0.2.2
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 +12 -1
- package/.vscode/launch.json +5 -4
- package/.vscode/settings.json +7 -1
- package/CMakeLists.txt +11 -6
- package/CMakeListsInstructions.md +39 -17
- package/Cargo.toml +1 -1
- package/Package.swift +40 -0
- package/README.md +36 -18
- package/VERSION +1 -1
- package/build.gradle +5 -1
- package/cmake/jemalloc.cmake +0 -2
- package/conanfile.py +18 -0
- package/include/usearch/usearch.hpp +337 -228
- package/javascript/lib.cpp +8 -2
- package/objc/USearchObjective.mm +150 -0
- package/objc/include/USearchObjective.h +69 -0
- package/package.json +1 -1
- package/pyproject.toml +18 -4
- package/python/bench.py +14 -2
- package/python/lib.cpp +353 -43
- package/python/test.py +10 -4
- package/python/usearch/__init__.py +0 -0
- package/rust/lib.cpp +13 -2
- package/setup.py +8 -3
- package/simsimd/include/simsimd/simsimd.h +15 -0
- package/src/advanced.hpp +232 -122
- package/src/bench.cpp +249 -79
- package/swift/Index+Sugar.swift +70 -0
- package/swift/Test.swift +20 -0
- package/swift/USearch.swift +8 -0
- package/java/gradle/wrapper/gradle-wrapper.jar +0 -0
- package/java/gradle/wrapper/gradle-wrapper.properties +0 -6
- package/java/gradlew +0 -244
package/python/lib.cpp
CHANGED
|
@@ -24,11 +24,41 @@ using namespace unum::usearch;
|
|
|
24
24
|
using namespace unum;
|
|
25
25
|
|
|
26
26
|
namespace py = pybind11;
|
|
27
|
+
using py_shape_t = py::array::ShapeContainer;
|
|
27
28
|
|
|
28
29
|
using label_t = Py_ssize_t;
|
|
29
30
|
using distance_t = punned_distance_t;
|
|
31
|
+
using id_t = std::uint32_t;
|
|
30
32
|
using native_index_t = auto_index_gt<label_t>;
|
|
31
33
|
|
|
34
|
+
using set_member_t = std::uint32_t;
|
|
35
|
+
using set_view_t = span_gt<set_member_t const>;
|
|
36
|
+
using sets_index_t = index_gt<jaccard_gt<set_member_t>, label_t, id_t, set_member_t>;
|
|
37
|
+
|
|
38
|
+
using hash_word_t = std::uint64_t;
|
|
39
|
+
using hash_index_t = index_gt<bit_hamming_gt<hash_word_t>, label_t, id_t, hash_word_t>;
|
|
40
|
+
static constexpr std::size_t bits_per_hash_word_k = sizeof(hash_word_t) * CHAR_BIT;
|
|
41
|
+
|
|
42
|
+
struct hash_index_w_meta_t : public hash_index_t {
|
|
43
|
+
using hash_index_t::add;
|
|
44
|
+
using hash_index_t::capacity;
|
|
45
|
+
using hash_index_t::reserve;
|
|
46
|
+
using hash_index_t::search;
|
|
47
|
+
using hash_index_t::size;
|
|
48
|
+
|
|
49
|
+
std::vector<hash_word_t> buffer_;
|
|
50
|
+
std::size_t words_;
|
|
51
|
+
std::size_t bits_;
|
|
52
|
+
|
|
53
|
+
hash_index_w_meta_t(config_t config, std::size_t bits) : hash_index_t(config) {
|
|
54
|
+
words_ = divide_round_up<bits_per_hash_word_k>(bits);
|
|
55
|
+
bits_ = words_ * bits_per_hash_word_k;
|
|
56
|
+
buffer_.resize(words_);
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
span_gt<hash_word_t const> buffer() const noexcept { return {buffer_.data(), buffer_.size()}; }
|
|
60
|
+
};
|
|
61
|
+
|
|
32
62
|
static native_index_t make_index( //
|
|
33
63
|
std::size_t dimensions, //
|
|
34
64
|
std::size_t capacity, //
|
|
@@ -56,33 +86,91 @@ static native_index_t make_index( //
|
|
|
56
86
|
return index_from_name<native_index_t>(metric.c_str(), metric.size(), dimensions, accuracy, config);
|
|
57
87
|
}
|
|
58
88
|
|
|
59
|
-
static
|
|
89
|
+
static std::unique_ptr<sets_index_t> make_sets_index( //
|
|
90
|
+
std::size_t capacity, //
|
|
91
|
+
std::size_t connectivity, //
|
|
92
|
+
std::size_t expansion_add, //
|
|
93
|
+
std::size_t expansion_search //
|
|
94
|
+
) {
|
|
95
|
+
config_t config;
|
|
96
|
+
config.expansion_add = expansion_add;
|
|
97
|
+
config.expansion_search = expansion_search;
|
|
98
|
+
config.connectivity = connectivity;
|
|
99
|
+
config.max_elements = capacity;
|
|
100
|
+
config.max_threads_add = 1;
|
|
101
|
+
config.max_threads_search = 1;
|
|
102
|
+
|
|
103
|
+
return std::unique_ptr<sets_index_t>(new sets_index_t(config));
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
static std::unique_ptr<hash_index_w_meta_t> make_hash_index( //
|
|
107
|
+
std::size_t bits, //
|
|
108
|
+
std::size_t capacity, //
|
|
109
|
+
std::size_t connectivity, //
|
|
110
|
+
std::size_t expansion_add, //
|
|
111
|
+
std::size_t expansion_search //
|
|
112
|
+
) {
|
|
113
|
+
config_t config;
|
|
114
|
+
config.expansion_add = expansion_add;
|
|
115
|
+
config.expansion_search = expansion_search;
|
|
116
|
+
config.connectivity = connectivity;
|
|
117
|
+
config.max_elements = capacity;
|
|
118
|
+
config.max_threads_add = 1;
|
|
119
|
+
config.max_threads_search = 1;
|
|
120
|
+
|
|
121
|
+
return std::unique_ptr<hash_index_w_meta_t>(new hash_index_w_meta_t(config, bits));
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
static void add_one_to_index(native_index_t& index, label_t label, py::buffer vector, bool copy) {
|
|
125
|
+
|
|
126
|
+
py::buffer_info vector_info = vector.request();
|
|
127
|
+
if (vector_info.ndim != 1)
|
|
128
|
+
throw std::invalid_argument("Expects a vector, not a higher-rank tensor!");
|
|
129
|
+
|
|
130
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
131
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
132
|
+
if (vector_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
133
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
134
|
+
|
|
135
|
+
if (index.size() + 1 >= index.capacity())
|
|
136
|
+
index.reserve(ceil2(index.size() + 1));
|
|
137
|
+
|
|
138
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
139
|
+
if (vector_info.format == "e")
|
|
140
|
+
index.add(label, reinterpret_cast<f16_bits_t const*>(vector_data), 0, copy);
|
|
141
|
+
else if (vector_info.format == "f")
|
|
142
|
+
index.add(label, reinterpret_cast<float const*>(vector_data), 0, copy);
|
|
143
|
+
else if (vector_info.format == "d")
|
|
144
|
+
index.add(label, reinterpret_cast<double const*>(vector_data), 0, copy);
|
|
145
|
+
else
|
|
146
|
+
throw std::invalid_argument("Incompatible scalars in the vector!");
|
|
147
|
+
}
|
|
148
|
+
|
|
149
|
+
static void add_many_to_index(native_index_t& index, py::buffer labels, py::buffer vectors, bool copy) {
|
|
60
150
|
|
|
61
151
|
py::buffer_info labels_info = labels.request();
|
|
62
152
|
py::buffer_info vectors_info = vectors.request();
|
|
63
153
|
|
|
64
154
|
if (labels_info.format != py::format_descriptor<label_t>::format())
|
|
65
|
-
throw std::
|
|
155
|
+
throw std::invalid_argument("Incompatible label type!");
|
|
66
156
|
|
|
67
157
|
if (labels_info.ndim != 1)
|
|
68
|
-
throw std::
|
|
158
|
+
throw std::invalid_argument("Labels must be placed in a single-dimensional array!");
|
|
69
159
|
|
|
70
160
|
if (vectors_info.ndim != 2)
|
|
71
|
-
throw std::
|
|
161
|
+
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
72
162
|
|
|
73
163
|
ssize_t labels_count = labels_info.shape[0];
|
|
74
164
|
ssize_t vectors_count = vectors_info.shape[0];
|
|
75
165
|
ssize_t vectors_dimensions = vectors_info.shape[1];
|
|
76
|
-
if (vectors_dimensions != index.dimensions())
|
|
77
|
-
throw std::
|
|
166
|
+
if (vectors_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
167
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
78
168
|
|
|
79
169
|
if (labels_count != vectors_count)
|
|
80
|
-
throw std::
|
|
170
|
+
throw std::invalid_argument("Number of labels and vectors must match!");
|
|
81
171
|
|
|
82
|
-
if (index.size() + vectors_count >= index.capacity())
|
|
83
|
-
|
|
84
|
-
index.reserve(next_capacity);
|
|
85
|
-
}
|
|
172
|
+
if (index.size() + vectors_count >= index.capacity())
|
|
173
|
+
index.reserve(ceil2(index.size() + vectors_count));
|
|
86
174
|
|
|
87
175
|
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
88
176
|
char const* labels_data = reinterpret_cast<char const*>(labels_info.ptr);
|
|
@@ -91,8 +179,8 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
|
|
|
91
179
|
if (vectors_info.format == "e")
|
|
92
180
|
multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
93
181
|
label_t label = *reinterpret_cast<label_t const*>(labels_data + task_idx * labels_info.strides[0]);
|
|
94
|
-
|
|
95
|
-
reinterpret_cast<
|
|
182
|
+
f16_bits_t const* vector =
|
|
183
|
+
reinterpret_cast<f16_bits_t const*>(vectors_data + task_idx * vectors_info.strides[0]);
|
|
96
184
|
index.add(label, vector, thread_idx, copy);
|
|
97
185
|
});
|
|
98
186
|
else if (vectors_info.format == "f")
|
|
@@ -108,11 +196,48 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
|
|
|
108
196
|
index.add(label, vector, thread_idx, copy);
|
|
109
197
|
});
|
|
110
198
|
else
|
|
111
|
-
throw std::
|
|
199
|
+
throw std::invalid_argument("Incompatible scalars in the vectors matrix!");
|
|
200
|
+
}
|
|
201
|
+
|
|
202
|
+
static py::tuple search_one_in_index(native_index_t& index, py::buffer vector, std::size_t wanted) {
|
|
203
|
+
|
|
204
|
+
py::buffer_info vector_info = vector.request();
|
|
205
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
206
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
207
|
+
if (vector_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
208
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
209
|
+
|
|
210
|
+
py::array_t<label_t> labels_py(static_cast<ssize_t>(wanted));
|
|
211
|
+
py::array_t<distance_t> distances_py(static_cast<ssize_t>(wanted));
|
|
212
|
+
std::size_t count{};
|
|
213
|
+
auto labels_py1d = labels_py.mutable_unchecked<1>();
|
|
214
|
+
auto distances_py1d = distances_py.mutable_unchecked<1>();
|
|
215
|
+
|
|
216
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
217
|
+
if (vector_info.format == "e")
|
|
218
|
+
count = index.search( //
|
|
219
|
+
reinterpret_cast<f16_bits_t const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
220
|
+
else if (vector_info.format == "f")
|
|
221
|
+
count = index.search( //
|
|
222
|
+
reinterpret_cast<float const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
223
|
+
else if (vector_info.format == "d")
|
|
224
|
+
count = index.search( //
|
|
225
|
+
reinterpret_cast<double const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
226
|
+
else
|
|
227
|
+
throw std::invalid_argument("Incompatible scalars in the query vector!");
|
|
228
|
+
|
|
229
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(count)});
|
|
230
|
+
distances_py.resize(py_shape_t{static_cast<ssize_t>(count)});
|
|
231
|
+
|
|
232
|
+
py::tuple results(3);
|
|
233
|
+
results[0] = labels_py;
|
|
234
|
+
results[1] = distances_py;
|
|
235
|
+
results[2] = static_cast<Py_ssize_t>(count);
|
|
236
|
+
return results;
|
|
112
237
|
}
|
|
113
238
|
|
|
114
239
|
/**
|
|
115
|
-
* @param vectors Matrix of vectors to search for
|
|
240
|
+
* @param vectors Matrix of vectors to search for.
|
|
116
241
|
* @param wanted Number of matches per request.
|
|
117
242
|
*
|
|
118
243
|
* @return Tuple with:
|
|
@@ -120,30 +245,34 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
|
|
|
120
245
|
* 2. matrix of distances,
|
|
121
246
|
* 3. array with match counts.
|
|
122
247
|
*/
|
|
123
|
-
static py::tuple
|
|
248
|
+
static py::tuple search_many_in_index(native_index_t& index, py::buffer vectors, std::size_t wanted) {
|
|
249
|
+
|
|
250
|
+
if (wanted == 0)
|
|
251
|
+
return py::tuple(3);
|
|
124
252
|
|
|
125
253
|
py::buffer_info vectors_info = vectors.request();
|
|
254
|
+
if (vectors_info.ndim == 1)
|
|
255
|
+
return search_one_in_index(index, vectors, wanted);
|
|
126
256
|
if (vectors_info.ndim != 2)
|
|
127
|
-
throw std::
|
|
257
|
+
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
128
258
|
|
|
129
259
|
ssize_t vectors_count = vectors_info.shape[0];
|
|
130
260
|
ssize_t vectors_dimensions = vectors_info.shape[1];
|
|
131
|
-
|
|
132
|
-
|
|
261
|
+
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
262
|
+
if (vectors_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
263
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
133
264
|
|
|
134
|
-
py::array_t<label_t> labels_py({vectors_count, wanted});
|
|
135
|
-
py::array_t<distance_t> distances_py({vectors_count, wanted});
|
|
265
|
+
py::array_t<label_t> labels_py({vectors_count, static_cast<ssize_t>(wanted)});
|
|
266
|
+
py::array_t<distance_t> distances_py({vectors_count, static_cast<ssize_t>(wanted)});
|
|
136
267
|
py::array_t<Py_ssize_t> counts_py(vectors_count);
|
|
137
268
|
auto labels_py2d = labels_py.mutable_unchecked<2>();
|
|
138
269
|
auto distances_py2d = distances_py.mutable_unchecked<2>();
|
|
139
270
|
auto counts_py1d = counts_py.mutable_unchecked<1>();
|
|
140
271
|
|
|
141
|
-
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
142
|
-
|
|
143
272
|
// https://docs.python.org/3/library/struct.html#format-characters
|
|
144
273
|
if (vectors_info.format == "e")
|
|
145
274
|
multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
146
|
-
|
|
275
|
+
f16_bits_t const* vector = (f16_bits_t const*)(vectors_data + task_idx * vectors_info.strides[0]);
|
|
147
276
|
counts_py1d(task_idx) = static_cast<Py_ssize_t>(
|
|
148
277
|
index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
|
|
149
278
|
});
|
|
@@ -160,7 +289,7 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
|
|
|
160
289
|
index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
|
|
161
290
|
});
|
|
162
291
|
else
|
|
163
|
-
throw std::
|
|
292
|
+
throw std::invalid_argument("Incompatible scalars in the query matrix!");
|
|
164
293
|
|
|
165
294
|
py::tuple results(3);
|
|
166
295
|
results[0] = labels_py;
|
|
@@ -169,18 +298,87 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
|
|
|
169
298
|
return results;
|
|
170
299
|
}
|
|
171
300
|
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
301
|
+
// clang-format off
|
|
302
|
+
template <typename index_at> void save_index(index_at const& index, std::string const& path) { index.save(path.c_str()); }
|
|
303
|
+
template <typename index_at> void load_index(index_at& index, std::string const& path) { index.load(path.c_str()); }
|
|
304
|
+
template <typename index_at> void view_index(index_at& index, std::string const& path) { index.view(path.c_str()); }
|
|
305
|
+
template <typename index_at> void clear_index(index_at& index) { index.clear(); }
|
|
306
|
+
// clang-format on
|
|
307
|
+
|
|
308
|
+
template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
|
|
309
|
+
if (begin == end)
|
|
310
|
+
return false;
|
|
311
|
+
element_at const* last = begin;
|
|
312
|
+
begin++;
|
|
313
|
+
for (; begin != end; ++begin, ++last) {
|
|
314
|
+
if (*begin == *last)
|
|
315
|
+
return true;
|
|
316
|
+
}
|
|
317
|
+
return false;
|
|
318
|
+
}
|
|
319
|
+
|
|
320
|
+
void validate_set(py::array_t<set_member_t> const& set) {
|
|
321
|
+
if (set.ndim() != 1)
|
|
322
|
+
throw std::invalid_argument("Set can't be multi-dimensional!");
|
|
323
|
+
if (set.strides(0) != sizeof(set_member_t))
|
|
324
|
+
throw std::invalid_argument("Set can't be strided!");
|
|
325
|
+
auto proxy = set.unchecked<1>();
|
|
326
|
+
set_member_t const* begin = &proxy(0);
|
|
327
|
+
set_member_t const* end = begin + proxy.size();
|
|
328
|
+
if (!std::is_sorted(begin, end))
|
|
329
|
+
throw std::invalid_argument("Set must be sorted!");
|
|
330
|
+
if (has_duplicates(begin, end))
|
|
331
|
+
throw std::invalid_argument("Set must be deduplicated!");
|
|
332
|
+
}
|
|
175
333
|
|
|
176
|
-
|
|
334
|
+
inline std::uint64_t hash_ror64(std::uint64_t v, int r) noexcept { return (v >> r) | (v << (64 - r)); }
|
|
335
|
+
|
|
336
|
+
inline std::uint64_t hash(std::uint64_t v) noexcept {
|
|
337
|
+
v ^= hash_ror64(v, 25) ^ hash_ror64(v, 50);
|
|
338
|
+
v *= 0xA24BAED4963EE407UL;
|
|
339
|
+
v ^= hash_ror64(v, 24) ^ hash_ror64(v, 49);
|
|
340
|
+
v *= 0x9FB21C651E98DF25UL;
|
|
341
|
+
return v ^ v >> 28;
|
|
342
|
+
}
|
|
343
|
+
|
|
344
|
+
template <typename scalar_at>
|
|
345
|
+
inline void hash_typed_buffer(hash_index_w_meta_t& index, py::buffer_info const& vector_info) noexcept {
|
|
346
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
347
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
348
|
+
ssize_t vector_stride = vector_info.strides[0];
|
|
349
|
+
std::memset(index.buffer_.data(), 0, index.words_ * sizeof(hash_word_t));
|
|
350
|
+
|
|
351
|
+
for (std::size_t i = 0; i != static_cast<std::size_t>(vector_dimensions); ++i) {
|
|
352
|
+
scalar_at scalar = *reinterpret_cast<scalar_at const*>(vector_data + i * vector_stride);
|
|
353
|
+
std::uint64_t scalar_hash = hash(scalar);
|
|
354
|
+
index.buffer_[scalar_hash % index.words_] |= hash_word_t(1) << (scalar_hash % bits_per_hash_word_k);
|
|
355
|
+
}
|
|
356
|
+
}
|
|
357
|
+
|
|
358
|
+
void hash_buffer(hash_index_w_meta_t& index, py::buffer vector) {
|
|
359
|
+
py::buffer_info info = vector.request();
|
|
360
|
+
if (info.ndim != 1)
|
|
361
|
+
throw std::invalid_argument("Array can't be multi-dimensional!");
|
|
362
|
+
|
|
363
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
364
|
+
if (info.format == "h" || info.format == "H")
|
|
365
|
+
return hash_typed_buffer<std::uint16_t>(index, info);
|
|
366
|
+
else if (info.format == "i" || info.format == "I" || info.format == "l" || info.format == "L")
|
|
367
|
+
return hash_typed_buffer<std::uint32_t>(index, info);
|
|
368
|
+
else if (info.format == "q" || info.format == "Q" || info.format == "n" || info.format == "N")
|
|
369
|
+
return hash_typed_buffer<std::uint64_t>(index, info);
|
|
370
|
+
else
|
|
371
|
+
throw std::invalid_argument("Array elements must be 16, 32, or 64 bit hashable integers!");
|
|
372
|
+
}
|
|
373
|
+
|
|
374
|
+
PYBIND11_MODULE(index, m) {
|
|
177
375
|
m.doc() = "Unum USearch Python bindings";
|
|
178
376
|
|
|
179
377
|
auto i = py::class_<native_index_t>(m, "Index");
|
|
180
378
|
|
|
181
379
|
i.def(py::init(&make_index), //
|
|
182
380
|
py::kw_only(), //
|
|
183
|
-
py::arg("ndim"),
|
|
381
|
+
py::arg("ndim") = 0, //
|
|
184
382
|
py::arg("capacity") = 0, //
|
|
185
383
|
py::arg("dtype") = std::string("f32"), //
|
|
186
384
|
py::arg("metric") = std::string("ip"), //
|
|
@@ -190,18 +388,26 @@ PYBIND11_MODULE(usearch, m) {
|
|
|
190
388
|
py::arg("metric_pointer") = 0 //
|
|
191
389
|
);
|
|
192
390
|
|
|
193
|
-
i.def(
|
|
194
|
-
"add", &
|
|
195
|
-
py::arg("labels"),
|
|
196
|
-
py::arg("vectors"),
|
|
197
|
-
py::kw_only(),
|
|
198
|
-
py::arg("copy") = true
|
|
391
|
+
i.def( //
|
|
392
|
+
"add", &add_many_to_index, //
|
|
393
|
+
py::arg("labels"), //
|
|
394
|
+
py::arg("vectors"), //
|
|
395
|
+
py::kw_only(), //
|
|
396
|
+
py::arg("copy") = true //
|
|
397
|
+
);
|
|
398
|
+
|
|
399
|
+
i.def( //
|
|
400
|
+
"add", &add_one_to_index, //
|
|
401
|
+
py::arg("label"), //
|
|
402
|
+
py::arg("vector"), //
|
|
403
|
+
py::kw_only(), //
|
|
404
|
+
py::arg("copy") = true //
|
|
199
405
|
);
|
|
200
406
|
|
|
201
|
-
i.def(
|
|
202
|
-
"search", &
|
|
203
|
-
py::arg("
|
|
204
|
-
py::arg("count") = 10
|
|
407
|
+
i.def( //
|
|
408
|
+
"search", &search_many_in_index, //
|
|
409
|
+
py::arg("query"), //
|
|
410
|
+
py::arg("count") = 10 //
|
|
205
411
|
);
|
|
206
412
|
|
|
207
413
|
i.def("__len__", &native_index_t::size);
|
|
@@ -210,7 +416,111 @@ PYBIND11_MODULE(usearch, m) {
|
|
|
210
416
|
i.def_property_readonly("connectivity", &native_index_t::connectivity);
|
|
211
417
|
i.def_property_readonly("capacity", &native_index_t::capacity);
|
|
212
418
|
|
|
213
|
-
i.def("save", &save_index
|
|
214
|
-
i.def("load", &load_index
|
|
215
|
-
i.def("view", &view_index
|
|
419
|
+
i.def("save", &save_index<native_index_t>, py::arg("path"));
|
|
420
|
+
i.def("load", &load_index<native_index_t>, py::arg("path"));
|
|
421
|
+
i.def("view", &view_index<native_index_t>, py::arg("path"));
|
|
422
|
+
i.def("clear", &clear_index<native_index_t>);
|
|
423
|
+
|
|
424
|
+
auto si = py::class_<sets_index_t>(m, "SetsIndex");
|
|
425
|
+
|
|
426
|
+
si.def( //
|
|
427
|
+
py::init(&make_sets_index), //
|
|
428
|
+
py::kw_only(), //
|
|
429
|
+
py::arg("capacity") = 0, //
|
|
430
|
+
py::arg("connectivity") = config_t::connectivity_default_k, //
|
|
431
|
+
py::arg("expansion_add") = config_t::expansion_add_default_k, //
|
|
432
|
+
py::arg("expansion_search") = config_t::expansion_search_default_k //
|
|
433
|
+
);
|
|
434
|
+
|
|
435
|
+
si.def( //
|
|
436
|
+
"add",
|
|
437
|
+
[](sets_index_t& index, label_t label, py::array_t<set_member_t> set, bool copy) {
|
|
438
|
+
validate_set(set);
|
|
439
|
+
if (index.size() + 1 >= index.capacity())
|
|
440
|
+
index.reserve(ceil2(index.size() + 1));
|
|
441
|
+
auto proxy = set.unchecked<1>();
|
|
442
|
+
auto view = set_view_t{proxy.data(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
443
|
+
index.add(label, view, 0, copy);
|
|
444
|
+
}, //
|
|
445
|
+
py::arg("label"), //
|
|
446
|
+
py::arg("set"), //
|
|
447
|
+
py::kw_only(), //
|
|
448
|
+
py::arg("copy") = true //
|
|
449
|
+
);
|
|
450
|
+
|
|
451
|
+
si.def( //
|
|
452
|
+
"search",
|
|
453
|
+
[](sets_index_t& index, py::array_t<set_member_t> set, std::size_t count) -> py::array_t<label_t> {
|
|
454
|
+
validate_set(set);
|
|
455
|
+
auto proxy = set.unchecked<1>();
|
|
456
|
+
auto view = set_view_t{proxy.data(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
457
|
+
auto labels_py = py::array_t<label_t>(py_shape_t{static_cast<ssize_t>(count)});
|
|
458
|
+
auto labels_proxy = labels_py.mutable_unchecked<1>();
|
|
459
|
+
auto found = index.search(view, count, &labels_proxy(0), nullptr, 0);
|
|
460
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(found)});
|
|
461
|
+
return labels_py;
|
|
462
|
+
},
|
|
463
|
+
py::arg("set"), //
|
|
464
|
+
py::arg("count") = 10 //
|
|
465
|
+
);
|
|
466
|
+
|
|
467
|
+
si.def("__len__", &sets_index_t::size);
|
|
468
|
+
si.def_property_readonly("size", &sets_index_t::size);
|
|
469
|
+
si.def_property_readonly("connectivity", &sets_index_t::connectivity);
|
|
470
|
+
si.def_property_readonly("capacity", &sets_index_t::capacity);
|
|
471
|
+
|
|
472
|
+
si.def("save", &save_index<sets_index_t>, py::arg("path"));
|
|
473
|
+
si.def("load", &load_index<sets_index_t>, py::arg("path"));
|
|
474
|
+
si.def("view", &view_index<sets_index_t>, py::arg("path"));
|
|
475
|
+
si.def("clear", &clear_index<sets_index_t>);
|
|
476
|
+
|
|
477
|
+
auto hi = py::class_<hash_index_w_meta_t>(m, "HashIndex");
|
|
478
|
+
|
|
479
|
+
hi.def( //
|
|
480
|
+
py::init(&make_hash_index), //
|
|
481
|
+
py::kw_only(), //
|
|
482
|
+
py::arg("bits"), //
|
|
483
|
+
py::arg("capacity") = 0, //
|
|
484
|
+
py::arg("connectivity") = config_t::connectivity_default_k, //
|
|
485
|
+
py::arg("expansion_add") = config_t::expansion_add_default_k, //
|
|
486
|
+
py::arg("expansion_search") = config_t::expansion_search_default_k //
|
|
487
|
+
);
|
|
488
|
+
|
|
489
|
+
hi.def( //
|
|
490
|
+
"add",
|
|
491
|
+
[](hash_index_w_meta_t& index, label_t label, py::buffer array) {
|
|
492
|
+
if (index.size() + 1 >= index.capacity())
|
|
493
|
+
index.reserve(ceil2(index.size() + 1));
|
|
494
|
+
hash_buffer(index, array);
|
|
495
|
+
index.add(label, index.buffer(), 0);
|
|
496
|
+
}, //
|
|
497
|
+
py::arg("label"), //
|
|
498
|
+
py::arg("array") //
|
|
499
|
+
);
|
|
500
|
+
|
|
501
|
+
hi.def( //
|
|
502
|
+
"search",
|
|
503
|
+
[](hash_index_w_meta_t& index, py::buffer array, std::size_t count) -> py::array_t<label_t> {
|
|
504
|
+
if (index.size() + 1 >= index.capacity())
|
|
505
|
+
index.reserve(ceil2(index.size() + 1));
|
|
506
|
+
hash_buffer(index, array);
|
|
507
|
+
auto labels_py = py::array_t<label_t>(py_shape_t{static_cast<ssize_t>(count)});
|
|
508
|
+
auto labels_proxy = labels_py.mutable_unchecked<1>();
|
|
509
|
+
auto found = index.search(index.buffer(), count, &labels_proxy(0), nullptr, 0);
|
|
510
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(found)});
|
|
511
|
+
return labels_py;
|
|
512
|
+
},
|
|
513
|
+
py::arg("array"), //
|
|
514
|
+
py::arg("count") = 10 //
|
|
515
|
+
);
|
|
516
|
+
|
|
517
|
+
hi.def("__len__", &hash_index_w_meta_t::size);
|
|
518
|
+
hi.def_property_readonly("size", &hash_index_w_meta_t::size);
|
|
519
|
+
hi.def_property_readonly("connectivity", &hash_index_w_meta_t::connectivity);
|
|
520
|
+
hi.def_property_readonly("capacity", &hash_index_w_meta_t::capacity);
|
|
521
|
+
|
|
522
|
+
hi.def("save", &save_index<hash_index_w_meta_t>, py::arg("path"));
|
|
523
|
+
hi.def("load", &load_index<hash_index_w_meta_t>, py::arg("path"));
|
|
524
|
+
hi.def("view", &view_index<hash_index_w_meta_t>, py::arg("path"));
|
|
525
|
+
hi.def("clear", &clear_index<hash_index_w_meta_t>);
|
|
216
526
|
}
|
package/python/test.py
CHANGED
|
@@ -1,7 +1,6 @@
|
|
|
1
1
|
import numpy as np
|
|
2
2
|
|
|
3
|
-
import
|
|
4
|
-
|
|
3
|
+
from usearch.index import Index, SetsIndex
|
|
5
4
|
from numba import cfunc, types, carray
|
|
6
5
|
|
|
7
6
|
# Showcases how to use Numba to JIT-compile similarity measures for USearch.
|
|
@@ -29,13 +28,13 @@ vectors = np.random.uniform(
|
|
|
29
28
|
0, 0.3, (count_vectors, count_dimensions)).astype(np.float32)
|
|
30
29
|
labels = np.array(range(count_vectors), dtype=np.longlong)
|
|
31
30
|
|
|
32
|
-
index_udf =
|
|
31
|
+
index_udf = Index(ndim=96, metric_pointer=python_dot.address)
|
|
33
32
|
index_udf.add(labels, vectors, copy=True)
|
|
34
33
|
results = index_udf.search(vectors, 10)
|
|
35
34
|
print('found', results[0].shape, results[1].shape, results[2].shape)
|
|
36
35
|
|
|
37
36
|
|
|
38
|
-
index =
|
|
37
|
+
index = Index(ndim=96)
|
|
39
38
|
index.add(labels, vectors, copy=True)
|
|
40
39
|
assert index.ndim == count_dimensions
|
|
41
40
|
assert index.capacity >= count_vectors
|
|
@@ -43,3 +42,10 @@ assert len(index) == count_vectors
|
|
|
43
42
|
|
|
44
43
|
results = index.search(vectors, 10)
|
|
45
44
|
print('found', results[0].shape, results[1].shape, results[2].shape)
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
index_sets = SetsIndex()
|
|
48
|
+
index_sets.add(10, np.array([10, 12, 15], dtype=np.uint32))
|
|
49
|
+
index_sets.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
|
|
50
|
+
results = index_sets.search(np.array([12, 15], dtype=np.uint32), 20)
|
|
51
|
+
print(results)
|
|
File without changes
|
package/rust/lib.cpp
CHANGED
|
@@ -24,9 +24,20 @@ Matches Index::search_in_thread(rust::Slice<float const> vector, size_t count, s
|
|
|
24
24
|
return results;
|
|
25
25
|
}
|
|
26
26
|
|
|
27
|
-
void Index::add(uint32_t label, rust::Slice<float const> vector) const {
|
|
27
|
+
void Index::add(uint32_t label, rust::Slice<float const> vector) const {
|
|
28
|
+
index_->add(label, vector.data()); //
|
|
29
|
+
}
|
|
30
|
+
|
|
28
31
|
Matches Index::search(rust::Slice<float const> vector, size_t count) const {
|
|
29
|
-
|
|
32
|
+
Matches results;
|
|
33
|
+
results.labels.reserve(count);
|
|
34
|
+
results.distances.reserve(count);
|
|
35
|
+
for (size_t i = 0; i != count; ++i)
|
|
36
|
+
results.labels.push_back(0), results.distances.push_back(0);
|
|
37
|
+
results.count = index_->search(vector.data(), count, results.labels.data(), results.distances.data());
|
|
38
|
+
results.labels.truncate(results.count);
|
|
39
|
+
results.distances.truncate(results.count);
|
|
40
|
+
return results;
|
|
30
41
|
}
|
|
31
42
|
|
|
32
43
|
void Index::reserve(size_t capacity) const { index_->reserve(capacity); }
|
package/setup.py
CHANGED
|
@@ -9,11 +9,15 @@ compile_args = ['-std=c++11', '-O3']
|
|
|
9
9
|
if sys.platform == 'darwin':
|
|
10
10
|
compile_args.append('-mmacosx-version-min=10.13')
|
|
11
11
|
|
|
12
|
+
if sys.platform == 'win32':
|
|
13
|
+
compile_args = ['/std:c++14', '/O2']
|
|
14
|
+
|
|
15
|
+
|
|
12
16
|
ext_modules = [
|
|
13
17
|
Pybind11Extension(
|
|
14
|
-
'usearch',
|
|
18
|
+
'usearch.index',
|
|
15
19
|
['python/lib.cpp'],
|
|
16
|
-
extra_compile_args=compile_args
|
|
20
|
+
extra_compile_args=compile_args
|
|
17
21
|
),
|
|
18
22
|
]
|
|
19
23
|
|
|
@@ -30,7 +34,8 @@ setup(
|
|
|
30
34
|
|
|
31
35
|
name=__lib_name__,
|
|
32
36
|
version=__version__,
|
|
33
|
-
packages=
|
|
37
|
+
packages=['usearch'],
|
|
38
|
+
package_dir={'usearch': 'python/usearch'},
|
|
34
39
|
|
|
35
40
|
description='Smaller & Faster Single-File Vector Search Engine from Unum',
|
|
36
41
|
long_description=long_description,
|
|
@@ -255,6 +255,21 @@ inline static simsimd_f32_t simsimd_dot_f32x4avx2(simsimd_f32_t const* a, simsim
|
|
|
255
255
|
#endif
|
|
256
256
|
}
|
|
257
257
|
|
|
258
|
+
inline static simsimd_f32_t simsimd_dot_f16x16avx512(simsimd_f16_t const* a, simsimd_f16_t const* b, size_t d) {
|
|
259
|
+
#if defined(__AVX512F__)
|
|
260
|
+
__m512 ab_vec = _mm512_set1_ps(0);
|
|
261
|
+
for (size_t i = 0; i != d; i += 16)
|
|
262
|
+
ab_vec = _mm512_fmadd_ps( //
|
|
263
|
+
_mm512_cvtxph_ps(_mm256_loadu_ph(a + i)), //
|
|
264
|
+
_mm512_cvtxph_ps(_mm256_loadu_ph(b + i)), //
|
|
265
|
+
ab_vec);
|
|
266
|
+
return _mm512_reduce_add_ps(ab_vec);
|
|
267
|
+
#else
|
|
268
|
+
(void)a, (void)b, (void)d;
|
|
269
|
+
return 0;
|
|
270
|
+
#endif
|
|
271
|
+
}
|
|
272
|
+
|
|
258
273
|
inline static int32_t simsimd_dot_i8x16avx2(int8_t const* a, int8_t const* b, size_t d) {
|
|
259
274
|
#if defined(__AVX2__)
|
|
260
275
|
__m256i ab_vec = _mm256_set1_epi32(0);
|