usearch 2.8.7 → 2.8.8

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.
Files changed (102) hide show
  1. package/CITATION.cff +1 -1
  2. package/CMakeLists.txt +3 -3
  3. package/CONTRIBUTING.md +3 -3
  4. package/Cargo.toml +1 -1
  5. package/Dockerfile +16 -0
  6. package/README.md +41 -38
  7. package/VERSION +1 -0
  8. package/binding.gyp +57 -8
  9. package/conanfile.py +1 -1
  10. package/include/usearch/index.hpp +1 -1
  11. package/include/usearch/index_plugins.hpp +27 -23
  12. package/package.json +1 -2
  13. package/setup.py +76 -15
  14. package/wasmer.toml +1 -1
  15. package/. npmignore +0 -14
  16. package/.gitignore +0 -95
  17. package/.vscode/launch.json +0 -181
  18. package/.vscode/settings.json +0 -180
  19. package/.vscode/tasks.json +0 -71
  20. package/Package.swift +0 -41
  21. package/csharp/.editorconfig +0 -361
  22. package/csharp/Cloud.Unum.USearch.sln +0 -34
  23. package/csharp/DEV_SETUP.md +0 -51
  24. package/csharp/Directory.Build.props +0 -29
  25. package/csharp/README.md +0 -74
  26. package/csharp/build_and_test.cmd +0 -31
  27. package/csharp/build_and_test.sh +0 -30
  28. package/csharp/nuget/nuget-package.props +0 -33
  29. package/csharp/pack.sh +0 -56
  30. package/csharp/src/Cloud.Unum.USearch.Tests/Cloud.Unum.USearch.Tests.csproj +0 -39
  31. package/csharp/src/Cloud.Unum.USearch.Tests/USearchIndexTests.cs +0 -1013
  32. package/csharp/src/Cloud.Unum.USearch.Tests/Usings.cs +0 -1
  33. package/docs/_static/custom.css +0 -26
  34. package/docs/_static/custom.js +0 -7
  35. package/docs/benchmarks.md +0 -245
  36. package/docs/benchmarks.rst +0 -5
  37. package/docs/c/index.rst +0 -11
  38. package/docs/c/reference.rst +0 -7
  39. package/docs/conf.dox +0 -2668
  40. package/docs/conf.py +0 -53
  41. package/docs/contributing.rst +0 -5
  42. package/docs/cpp/index.rst +0 -11
  43. package/docs/cpp/reference.rst +0 -17
  44. package/docs/csharp/index.rst +0 -6
  45. package/docs/format.md +0 -86
  46. package/docs/format.rst +0 -5
  47. package/docs/golang/index.rst +0 -6
  48. package/docs/index.rst +0 -33
  49. package/docs/java/index.rst +0 -11
  50. package/docs/java/reference.rst +0 -7
  51. package/docs/javascript/index.rst +0 -11
  52. package/docs/javascript/reference.rst +0 -12
  53. package/docs/python/index.rst +0 -11
  54. package/docs/python/reference.rst +0 -42
  55. package/docs/rust/index.rst +0 -6
  56. package/docs/swift/index.rst +0 -6
  57. package/docs/wolfram/index.rst +0 -6
  58. package/golang/README.md +0 -95
  59. package/golang/go.mod +0 -3
  60. package/golang/lib.go +0 -379
  61. package/golang/lib_test.go +0 -108
  62. package/java/README.md +0 -22
  63. package/java/cloud/unum/usearch/Index.java +0 -337
  64. package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +0 -197
  65. package/java/cloud/unum/usearch/cloud_unum_usearch_Index.h +0 -125
  66. package/java/test/IndexTest.java +0 -30
  67. package/objc/USearchObjective.mm +0 -299
  68. package/objc/include/USearchObjective.h +0 -154
  69. package/python/README.md +0 -292
  70. package/python/lib.cpp +0 -1216
  71. package/python/scripts/bench.ipynb +0 -193
  72. package/python/scripts/bench.py +0 -137
  73. package/python/scripts/bench_exact.py +0 -44
  74. package/python/scripts/cluster.ipynb +0 -160
  75. package/python/scripts/index_faiss.py +0 -100
  76. package/python/scripts/join.py +0 -155
  77. package/python/scripts/test_index.py +0 -247
  78. package/python/scripts/test_jit.py +0 -231
  79. package/python/scripts/test_sparse.py +0 -32
  80. package/python/scripts/test_tooling.py +0 -133
  81. package/python/usearch/__init__.py +0 -0
  82. package/python/usearch/client.py +0 -124
  83. package/python/usearch/eval.py +0 -470
  84. package/python/usearch/index.py +0 -1401
  85. package/python/usearch/io.py +0 -113
  86. package/python/usearch/numba.py +0 -121
  87. package/python/usearch/server.py +0 -137
  88. package/rust/README.md +0 -66
  89. package/rust/lib.cpp +0 -136
  90. package/rust/lib.hpp +0 -69
  91. package/rust/lib.rs +0 -610
  92. package/swift/Index+Sugar.swift +0 -153
  93. package/swift/README.md +0 -48
  94. package/swift/Test.swift +0 -37
  95. package/swift/USearch.swift +0 -8
  96. package/wasm/CMakeLists.txt +0 -7
  97. package/wasm/README.md +0 -13
  98. package/wasm/lib.cpp +0 -5
  99. package/wolfram/CMakeLists.txt +0 -18
  100. package/wolfram/README.md +0 -13
  101. package/wolfram/lib.cpp +0 -160
  102. package/wolfram/test.wls +0 -27
@@ -1,155 +0,0 @@
1
- """
2
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clip_images.fbin -P datasets/cc_3M/
3
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clip_texts.fbin -P datasets/cc_3M/
4
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/images.fbin -P datasets/cc_3M/
5
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/texts.fbin -P datasets/cc_3M/
6
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_images.fbin -P datasets/cc_3M/
7
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/uform_english_texts.fbin -P datasets/cc_3M/
8
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clipbigg_images.fbin -P datasets/cc_3M/
9
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-cc-3m/resolve/main/clipbigg_texts.fbin -P datasets/cc_3M/
10
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-arxiv-2m/resolve/main/e5base_abstract.fbin -P datasets/arxiv_2M/
11
- wget -nc https://huggingface.co/datasets/unum-cloud/ann-arxiv-2m/resolve/main/e5base_title.fbin -P datasets/arxiv_2M/
12
- rm -rf datasets/cc_3M/*.usearch datasets/arxiv_2M/*.usearch
13
- python python/scripts/join.py
14
- """
15
- from numpy import dot
16
- from numpy.linalg import norm
17
- from tqdm import tqdm
18
- from simsimd import cos_f32x4_neon, to_int
19
-
20
- from usearch.index import Index, MetricKind, CompiledMetric, MetricSignature
21
- from usearch.io import load_matrix
22
- from usearch.eval import measure_seconds
23
-
24
- count = 10
25
- exact = False
26
- batch_size = 1024 * 4
27
- max_elements = 1000000
28
-
29
- a_name = "cc_3M/texts"
30
- b_name = "cc_3M/images"
31
-
32
- a_mat = load_matrix(f"datasets/{a_name}.fbin", view=True)
33
- b_mat = load_matrix(f"datasets/{b_name}.fbin", view=True)
34
-
35
- a_mat = a_mat[:max_elements]
36
- b_mat = b_mat[:max_elements]
37
-
38
- print(f"Loaded two datasets of shape: {a_mat.shape}, {b_mat.shape}")
39
- print("--------------------------------------")
40
- print("---------------Indexing---------------")
41
- print("--------------------------------------")
42
-
43
- metric = CompiledMetric(
44
- pointer=to_int(cos_f32x4_neon),
45
- kind=MetricKind.Cos,
46
- signature=MetricSignature.ArrayArraySize,
47
- )
48
-
49
- a = Index(
50
- a_mat.shape[1],
51
- metric=metric,
52
- path=f"datasets/{a_name}-{max_elements}.f32.usearch",
53
- dtype="f32",
54
- )
55
- b = Index(
56
- b_mat.shape[1],
57
- metric=metric,
58
- path=f"datasets/{b_name}-{max_elements}.f32.usearch",
59
- dtype="f32",
60
- )
61
-
62
- if len(a) != a_mat.shape[0]:
63
- a.clear()
64
- a.add(None, a_mat, log=True, batch_size=batch_size)
65
- a.save()
66
-
67
- if len(b) != b_mat.shape[0]:
68
- b.clear()
69
- b.add(None, b_mat, log=True, batch_size=batch_size)
70
- b.save()
71
-
72
-
73
- print(
74
- f"Loaded two indexes of size: {len(a):,} for {a_name} and {len(b):,} for {b_name}"
75
- )
76
- min_elements = min(len(a), len(b))
77
-
78
- run_diagnostics = input("Would you like to run diagnostics? [Y/n]: ")
79
- if len(run_diagnostics) == 0 or run_diagnostics.lower() == "y":
80
- print("--------------------------------------")
81
- print("-------------Diagnostics--------------")
82
- print("--------------------------------------")
83
-
84
- mean_similarity = 0.0
85
- mean_recovered_similarity = 0.0
86
-
87
- for i in tqdm(range(min_elements), desc="Pairwise Similarity"):
88
- a_vec = a_mat[i]
89
- b_vec = b_mat[i]
90
- cos_similarity = dot(a_vec, b_vec) / (norm(a_vec) * norm(b_vec))
91
- mean_similarity += cos_similarity
92
-
93
- a_vec = a[i]
94
- b_vec = b[i]
95
- cos_similarity = dot(a_vec, b_vec) / (norm(a_vec) * norm(b_vec))
96
- mean_recovered_similarity += cos_similarity
97
-
98
- mean_similarity /= min_elements
99
- mean_recovered_similarity /= min_elements
100
- print(
101
- f"Average vector similarity is {mean_similarity:.4f} in original dataset, "
102
- f"and {mean_recovered_similarity:.4f} in recovered state in index"
103
- )
104
-
105
- dt = measure_seconds
106
- args = dict(
107
- count=count,
108
- batch_size=batch_size,
109
- log=True,
110
- exact=exact,
111
- )
112
-
113
- secs, a_self_recall = dt(lambda: a.search(a.vectors, **args).recall(a.keys))
114
- print(
115
- "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
116
- count, a_name, a_self_recall * 100, secs
117
- )
118
- )
119
-
120
- secs, b_self_recall = dt(lambda: b.search(b.vectors, **args).recall(b.keys))
121
- print(
122
- "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
123
- count, b_name, b_self_recall * 100, secs
124
- )
125
- )
126
-
127
- secs, ab_recall = dt(lambda: b.search(a.vectors, **args).recall(b.keys))
128
- print(
129
- "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
130
- count, a_name, b_name, ab_recall * 100, secs
131
- )
132
- )
133
-
134
- secs, ba_recall = dt(lambda: a.search(b.vectors, **args).recall(a.keys))
135
- print(
136
- "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
137
- count, b_name, a_name, ba_recall * 100, secs
138
- )
139
- )
140
-
141
-
142
- print("--------------------------------------")
143
- print("-----------------Join-----------------")
144
- print("--------------------------------------")
145
-
146
- secs, bimapping = measure_seconds(lambda: a.join(b, max_proposals=100))
147
- mapping_size = len(bimapping)
148
- recall = 0
149
- for i, j in bimapping.items():
150
- recall += i == j
151
-
152
- recall *= 100.0 / min_elements
153
- print(
154
- f"Took {secs:.2f}s to find {mapping_size:,} pairings with {recall:.2f}% being exact"
155
- )
@@ -1,247 +0,0 @@
1
- import os
2
-
3
- import pytest
4
- import numpy as np
5
-
6
- from usearch.io import load_matrix, save_matrix
7
- from usearch.eval import random_vectors
8
- from usearch.index import search
9
-
10
- from usearch.index import (
11
- Index,
12
- MetricKind,
13
- ScalarKind,
14
- Match,
15
- Matches,
16
- BatchMatches,
17
- Clustering,
18
- )
19
- from usearch.index import (
20
- DEFAULT_CONNECTIVITY,
21
- )
22
-
23
-
24
- ndims = [3, 97, 256]
25
- batch_sizes = [1, 11, 77]
26
- quantizations = [
27
- ScalarKind.F32,
28
- ScalarKind.F64,
29
- ScalarKind.F16,
30
- ScalarKind.I8,
31
- ]
32
- dtypes = [np.float32, np.float64, np.float16]
33
- threads = 2
34
-
35
- connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
36
- continuous_metrics = [MetricKind.Cos, MetricKind.L2sq]
37
- hash_metrics = [
38
- MetricKind.Hamming,
39
- MetricKind.Tanimoto,
40
- MetricKind.Sorensen,
41
- ]
42
-
43
-
44
- @pytest.mark.parametrize("ndim", [3, 97, 256])
45
- @pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
46
- @pytest.mark.parametrize("batch_size", [1, 7, 1024])
47
- @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
48
- @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
49
- def test_index_initialization_and_addition(
50
- ndim, metric, quantization, dtype, batch_size
51
- ):
52
- index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
53
- keys = np.arange(batch_size)
54
- vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
55
- index.add(keys, vectors, threads=threads)
56
- assert len(index) == batch_size
57
-
58
-
59
- @pytest.mark.parametrize("ndim", [3, 97, 256])
60
- @pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
61
- @pytest.mark.parametrize("batch_size", [1, 7, 1024])
62
- @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
63
- @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
64
- def test_index_retrieval(ndim, metric, quantization, dtype, batch_size):
65
- index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
66
- keys = np.arange(batch_size)
67
- vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
68
- index.add(keys, vectors, threads=threads)
69
- vectors_retrived = np.vstack(index.get(keys, dtype))
70
- assert np.allclose(vectors_retrived, vectors, atol=0.1)
71
-
72
-
73
- @pytest.mark.parametrize("ndim", [3, 97, 256])
74
- @pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
75
- @pytest.mark.parametrize("batch_size", [1, 7, 1024])
76
- @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
77
- @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
78
- def test_index_search(ndim, metric, quantization, dtype, batch_size):
79
- index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
80
- keys = np.arange(batch_size)
81
- vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
82
- index.add(keys, vectors, threads=threads)
83
-
84
- if batch_size == 1:
85
- matches: Matches = index.search(vectors, 10, threads=threads)
86
- assert matches.keys.ndim == 1
87
- assert matches.keys.shape[0] == matches.distances.shape[0]
88
- assert len(matches) == batch_size
89
- assert np.all(np.sort(index.keys) == np.sort(keys))
90
-
91
- else:
92
- matches: BatchMatches = index.search(vectors, 10, threads=threads)
93
- assert matches.keys.ndim == 2
94
- assert matches.keys.shape[0] == matches.distances.shape[0]
95
- assert len(matches) == batch_size
96
- assert np.all(np.sort(index.keys) == np.sort(keys))
97
-
98
-
99
- @pytest.mark.parametrize("batch_size", [1, 7, 1024])
100
- def test_index_duplicates(batch_size):
101
- ndim = 8
102
- index = Index(ndim=ndim, multi=False)
103
- keys = np.arange(batch_size)
104
- vectors = random_vectors(count=batch_size, ndim=ndim)
105
- index.add(keys, vectors, threads=threads)
106
- with pytest.raises(Exception):
107
- index.add(keys, vectors, threads=threads)
108
-
109
- index = Index(ndim=ndim, multi=True)
110
- keys = np.arange(batch_size)
111
- vectors = random_vectors(count=batch_size, ndim=ndim)
112
- index.add(keys, vectors, threads=threads)
113
- index.add(keys, vectors, threads=threads)
114
- assert len(index) == batch_size * 2
115
-
116
- two_per_key = index.get(keys)
117
- assert np.vstack(two_per_key).shape == (2 * batch_size, ndim)
118
-
119
-
120
- @pytest.mark.parametrize("batch_size", [1, 7, 1024])
121
- def test_index_stats(batch_size):
122
- ndim = 8
123
- index = Index(ndim=ndim, multi=False)
124
- keys = np.arange(batch_size)
125
- vectors = random_vectors(count=batch_size, ndim=ndim)
126
- index.add(keys, vectors, threads=threads)
127
-
128
- assert index.max_level >= 0
129
- assert index.stats.nodes >= batch_size
130
- assert index.levels_stats[0].nodes == batch_size
131
- assert index.level_stats(0).nodes == batch_size
132
-
133
- assert index.levels_stats[index.max_level].nodes > 0
134
-
135
-
136
- @pytest.mark.parametrize("ndim", [1, 3, 8, 32, 256, 4096])
137
- @pytest.mark.parametrize("batch_size", [0, 1, 7, 1024])
138
- @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
139
- def test_index_save_load_restore_copy(ndim, quantization, batch_size):
140
- index = Index(ndim=ndim, dtype=quantization, multi=False)
141
-
142
- if batch_size > 0:
143
- keys = np.arange(batch_size)
144
- vectors = random_vectors(count=batch_size, ndim=ndim)
145
- index.add(keys, vectors, threads=threads)
146
-
147
- index.save("tmp.usearch")
148
- index.clear()
149
- assert len(index) == 0
150
- assert os.path.exists("tmp.usearch")
151
-
152
- index.load("tmp.usearch")
153
- assert len(index) == batch_size
154
- if batch_size > 0:
155
- assert len(index[0].flatten()) == ndim
156
-
157
- index_meta = Index.metadata("tmp.usearch")
158
- assert index_meta is not None
159
-
160
- index = Index.restore("tmp.usearch", view=False)
161
- assert len(index) == batch_size
162
- if batch_size > 0:
163
- assert len(index[0].flatten()) == ndim
164
-
165
- copied_index = index.copy()
166
- assert len(copied_index) == len(index)
167
- if batch_size > 0:
168
- assert np.allclose(
169
- np.vstack(copied_index.get(keys)), np.vstack(index.get(keys))
170
- )
171
-
172
- # Perform the same operations in RAM, without touching the filesystem
173
- serialized_index = index.save()
174
- deserialized_metadata = Index.metadata(serialized_index)
175
- assert deserialized_metadata is not None
176
-
177
- deserialized_index = Index.restore(serialized_index)
178
- assert len(deserialized_index) == len(index)
179
- assert set(np.array(deserialized_index.keys)) == set(np.array(index.keys))
180
- if batch_size > 0:
181
- assert np.allclose(
182
- np.vstack(deserialized_index.get(keys)), np.vstack(index.get(keys))
183
- )
184
-
185
- deserialized_index.reset()
186
- index.reset()
187
- os.remove("tmp.usearch")
188
-
189
-
190
- @pytest.mark.parametrize("batch_size", [32])
191
- def test_index_contains_remove_rename(batch_size):
192
- if batch_size <= 1:
193
- return
194
-
195
- ndim = 8
196
- index = Index(ndim=ndim, multi=False)
197
- keys = np.arange(batch_size)
198
- vectors = random_vectors(count=batch_size, ndim=ndim)
199
-
200
- index.add(keys, vectors, threads=threads)
201
- assert np.all(index.contains(keys))
202
- assert np.all(index.count(keys) == np.ones(batch_size))
203
-
204
- removed_keys = keys[: batch_size // 2]
205
- remaining_keys = keys[batch_size // 2 :]
206
- index.remove(removed_keys)
207
- assert len(index) == (len(keys) - len(removed_keys))
208
- assert np.sum(index.contains(keys)) == len(remaining_keys)
209
- assert np.sum(index.count(keys)) == len(remaining_keys)
210
- assert np.sum(index.count(removed_keys)) == 0
211
-
212
- assert keys[0] not in index
213
- assert keys[-1] in index
214
-
215
- renamed_counts = index.rename(removed_keys, removed_keys)
216
- assert np.sum(index.count(renamed_counts)) == 0
217
-
218
- renamed_counts = index.rename(remaining_keys, removed_keys)
219
- assert np.sum(index.count(removed_keys)) == len(index)
220
-
221
-
222
- @pytest.mark.parametrize("ndim", [3, 97, 256])
223
- @pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
224
- @pytest.mark.parametrize("batch_size", [500, 1024])
225
- @pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
226
- @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
227
- def test_index_clustering(ndim, metric, quantization, dtype, batch_size):
228
- index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
229
- keys = np.arange(batch_size)
230
- vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
231
- index.add(keys, vectors, threads=threads)
232
-
233
- clusters: Clustering = index.cluster(vectors=vectors, threads=threads)
234
- assert len(clusters.matches.keys) == batch_size
235
-
236
- # If no argument is provided, we cluster the present entries
237
- clusters: Clustering = index.cluster(threads=threads)
238
- assert len(clusters.matches.keys) == batch_size
239
-
240
- # If no argument is provided, we cluster the present entries
241
- clusters: Clustering = index.cluster(keys=keys[:50], threads=threads)
242
- assert len(clusters.matches.keys) == 50
243
-
244
- # If no argument is provided, we cluster the present entries
245
- clusters: Clustering = index.cluster(min_count=3, max_count=10, threads=threads)
246
- unique_clusters = set(clusters.matches.keys.flatten().tolist())
247
- assert len(unique_clusters) >= 3 and len(unique_clusters) <= 10
@@ -1,231 +0,0 @@
1
- import pytest
2
- import numpy as np
3
-
4
- from usearch.eval import random_vectors
5
-
6
- from usearch.index import (
7
- Index,
8
- MetricKind,
9
- MetricSignature,
10
- CompiledMetric,
11
- )
12
-
13
- dimensions = [3, 97, 256]
14
- batch_sizes = [1, 77]
15
-
16
-
17
- @pytest.mark.parametrize("ndim", dimensions)
18
- @pytest.mark.parametrize("batch_size", batch_sizes)
19
- def test_index_numba(ndim: int, batch_size: int):
20
- try:
21
- from numba import cfunc, types, carray
22
- except ImportError:
23
- return
24
-
25
- # Showcases how to use Numba to JIT-compile similarity measures for USearch.
26
- # https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
27
- signature_two_args = types.float32(
28
- types.CPointer(types.float32),
29
- types.CPointer(types.float32),
30
- )
31
- signature_three_args = types.float32(
32
- types.CPointer(types.float32),
33
- types.CPointer(types.float32),
34
- types.uint64,
35
- )
36
- signature_four_args = types.float32(
37
- types.CPointer(types.float32),
38
- types.uint64,
39
- types.CPointer(types.float32),
40
- types.uint64,
41
- )
42
-
43
- @cfunc(signature_two_args)
44
- def python_inner_product_two_args(a, b):
45
- a_array = carray(a, ndim)
46
- b_array = carray(b, ndim)
47
- c = 0.0
48
- for i in range(ndim):
49
- c += a_array[i] * b_array[i]
50
- return 1 - c
51
-
52
- @cfunc(signature_three_args)
53
- def python_inner_product_three_args(a, b, ndim):
54
- a_array = carray(a, ndim)
55
- b_array = carray(b, ndim)
56
- c = 0.0
57
- for i in range(ndim):
58
- c += a_array[i] * b_array[i]
59
- return 1 - c
60
-
61
- functions = [
62
- python_inner_product_two_args,
63
- python_inner_product_three_args,
64
- ]
65
- signatures = [
66
- MetricSignature.ArrayArray,
67
- MetricSignature.ArrayArraySize,
68
- ]
69
- for function, signature in zip(functions, signatures):
70
- metric = CompiledMetric(
71
- pointer=function.address,
72
- kind=MetricKind.IP,
73
- signature=signature,
74
- )
75
- index = Index(ndim=ndim, metric=metric)
76
-
77
- keys = np.arange(batch_size)
78
- vectors = random_vectors(count=batch_size, ndim=ndim)
79
-
80
- index.add(keys, vectors)
81
- matches = index.search(vectors, 10)
82
- assert len(matches) == batch_size
83
-
84
-
85
- @pytest.mark.parametrize("ndim", dimensions[-1:])
86
- @pytest.mark.parametrize("batch_size", batch_sizes[-1:])
87
- def test_index_cppyy(ndim: int, batch_size: int):
88
- try:
89
- import cppyy
90
- import cppyy.ll
91
- except ImportError:
92
- return
93
-
94
- cppyy.cppdef(
95
- """
96
- float inner_product_two_args(float *a, float *b) {
97
- float result = 0;
98
- #pragma unroll
99
- for (size_t i = 0; i != ndim; ++i)
100
- result += a[i] * b[i];
101
- return 1 - result;
102
- }
103
-
104
- float inner_product_three_args(float *a, float *b, size_t n) {
105
- float result = 0;
106
- for (size_t i = 0; i != n; ++i)
107
- result += a[i] * b[i];
108
- return 1 - result;
109
- }
110
-
111
- float inner_product_four_args(float *a, size_t an, float *b, size_t) {
112
- float result = 0;
113
- for (size_t i = 0; i != an; ++i)
114
- result += a[i] * b[i];
115
- return 1 - result;
116
- }
117
- """.replace(
118
- "ndim", str(ndim)
119
- )
120
- )
121
-
122
- functions = [
123
- cppyy.gbl.inner_product_two_args,
124
- cppyy.gbl.inner_product_three_args,
125
- # cppyy.gbl.inner_product_four_args,
126
- ]
127
- signatures = [
128
- MetricSignature.ArrayArray,
129
- MetricSignature.ArrayArraySize,
130
- # MetricSignature.ArraySizeArraySize,
131
- ]
132
- for function, signature in zip(functions, signatures):
133
- metric = CompiledMetric(
134
- pointer=cppyy.ll.addressof(function),
135
- kind=MetricKind.IP,
136
- signature=signature,
137
- )
138
- index = Index(ndim=ndim, metric=metric)
139
-
140
- keys = np.arange(batch_size)
141
- vectors = random_vectors(count=batch_size, ndim=ndim)
142
-
143
- index.add(keys, vectors)
144
- matches = index.search(vectors, 10)
145
- assert len(matches) == batch_size
146
-
147
-
148
- @pytest.mark.parametrize("ndim", [8])
149
- @pytest.mark.parametrize("batch_size", batch_sizes)
150
- def test_index_peachpy(ndim: int, batch_size: int):
151
- try:
152
- from peachpy import (
153
- Argument,
154
- ptr,
155
- float_,
156
- const_float_,
157
- )
158
- from peachpy.x86_64 import (
159
- abi,
160
- Function,
161
- uarch,
162
- isa,
163
- GeneralPurposeRegister64,
164
- LOAD,
165
- YMMRegister,
166
- VSUBPS,
167
- VADDPS,
168
- VHADDPS,
169
- VMOVUPS,
170
- VFMADD231PS,
171
- VPERM2F128,
172
- VXORPS,
173
- RETURN,
174
- )
175
- except ImportError:
176
- return
177
-
178
- a = Argument(ptr(const_float_), name="a")
179
- b = Argument(ptr(const_float_), name="b")
180
-
181
- with Function(
182
- "InnerProduct", (a, b), float_, target=uarch.default + isa.avx + isa.avx2
183
- ) as asm_function:
184
- # Request two 64-bit general-purpose registers for addresses
185
- reg_a, reg_b = GeneralPurposeRegister64(), GeneralPurposeRegister64()
186
- LOAD.ARGUMENT(reg_a, a)
187
- LOAD.ARGUMENT(reg_b, b)
188
-
189
- # Load the vectors
190
- ymm_a = YMMRegister()
191
- ymm_b = YMMRegister()
192
- VMOVUPS(ymm_a, [reg_a])
193
- VMOVUPS(ymm_b, [reg_b])
194
-
195
- # Prepare the accumulator
196
- ymm_c = YMMRegister()
197
- ymm_one = YMMRegister()
198
- VXORPS(ymm_c, ymm_c, ymm_c)
199
- VXORPS(ymm_one, ymm_one, ymm_one)
200
-
201
- # Accumulate A and B product into C
202
- VFMADD231PS(ymm_c, ymm_a, ymm_b)
203
-
204
- # Reduce the contents of a YMM register
205
- ymm_c_permuted = YMMRegister()
206
- VPERM2F128(ymm_c_permuted, ymm_c, ymm_c, 1)
207
- VADDPS(ymm_c, ymm_c, ymm_c_permuted)
208
- VHADDPS(ymm_c, ymm_c, ymm_c)
209
- VHADDPS(ymm_c, ymm_c, ymm_c)
210
-
211
- # Negate the values, to go from "similarity" to "distance"
212
- VSUBPS(ymm_c, ymm_one, ymm_c)
213
-
214
- # A common convention is to return floats in XMM registers
215
- RETURN(ymm_c.as_xmm)
216
-
217
- python_function = asm_function.finalize(abi.detect()).encode().load()
218
- metric = CompiledMetric(
219
- pointer=python_function.loader.code_address,
220
- kind=MetricKind.IP,
221
- signature=MetricSignature.ArrayArray,
222
- )
223
- index = Index(ndim=ndim, metric=metric)
224
-
225
- keys = np.arange(batch_size)
226
- vectors = random_vectors(count=batch_size, ndim=ndim)
227
-
228
- index.add(keys, vectors)
229
- matches, distances, count = index.search(vectors, 10)
230
- assert matches.shape[0] == distances.shape[0]
231
- assert count.shape[0] == batch_size
@@ -1,32 +0,0 @@
1
- import pytest
2
- import numpy as np
3
-
4
- from usearch.index import (
5
- Index,
6
- MetricKind,
7
- ScalarKind,
8
- )
9
-
10
-
11
- @pytest.mark.parametrize("bits", [7, 97, 256, 4097])
12
- @pytest.mark.parametrize("metric", [MetricKind.Tanimoto])
13
- @pytest.mark.parametrize("connectivity", [3, 13, 50])
14
- @pytest.mark.parametrize("batch_size", [3, 77])
15
- def test_bitwise_index(
16
- bits: int,
17
- metric: MetricKind,
18
- connectivity: int,
19
- batch_size: int,
20
- ):
21
- index = Index(ndim=bits, metric=metric, connectivity=connectivity)
22
-
23
- keys = np.arange(batch_size)
24
- byte_vectors = np.random.randint(2, size=(batch_size, bits))
25
- bit_vectors = np.packbits(byte_vectors, axis=1)
26
-
27
- index.add(keys, bit_vectors)
28
-
29
- byte_vectors_retrieved = np.vstack(index.get(keys, ScalarKind.B1))
30
- assert np.all(byte_vectors_retrieved == bit_vectors)
31
-
32
- index.search(bit_vectors, 10)