vicinity 0.3.1__tar.gz → 0.3.2__tar.gz
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.
- {vicinity-0.3.1 → vicinity-0.3.2}/Makefile +1 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/PKG-INFO +27 -11
- {vicinity-0.3.1 → vicinity-0.3.2}/README.md +23 -10
- vicinity-0.3.2/assets/images/vicinity_logo.png +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/pyproject.toml +4 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/tests/conftest.py +51 -3
- {vicinity-0.3.1 → vicinity-0.3.2}/tests/test_vicinity.py +41 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/__init__.py +7 -2
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/annoy.py +13 -11
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/base.py +9 -3
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/basic.py +81 -60
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/faiss.py +2 -3
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/hnsw.py +5 -4
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/pynndescent.py +3 -6
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/backends/usearch.py +5 -5
- vicinity-0.3.2/vicinity/backends/voyager.py +116 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/datatypes.py +1 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/version.py +1 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/vicinity.py +33 -5
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity.egg-info/PKG-INFO +27 -11
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity.egg-info/SOURCES.txt +3 -1
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity.egg-info/requires.txt +4 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/.github/workflows/ci.yaml +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/.gitignore +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/.pre-commit-config.yaml +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/LICENSE +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/setup.cfg +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/tests/test_utils.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/uv.lock +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/__init__.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/py.typed +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity/utils.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity.egg-info/dependency_links.txt +0 -0
- {vicinity-0.3.1 → vicinity-0.3.2}/vicinity.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: vicinity
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: Lightweight Nearest Neighbors with Flexible Backends
|
|
5
5
|
Author-email: Stéphan Tulkens <stephantul@gmail.com>, Thomas van Dongen <thomas123@live.nl>
|
|
6
6
|
License: MIT License
|
|
@@ -67,6 +67,8 @@ Provides-Extra: faiss
|
|
|
67
67
|
Requires-Dist: faiss-cpu; extra == "faiss"
|
|
68
68
|
Provides-Extra: usearch
|
|
69
69
|
Requires-Dist: usearch; extra == "usearch"
|
|
70
|
+
Provides-Extra: voyager
|
|
71
|
+
Requires-Dist: voyager; extra == "voyager"
|
|
70
72
|
Provides-Extra: all
|
|
71
73
|
Requires-Dist: hnswlib; extra == "all"
|
|
72
74
|
Requires-Dist: pynndescent>=0.5.10; extra == "all"
|
|
@@ -76,13 +78,19 @@ Requires-Dist: numpy>=1.24.0; extra == "all"
|
|
|
76
78
|
Requires-Dist: annoy; extra == "all"
|
|
77
79
|
Requires-Dist: faiss-cpu; extra == "all"
|
|
78
80
|
Requires-Dist: usearch; extra == "all"
|
|
81
|
+
Requires-Dist: voyager; extra == "all"
|
|
79
82
|
|
|
80
|
-
<div align="center">
|
|
81
|
-
|
|
82
|
-
# Vicinity: Lightweight Nearest Neighbors
|
|
83
83
|
|
|
84
|
+
<div align="center">
|
|
85
|
+
<picture>
|
|
86
|
+
<img width="40%" alt="Model2Vec logo" src="assets/images/vicinity_logo.png">
|
|
87
|
+
</picture>
|
|
88
|
+
</a>
|
|
84
89
|
</div>
|
|
85
90
|
|
|
91
|
+
<div align="center">
|
|
92
|
+
<h2>Lightweight Nearest Neighbors with Flexible Backends</h2>
|
|
93
|
+
</div>
|
|
86
94
|
|
|
87
95
|
<div align="center">
|
|
88
96
|
<h2>
|
|
@@ -118,7 +126,7 @@ Install the package with:
|
|
|
118
126
|
```bash
|
|
119
127
|
pip install vicinity
|
|
120
128
|
```
|
|
121
|
-
|
|
129
|
+
Optionally, [install any of the supported backends](#installation), or simply install all of them with:
|
|
122
130
|
```bash
|
|
123
131
|
pip install vicinity[all]
|
|
124
132
|
```
|
|
@@ -134,7 +142,12 @@ items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
|
|
|
134
142
|
vectors = np.random.rand(len(items), 128)
|
|
135
143
|
|
|
136
144
|
# Initialize the Vicinity instance (using the basic backend and cosine metric)
|
|
137
|
-
vicinity = Vicinity.from_vectors_and_items(
|
|
145
|
+
vicinity = Vicinity.from_vectors_and_items(
|
|
146
|
+
vectors=vectors,
|
|
147
|
+
items=items,
|
|
148
|
+
backend_type=Backend.BASIC,
|
|
149
|
+
metric=Metric.COSINE
|
|
150
|
+
)
|
|
138
151
|
|
|
139
152
|
# Create a query vector
|
|
140
153
|
query_vector = np.random.rand(128)
|
|
@@ -189,7 +202,7 @@ The following backends are supported:
|
|
|
189
202
|
- `ivf_scalar`: Inverted file search with scalar quantizer.
|
|
190
203
|
- `ivfpq`: Inverted file search with product quantizer.
|
|
191
204
|
- `ivfpqr`: Inverted file search with product quantizer and refinement.
|
|
192
|
-
|
|
205
|
+
- [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors.
|
|
193
206
|
|
|
194
207
|
|
|
195
208
|
|
|
@@ -200,7 +213,7 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
200
213
|
|
|
201
214
|
| Backend | Parameter | Description | Default Value |
|
|
202
215
|
|-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
|
|
203
|
-
| **
|
|
216
|
+
| **ANNOY** | `metric` | Similarity metric to use (`dot`, `euclidean`, `cosine`). | `"cosine"` |
|
|
204
217
|
| | `trees` | Number of trees to use for indexing. | `100` |
|
|
205
218
|
| | `length` | Optional length of the dataset. | `None` |
|
|
206
219
|
| **FAISS** | `metric` | Similarity metric to use (`cosine`, `l2`). | `"cosine"` |
|
|
@@ -212,13 +225,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
212
225
|
| **HNSW** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
213
226
|
| | `ef_construction` | Size of the dynamic list during index construction. | `200` |
|
|
214
227
|
| | `m` | Number of connections per layer. | `16` |
|
|
215
|
-
| **
|
|
228
|
+
| **PYNNDESCENT** | `metric` | Similarity metric to use (`cosine`, `euclidean`, `manhattan`). | `"cosine"` |
|
|
216
229
|
| | `n_neighbors` | Number of neighbors to use for search. | `15` |
|
|
217
|
-
| **
|
|
230
|
+
| **USEARCH** | `metric` | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`). | `"cos"` |
|
|
218
231
|
| | `connectivity` | Number of connections per node in the graph. | `16` |
|
|
219
232
|
| | `expansion_add` | Number of candidates considered during graph construction. | `128` |
|
|
220
233
|
| | `expansion_search` | Number of candidates considered during search. | `64` |
|
|
221
|
-
|
|
234
|
+
| **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
235
|
+
| | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` |
|
|
236
|
+
| | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` |
|
|
222
237
|
|
|
223
238
|
## Installation
|
|
224
239
|
The following installation options are available:
|
|
@@ -235,6 +250,7 @@ pip install vicinity[faiss]
|
|
|
235
250
|
pip install vicinity[hnsw]
|
|
236
251
|
pip install vicinity[pynndescent]
|
|
237
252
|
pip install vicinity[usearch]
|
|
253
|
+
pip install vicinity[voyager]
|
|
238
254
|
```
|
|
239
255
|
|
|
240
256
|
## License
|
|
@@ -1,9 +1,14 @@
|
|
|
1
|
-
<div align="center">
|
|
2
|
-
|
|
3
|
-
# Vicinity: Lightweight Nearest Neighbors
|
|
4
1
|
|
|
2
|
+
<div align="center">
|
|
3
|
+
<picture>
|
|
4
|
+
<img width="40%" alt="Model2Vec logo" src="assets/images/vicinity_logo.png">
|
|
5
|
+
</picture>
|
|
6
|
+
</a>
|
|
5
7
|
</div>
|
|
6
8
|
|
|
9
|
+
<div align="center">
|
|
10
|
+
<h2>Lightweight Nearest Neighbors with Flexible Backends</h2>
|
|
11
|
+
</div>
|
|
7
12
|
|
|
8
13
|
<div align="center">
|
|
9
14
|
<h2>
|
|
@@ -39,7 +44,7 @@ Install the package with:
|
|
|
39
44
|
```bash
|
|
40
45
|
pip install vicinity
|
|
41
46
|
```
|
|
42
|
-
|
|
47
|
+
Optionally, [install any of the supported backends](#installation), or simply install all of them with:
|
|
43
48
|
```bash
|
|
44
49
|
pip install vicinity[all]
|
|
45
50
|
```
|
|
@@ -55,7 +60,12 @@ items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
|
|
|
55
60
|
vectors = np.random.rand(len(items), 128)
|
|
56
61
|
|
|
57
62
|
# Initialize the Vicinity instance (using the basic backend and cosine metric)
|
|
58
|
-
vicinity = Vicinity.from_vectors_and_items(
|
|
63
|
+
vicinity = Vicinity.from_vectors_and_items(
|
|
64
|
+
vectors=vectors,
|
|
65
|
+
items=items,
|
|
66
|
+
backend_type=Backend.BASIC,
|
|
67
|
+
metric=Metric.COSINE
|
|
68
|
+
)
|
|
59
69
|
|
|
60
70
|
# Create a query vector
|
|
61
71
|
query_vector = np.random.rand(128)
|
|
@@ -110,7 +120,7 @@ The following backends are supported:
|
|
|
110
120
|
- `ivf_scalar`: Inverted file search with scalar quantizer.
|
|
111
121
|
- `ivfpq`: Inverted file search with product quantizer.
|
|
112
122
|
- `ivfpqr`: Inverted file search with product quantizer and refinement.
|
|
113
|
-
|
|
123
|
+
- [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors.
|
|
114
124
|
|
|
115
125
|
|
|
116
126
|
|
|
@@ -121,7 +131,7 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
121
131
|
|
|
122
132
|
| Backend | Parameter | Description | Default Value |
|
|
123
133
|
|-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
|
|
124
|
-
| **
|
|
134
|
+
| **ANNOY** | `metric` | Similarity metric to use (`dot`, `euclidean`, `cosine`). | `"cosine"` |
|
|
125
135
|
| | `trees` | Number of trees to use for indexing. | `100` |
|
|
126
136
|
| | `length` | Optional length of the dataset. | `None` |
|
|
127
137
|
| **FAISS** | `metric` | Similarity metric to use (`cosine`, `l2`). | `"cosine"` |
|
|
@@ -133,13 +143,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
133
143
|
| **HNSW** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
134
144
|
| | `ef_construction` | Size of the dynamic list during index construction. | `200` |
|
|
135
145
|
| | `m` | Number of connections per layer. | `16` |
|
|
136
|
-
| **
|
|
146
|
+
| **PYNNDESCENT** | `metric` | Similarity metric to use (`cosine`, `euclidean`, `manhattan`). | `"cosine"` |
|
|
137
147
|
| | `n_neighbors` | Number of neighbors to use for search. | `15` |
|
|
138
|
-
| **
|
|
148
|
+
| **USEARCH** | `metric` | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`). | `"cos"` |
|
|
139
149
|
| | `connectivity` | Number of connections per node in the graph. | `16` |
|
|
140
150
|
| | `expansion_add` | Number of candidates considered during graph construction. | `128` |
|
|
141
151
|
| | `expansion_search` | Number of candidates considered during search. | `64` |
|
|
142
|
-
|
|
152
|
+
| **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
153
|
+
| | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` |
|
|
154
|
+
| | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` |
|
|
143
155
|
|
|
144
156
|
## Installation
|
|
145
157
|
The following installation options are available:
|
|
@@ -156,6 +168,7 @@ pip install vicinity[faiss]
|
|
|
156
168
|
pip install vicinity[hnsw]
|
|
157
169
|
pip install vicinity[pynndescent]
|
|
158
170
|
pip install vicinity[usearch]
|
|
171
|
+
pip install vicinity[voyager]
|
|
159
172
|
```
|
|
160
173
|
|
|
161
174
|
## License
|
|
Binary file
|
|
@@ -52,6 +52,7 @@ pynndescent = [
|
|
|
52
52
|
annoy = ["annoy"]
|
|
53
53
|
faiss = ["faiss-cpu"]
|
|
54
54
|
usearch = ["usearch"]
|
|
55
|
+
voyager = ["voyager"]
|
|
55
56
|
all = [
|
|
56
57
|
"hnswlib",
|
|
57
58
|
"pynndescent>=0.5.10",
|
|
@@ -60,7 +61,8 @@ all = [
|
|
|
60
61
|
"numpy>=1.24.0",
|
|
61
62
|
"annoy",
|
|
62
63
|
"faiss-cpu",
|
|
63
|
-
"usearch"
|
|
64
|
+
"usearch",
|
|
65
|
+
"voyager"
|
|
64
66
|
]
|
|
65
67
|
|
|
66
68
|
[project.urls]
|
|
@@ -87,6 +89,7 @@ select = [
|
|
|
87
89
|
"NPY",
|
|
88
90
|
# Print: Forbid print statements
|
|
89
91
|
"T20",
|
|
92
|
+
"F",
|
|
90
93
|
]
|
|
91
94
|
ignore = [
|
|
92
95
|
# Allow self and cls to be untyped, and allow Any type
|
|
@@ -8,7 +8,17 @@ from vicinity.datatypes import Backend
|
|
|
8
8
|
|
|
9
9
|
random_gen = np.random.default_rng(42)
|
|
10
10
|
|
|
11
|
-
_faiss_index_types = [
|
|
11
|
+
_faiss_index_types = [
|
|
12
|
+
"flat",
|
|
13
|
+
"ivf",
|
|
14
|
+
"hnsw",
|
|
15
|
+
"lsh",
|
|
16
|
+
"scalar",
|
|
17
|
+
"pq",
|
|
18
|
+
"ivf_scalar",
|
|
19
|
+
"ivfpq",
|
|
20
|
+
"ivfpqr",
|
|
21
|
+
]
|
|
12
22
|
|
|
13
23
|
|
|
14
24
|
@pytest.fixture(scope="session")
|
|
@@ -35,6 +45,7 @@ BACKEND_PARAMS = [(Backend.FAISS, index_type) for index_type in _faiss_index_typ
|
|
|
35
45
|
(Backend.ANNOY, None),
|
|
36
46
|
(Backend.PYNNDESCENT, None),
|
|
37
47
|
(Backend.USEARCH, None),
|
|
48
|
+
(Backend.VOYAGER, None),
|
|
38
49
|
]
|
|
39
50
|
|
|
40
51
|
|
|
@@ -57,11 +68,48 @@ def vicinity_instance(request: pytest.FixtureRequest, items: list[str], vectors:
|
|
|
57
68
|
if index_type in ("pq", "ivfpq", "ivfpqr"):
|
|
58
69
|
# Use smaller values for pq indexes since the dataset is small
|
|
59
70
|
return Vicinity.from_vectors_and_items(
|
|
60
|
-
vectors,
|
|
71
|
+
vectors,
|
|
72
|
+
items,
|
|
73
|
+
backend_type=backend_type,
|
|
74
|
+
index_type=index_type,
|
|
75
|
+
m=2,
|
|
76
|
+
nbits=4,
|
|
61
77
|
)
|
|
62
78
|
else:
|
|
63
79
|
return Vicinity.from_vectors_and_items(
|
|
64
|
-
vectors,
|
|
80
|
+
vectors,
|
|
81
|
+
items,
|
|
82
|
+
backend_type=backend_type,
|
|
83
|
+
index_type=index_type,
|
|
84
|
+
nlist=2,
|
|
85
|
+
nbits=32,
|
|
65
86
|
)
|
|
66
87
|
|
|
67
88
|
return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
@pytest.fixture(params=BACKEND_PARAMS, ids=BACKEND_IDS)
|
|
92
|
+
def vicinity_instance_with_stored_vectors(
|
|
93
|
+
request: pytest.FixtureRequest, items: list[str], vectors: np.ndarray
|
|
94
|
+
) -> Vicinity:
|
|
95
|
+
"""Fixture providing a Vicinity instance for each backend type."""
|
|
96
|
+
backend_type, index_type = request.param
|
|
97
|
+
# Handle FAISS backend with specific FAISS index types
|
|
98
|
+
if backend_type == Backend.FAISS:
|
|
99
|
+
if index_type in ("pq", "ivfpq", "ivfpqr"):
|
|
100
|
+
# Use smaller values for pq indexes since the dataset is small
|
|
101
|
+
return Vicinity.from_vectors_and_items(
|
|
102
|
+
vectors, items, backend_type=backend_type, index_type=index_type, m=2, nbits=4, store_vectors=True
|
|
103
|
+
)
|
|
104
|
+
else:
|
|
105
|
+
return Vicinity.from_vectors_and_items(
|
|
106
|
+
vectors, items, backend_type=backend_type, index_type=index_type, nlist=2, nbits=32, store_vectors=True
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type, store_vectors=True)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
@pytest.fixture()
|
|
113
|
+
def vicinity_with_basic_backend(vectors: np.ndarray, items: list[str]) -> Vicinity:
|
|
114
|
+
"""Fixture providing a BasicBackend instance."""
|
|
115
|
+
return Vicinity.from_vectors_and_items(vectors, items, backend_type=Backend.BASIC, store_vectors=True)
|
|
@@ -131,8 +131,48 @@ def test_vicinity_save_and_load(tmp_path: Path, vicinity_instance: Vicinity) ->
|
|
|
131
131
|
"""
|
|
132
132
|
save_path = tmp_path / "vicinity_data"
|
|
133
133
|
vicinity_instance.save(save_path)
|
|
134
|
+
assert vicinity_instance.vector_store is None
|
|
134
135
|
|
|
135
|
-
Vicinity.load(save_path)
|
|
136
|
+
v = Vicinity.load(save_path)
|
|
137
|
+
assert v.vector_store is None
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
def test_vicinity_save_and_load_vector_store(tmp_path: Path, vicinity_instance_with_stored_vectors: Vicinity) -> None:
|
|
141
|
+
"""
|
|
142
|
+
Test Vicinity.save and Vicinity.load.
|
|
143
|
+
|
|
144
|
+
:param tmp_path: Temporary directory provided by pytest.
|
|
145
|
+
:param vicinity_instance: A Vicinity instance.
|
|
146
|
+
"""
|
|
147
|
+
save_path = tmp_path / "vicinity_data"
|
|
148
|
+
vicinity_instance_with_stored_vectors.save(save_path)
|
|
149
|
+
|
|
150
|
+
assert (save_path / "store").exists()
|
|
151
|
+
assert (save_path / "store" / "vectors.npy").exists()
|
|
152
|
+
|
|
153
|
+
v = Vicinity.load(save_path)
|
|
154
|
+
assert v.vector_store is not None
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def test_index_vector_store(vicinity_with_basic_backend: Vicinity, vectors: np.ndarray) -> None:
|
|
158
|
+
"""
|
|
159
|
+
Index vectors in the Vicinity instance.
|
|
160
|
+
|
|
161
|
+
:param vicinity_instance: A Vicinity instance.
|
|
162
|
+
:param vectors: Array of vectors to index.
|
|
163
|
+
"""
|
|
164
|
+
v = vicinity_with_basic_backend.get_vector_by_index(0)
|
|
165
|
+
assert np.allclose(v, vectors[0])
|
|
166
|
+
|
|
167
|
+
idx = [0, 1, 2, 3, 4, 10]
|
|
168
|
+
v = vicinity_with_basic_backend.get_vector_by_index(idx)
|
|
169
|
+
assert np.allclose(v, vectors[idx])
|
|
170
|
+
|
|
171
|
+
with pytest.raises(ValueError):
|
|
172
|
+
vicinity_with_basic_backend.get_vector_by_index([10_000])
|
|
173
|
+
|
|
174
|
+
with pytest.raises(ValueError):
|
|
175
|
+
vicinity_with_basic_backend.get_vector_by_index([-1])
|
|
136
176
|
|
|
137
177
|
|
|
138
178
|
def test_vicinity_insert_duplicate(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:
|
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
from typing import Union
|
|
2
2
|
|
|
3
3
|
from vicinity.backends.base import AbstractBackend
|
|
4
|
-
from vicinity.backends.basic import BasicBackend
|
|
4
|
+
from vicinity.backends.basic import BasicBackend, BasicVectorStore
|
|
5
5
|
from vicinity.datatypes import Backend
|
|
6
6
|
|
|
7
7
|
|
|
@@ -33,5 +33,10 @@ def get_backend_class(backend: Union[Backend, str]) -> type[AbstractBackend]:
|
|
|
33
33
|
|
|
34
34
|
return UsearchBackend
|
|
35
35
|
|
|
36
|
+
elif backend == Backend.VOYAGER:
|
|
37
|
+
from vicinity.backends.voyager import VoyagerBackend
|
|
36
38
|
|
|
37
|
-
|
|
39
|
+
return VoyagerBackend
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
__all__ = ["get_backend_class", "AbstractBackend", "BasicVectorStore"]
|
|
@@ -16,18 +16,18 @@ from vicinity.utils import Metric, normalize
|
|
|
16
16
|
@dataclass
|
|
17
17
|
class AnnoyArgs(BaseArgs):
|
|
18
18
|
dim: int = 0
|
|
19
|
-
metric:
|
|
19
|
+
metric: Metric = Metric.COSINE
|
|
20
|
+
internal_metric: str = "dot"
|
|
20
21
|
trees: int = 100
|
|
21
22
|
length: int | None = None
|
|
22
23
|
|
|
23
24
|
|
|
24
25
|
class AnnoyBackend(AbstractBackend[AnnoyArgs]):
|
|
25
26
|
argument_class = AnnoyArgs
|
|
26
|
-
supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN
|
|
27
|
-
inverse_metric_mapping = {
|
|
27
|
+
supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
|
|
28
|
+
inverse_metric_mapping: dict[Metric, str] = {
|
|
28
29
|
Metric.COSINE: "dot",
|
|
29
30
|
Metric.EUCLIDEAN: "euclidean",
|
|
30
|
-
Metric.INNER_PRODUCT: "dot",
|
|
31
31
|
}
|
|
32
32
|
|
|
33
33
|
def __init__(
|
|
@@ -56,18 +56,18 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
|
|
|
56
56
|
if metric_enum not in cls.supported_metrics:
|
|
57
57
|
raise ValueError(f"Metric '{metric_enum.value}' is not supported by AnnoyBackend.")
|
|
58
58
|
|
|
59
|
-
|
|
59
|
+
internal_metric = cls._map_metric_to_string(metric_enum)
|
|
60
60
|
|
|
61
|
-
if
|
|
61
|
+
if metric_enum == Metric.COSINE:
|
|
62
62
|
vectors = normalize(vectors)
|
|
63
63
|
|
|
64
64
|
dim = vectors.shape[1]
|
|
65
|
-
index = AnnoyIndex(f=dim, metric=
|
|
65
|
+
index = AnnoyIndex(f=dim, metric=internal_metric) # type: ignore
|
|
66
66
|
for i, vector in enumerate(vectors):
|
|
67
67
|
index.add_item(i, vector)
|
|
68
68
|
index.build(trees)
|
|
69
69
|
|
|
70
|
-
arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors)) # type: ignore
|
|
70
|
+
arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors), internal_metric=internal_metric) # type: ignore
|
|
71
71
|
return AnnoyBackend(index, arguments=arguments)
|
|
72
72
|
|
|
73
73
|
@property
|
|
@@ -88,8 +88,10 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
|
|
|
88
88
|
def load(cls: type[AnnoyBackend], base_path: Path) -> AnnoyBackend:
|
|
89
89
|
"""Load the vectors from a path."""
|
|
90
90
|
path = Path(base_path) / "index.bin"
|
|
91
|
+
|
|
91
92
|
arguments = AnnoyArgs.load(base_path / "arguments.json")
|
|
92
|
-
|
|
93
|
+
metric = cls._map_metric_to_string(arguments.metric)
|
|
94
|
+
index = AnnoyIndex(arguments.dim, metric) # type: ignore
|
|
93
95
|
index.load(str(path))
|
|
94
96
|
|
|
95
97
|
return cls(index, arguments=arguments)
|
|
@@ -106,11 +108,11 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
|
|
|
106
108
|
"""Query the backend."""
|
|
107
109
|
out = []
|
|
108
110
|
for vec in vectors:
|
|
109
|
-
if self.arguments.metric ==
|
|
111
|
+
if self.arguments.metric == Metric.COSINE:
|
|
110
112
|
vec = normalize(vec)
|
|
111
113
|
indices, scores = self.index.get_nns_by_vector(vec, k, include_distances=True)
|
|
112
114
|
scores_array = np.asarray(scores)
|
|
113
|
-
if self.arguments.metric ==
|
|
115
|
+
if self.arguments.metric == Metric.COSINE:
|
|
114
116
|
# Convert cosine similarity to cosine distance
|
|
115
117
|
scores_array = 1 - scores_array
|
|
116
118
|
out.append((np.asarray(indices), scores_array))
|
|
@@ -14,19 +14,25 @@ from vicinity.datatypes import Backend, QueryResult
|
|
|
14
14
|
|
|
15
15
|
@dataclass
|
|
16
16
|
class BaseArgs:
|
|
17
|
+
metric: Metric
|
|
18
|
+
|
|
17
19
|
def dump(self, file: Path) -> None:
|
|
18
20
|
"""Dump the arguments to a file."""
|
|
19
21
|
with open(file, "w") as f:
|
|
20
|
-
|
|
22
|
+
d = self.dict()
|
|
23
|
+
d["metric"] = d["metric"].value
|
|
24
|
+
json.dump(d, f)
|
|
21
25
|
|
|
22
26
|
@classmethod
|
|
23
27
|
def load(cls: type[ArgType], file: Path) -> ArgType:
|
|
24
28
|
"""Load the arguments from a file."""
|
|
25
29
|
with open(file, "r") as f:
|
|
26
|
-
|
|
30
|
+
data = json.load(f)
|
|
31
|
+
data["metric"] = Metric.from_string(data["metric"])
|
|
32
|
+
return cls(**data)
|
|
27
33
|
|
|
28
34
|
def dict(self) -> dict[str, Any]:
|
|
29
|
-
"""Dump the arguments to a
|
|
35
|
+
"""Dump the arguments to a dict."""
|
|
30
36
|
return asdict(self)
|
|
31
37
|
|
|
32
38
|
|
|
@@ -3,7 +3,7 @@ from __future__ import annotations
|
|
|
3
3
|
from abc import ABC, abstractmethod
|
|
4
4
|
from dataclasses import dataclass
|
|
5
5
|
from pathlib import Path
|
|
6
|
-
from typing import Any,
|
|
6
|
+
from typing import Any, Union
|
|
7
7
|
|
|
8
8
|
import numpy as np
|
|
9
9
|
from numpy import typing as npt
|
|
@@ -15,26 +15,61 @@ from vicinity.utils import Metric, normalize, normalize_or_copy
|
|
|
15
15
|
|
|
16
16
|
@dataclass
|
|
17
17
|
class BasicArgs(BaseArgs):
|
|
18
|
-
metric:
|
|
18
|
+
metric: Metric = Metric.COSINE
|
|
19
19
|
|
|
20
20
|
|
|
21
|
-
class
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
21
|
+
class BasicVectorStore:
|
|
22
|
+
def __init__(self, *, vectors: npt.NDArray, **kwargs: Any) -> None:
|
|
23
|
+
"""
|
|
24
|
+
A basic vector store that just stores vectors.
|
|
25
25
|
|
|
26
|
-
|
|
27
|
-
"""Initialize the backend."""
|
|
28
|
-
super().__init__(arguments)
|
|
26
|
+
Note that we use kwargs in order to use this class as a mixin.
|
|
29
27
|
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
28
|
+
:param vectors: The vectors to store.
|
|
29
|
+
:param **kwargs: Additional arguments. These are passed on to the super class.
|
|
30
|
+
"""
|
|
31
|
+
super().__init__(**kwargs)
|
|
32
|
+
self._vectors = vectors
|
|
33
33
|
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
return
|
|
34
|
+
def _update_precomputed_data(self) -> None:
|
|
35
|
+
"""Update precomputed data based on the metric."""
|
|
36
|
+
# NOTE: this is a no-op in the base implementation.
|
|
37
|
+
return
|
|
38
|
+
|
|
39
|
+
def get_by_index(self, indices: list[int]) -> npt.NDArray:
|
|
40
|
+
"""Get vectors by index."""
|
|
41
|
+
return self._vectors[indices]
|
|
42
|
+
|
|
43
|
+
def insert(self, vectors: npt.NDArray) -> None:
|
|
44
|
+
"""Insert vectors into the vector space."""
|
|
45
|
+
self._vectors = np.vstack([self._vectors, vectors])
|
|
46
|
+
self._update_precomputed_data()
|
|
47
|
+
|
|
48
|
+
def delete(self, indices: list[int]) -> None:
|
|
49
|
+
"""Deletes specific indices from the vector space."""
|
|
50
|
+
self._vectors = np.delete(self._vectors, indices, axis=0)
|
|
51
|
+
self._update_precomputed_data()
|
|
52
|
+
|
|
53
|
+
def save(self, folder: Path) -> None:
|
|
54
|
+
"""Save the vectors to a path."""
|
|
55
|
+
path = folder / "vectors.npy"
|
|
56
|
+
with open(path, "wb") as f:
|
|
57
|
+
np.save(f, self._vectors)
|
|
58
|
+
|
|
59
|
+
@staticmethod
|
|
60
|
+
def _load_vectors(folder: Path) -> npt.NDArray:
|
|
61
|
+
"""Load the vectors from a path."""
|
|
62
|
+
path = folder / "vectors.npy"
|
|
63
|
+
with open(path, "rb") as f:
|
|
64
|
+
vectors = np.load(f)
|
|
65
|
+
|
|
66
|
+
return vectors
|
|
67
|
+
|
|
68
|
+
@classmethod
|
|
69
|
+
def load(cls, folder: Path) -> BasicVectorStore:
|
|
70
|
+
"""Load the vectors from a path."""
|
|
71
|
+
vectors = cls._load_vectors(folder)
|
|
72
|
+
return cls(vectors=vectors)
|
|
38
73
|
|
|
39
74
|
@property
|
|
40
75
|
def dim(self) -> int:
|
|
@@ -55,10 +90,24 @@ class BasicBackend(AbstractBackend[BasicArgs], ABC):
|
|
|
55
90
|
self._vectors = matrix
|
|
56
91
|
self._update_precomputed_data()
|
|
57
92
|
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
93
|
+
|
|
94
|
+
class BasicBackend(BasicVectorStore, AbstractBackend[BasicArgs], ABC):
|
|
95
|
+
argument_class = BasicArgs
|
|
96
|
+
_vectors: npt.NDArray
|
|
97
|
+
supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
|
|
98
|
+
|
|
99
|
+
def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
|
|
100
|
+
"""Initialize the backend."""
|
|
101
|
+
super().__init__(vectors=vectors, arguments=arguments)
|
|
102
|
+
|
|
103
|
+
def __len__(self) -> int:
|
|
104
|
+
"""Get the number of vectors."""
|
|
105
|
+
return self.vectors.shape[0]
|
|
106
|
+
|
|
107
|
+
@property
|
|
108
|
+
def backend_type(self) -> Backend:
|
|
109
|
+
"""The type of the backend."""
|
|
110
|
+
return Backend.BASIC
|
|
62
111
|
|
|
63
112
|
@abstractmethod
|
|
64
113
|
def _dist(self, x: npt.NDArray) -> npt.NDArray:
|
|
@@ -72,11 +121,10 @@ class BasicBackend(AbstractBackend[BasicArgs], ABC):
|
|
|
72
121
|
if metric_enum not in cls.supported_metrics:
|
|
73
122
|
raise ValueError(f"Metric '{metric_enum.value}' is not supported by BasicBackend.")
|
|
74
123
|
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
if metric == "cosine":
|
|
124
|
+
arguments = BasicArgs(metric=metric_enum)
|
|
125
|
+
if metric_enum == Metric.COSINE:
|
|
78
126
|
return CosineBasicBackend(vectors, arguments)
|
|
79
|
-
elif
|
|
127
|
+
elif metric_enum == Metric.EUCLIDEAN:
|
|
80
128
|
return EuclideanBasicBackend(vectors, arguments)
|
|
81
129
|
else:
|
|
82
130
|
raise ValueError(f"Unsupported metric: {metric}")
|
|
@@ -84,23 +132,19 @@ class BasicBackend(AbstractBackend[BasicArgs], ABC):
|
|
|
84
132
|
@classmethod
|
|
85
133
|
def load(cls, folder: Path) -> BasicBackend:
|
|
86
134
|
"""Load the vectors from a path."""
|
|
87
|
-
path = folder / "vectors.npy"
|
|
88
135
|
arguments = BasicArgs.load(folder / "arguments.json")
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
if arguments.metric == "cosine":
|
|
136
|
+
vectors = cls._load_vectors(folder)
|
|
137
|
+
if arguments.metric == Metric.COSINE:
|
|
92
138
|
return CosineBasicBackend(vectors, arguments)
|
|
93
|
-
elif arguments.metric ==
|
|
139
|
+
elif arguments.metric == Metric.EUCLIDEAN:
|
|
94
140
|
return EuclideanBasicBackend(vectors, arguments)
|
|
95
141
|
else:
|
|
96
142
|
raise ValueError(f"Unsupported metric: {arguments.metric}")
|
|
97
143
|
|
|
98
144
|
def save(self, folder: Path) -> None:
|
|
99
145
|
"""Save the vectors to a path."""
|
|
100
|
-
|
|
146
|
+
super().save(folder)
|
|
101
147
|
self.arguments.dump(folder / "arguments.json")
|
|
102
|
-
with open(path, "wb") as f:
|
|
103
|
-
np.save(f, self._vectors)
|
|
104
148
|
|
|
105
149
|
def threshold(
|
|
106
150
|
self,
|
|
@@ -161,26 +205,12 @@ class BasicBackend(AbstractBackend[BasicArgs], ABC):
|
|
|
161
205
|
|
|
162
206
|
return out
|
|
163
207
|
|
|
164
|
-
def insert(self, vectors: npt.NDArray) -> None:
|
|
165
|
-
"""Insert vectors into the vector space."""
|
|
166
|
-
self._vectors = np.vstack([self._vectors, vectors])
|
|
167
|
-
self._update_precomputed_data()
|
|
168
|
-
|
|
169
|
-
def delete(self, indices: list[int]) -> None:
|
|
170
|
-
"""Deletes specific indices from the vector space."""
|
|
171
|
-
self._vectors = np.delete(self._vectors, indices, axis=0)
|
|
172
|
-
self._update_precomputed_data()
|
|
173
|
-
|
|
174
208
|
|
|
175
209
|
class CosineBasicBackend(BasicBackend):
|
|
176
210
|
def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
|
|
177
211
|
"""Initialize the cosine basic backend."""
|
|
178
|
-
super().__init__(arguments)
|
|
179
|
-
self._vectors = normalize_or_copy(
|
|
180
|
-
|
|
181
|
-
def _update_precomputed_data(self) -> None:
|
|
182
|
-
"""Update precomputed data for cosine similarity."""
|
|
183
|
-
pass
|
|
212
|
+
super().__init__(vectors=vectors, arguments=arguments)
|
|
213
|
+
self._vectors = normalize_or_copy(self._vectors)
|
|
184
214
|
|
|
185
215
|
def _dist(self, x: npt.NDArray) -> npt.NDArray:
|
|
186
216
|
"""Compute cosine distance."""
|
|
@@ -198,21 +228,12 @@ class CosineBasicBackend(BasicBackend):
|
|
|
198
228
|
class EuclideanBasicBackend(BasicBackend):
|
|
199
229
|
def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
|
|
200
230
|
"""Initialize the Euclidean basic backend."""
|
|
201
|
-
super().__init__(arguments)
|
|
202
|
-
self.
|
|
203
|
-
self._squared_norm_vectors: npt.NDArray | None = None
|
|
204
|
-
self._update_precomputed_data()
|
|
231
|
+
super().__init__(vectors=vectors, arguments=arguments)
|
|
232
|
+
self.squared_norm_vectors = (self._vectors**2).sum(1)
|
|
205
233
|
|
|
206
234
|
def _update_precomputed_data(self) -> None:
|
|
207
235
|
"""Update precomputed data for Euclidean distance."""
|
|
208
|
-
self.
|
|
209
|
-
|
|
210
|
-
@property
|
|
211
|
-
def squared_norm_vectors(self) -> npt.NDArray:
|
|
212
|
-
"""Return squared norms of vectors."""
|
|
213
|
-
if self._squared_norm_vectors is None:
|
|
214
|
-
self._squared_norm_vectors = (self._vectors**2).sum(1)
|
|
215
|
-
return self._squared_norm_vectors
|
|
236
|
+
self.squared_norm_vectors = (self._vectors**2).sum(1)
|
|
216
237
|
|
|
217
238
|
def _dist(self, x: npt.NDArray) -> npt.NDArray:
|
|
218
239
|
"""Compute Euclidean distance."""
|
|
@@ -6,7 +6,6 @@ from pathlib import Path
|
|
|
6
6
|
from typing import Any, Union
|
|
7
7
|
|
|
8
8
|
import faiss
|
|
9
|
-
import numpy as np
|
|
10
9
|
from numpy import typing as npt
|
|
11
10
|
|
|
12
11
|
from vicinity.backends.base import AbstractBackend, BaseArgs
|
|
@@ -37,7 +36,7 @@ TRAINABLE_INDEXES = (
|
|
|
37
36
|
class FaissArgs(BaseArgs):
|
|
38
37
|
dim: int = 0
|
|
39
38
|
index_type: str = "flat"
|
|
40
|
-
metric:
|
|
39
|
+
metric: Metric = Metric.COSINE
|
|
41
40
|
nlist: int = 100
|
|
42
41
|
m: int = 8
|
|
43
42
|
nbits: int = 8
|
|
@@ -123,7 +122,7 @@ class FaissBackend(AbstractBackend[FaissArgs]):
|
|
|
123
122
|
arguments = FaissArgs(
|
|
124
123
|
dim=dim,
|
|
125
124
|
index_type=index_type,
|
|
126
|
-
metric=metric_enum
|
|
125
|
+
metric=metric_enum,
|
|
127
126
|
nlist=nlist,
|
|
128
127
|
m=m,
|
|
129
128
|
nbits=nbits,
|
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
from dataclasses import dataclass
|
|
4
4
|
from pathlib import Path
|
|
5
|
-
from typing import Any,
|
|
5
|
+
from typing import Any, Union
|
|
6
6
|
|
|
7
7
|
from hnswlib import Index as HnswIndex
|
|
8
8
|
from numpy import typing as npt
|
|
@@ -15,7 +15,7 @@ from vicinity.utils import Metric
|
|
|
15
15
|
@dataclass
|
|
16
16
|
class HNSWArgs(BaseArgs):
|
|
17
17
|
dim: int = 0
|
|
18
|
-
metric:
|
|
18
|
+
metric: Metric = Metric.COSINE
|
|
19
19
|
ef_construction: int = 200
|
|
20
20
|
m: int = 16
|
|
21
21
|
|
|
@@ -58,7 +58,7 @@ class HNSWBackend(AbstractBackend[HNSWArgs]):
|
|
|
58
58
|
index = HnswIndex(space=metric, dim=dim)
|
|
59
59
|
index.init_index(max_elements=vectors.shape[0], ef_construction=ef_construction, M=m)
|
|
60
60
|
index.add_items(vectors)
|
|
61
|
-
arguments = HNSWArgs(dim=dim, metric=
|
|
61
|
+
arguments = HNSWArgs(dim=dim, metric=metric_enum, ef_construction=ef_construction, m=m)
|
|
62
62
|
return HNSWBackend(index, arguments=arguments)
|
|
63
63
|
|
|
64
64
|
@property
|
|
@@ -80,7 +80,8 @@ class HNSWBackend(AbstractBackend[HNSWArgs]):
|
|
|
80
80
|
"""Load the vectors from a path."""
|
|
81
81
|
path = Path(base_path) / "index.bin"
|
|
82
82
|
arguments = HNSWArgs.load(base_path / "arguments.json")
|
|
83
|
-
|
|
83
|
+
mapped_metric = cls.inverse_metric_mapping[arguments.metric]
|
|
84
|
+
index = HnswIndex(space=mapped_metric, dim=arguments.dim)
|
|
84
85
|
index.load_index(str(path))
|
|
85
86
|
return cls(index, arguments=arguments)
|
|
86
87
|
|
|
@@ -16,7 +16,7 @@ from vicinity.utils import Metric, normalize_or_copy
|
|
|
16
16
|
@dataclass
|
|
17
17
|
class PyNNDescentArgs(BaseArgs):
|
|
18
18
|
n_neighbors: int = 15
|
|
19
|
-
metric:
|
|
19
|
+
metric: Metric = Metric.COSINE
|
|
20
20
|
|
|
21
21
|
|
|
22
22
|
class PyNNDescentBackend(AbstractBackend[PyNNDescentArgs]):
|
|
@@ -49,7 +49,7 @@ class PyNNDescentBackend(AbstractBackend[PyNNDescentArgs]):
|
|
|
49
49
|
metric = metric_enum.value
|
|
50
50
|
|
|
51
51
|
index = NNDescent(vectors, n_neighbors=n_neighbors, metric=metric, **kwargs)
|
|
52
|
-
arguments = PyNNDescentArgs(n_neighbors=n_neighbors, metric=
|
|
52
|
+
arguments = PyNNDescentArgs(n_neighbors=n_neighbors, metric=metric_enum)
|
|
53
53
|
return cls(index=index, arguments=arguments)
|
|
54
54
|
|
|
55
55
|
def __len__(self) -> int:
|
|
@@ -105,10 +105,7 @@ class PyNNDescentBackend(AbstractBackend[PyNNDescentArgs]):
|
|
|
105
105
|
arguments = PyNNDescentArgs.load(base_path / "arguments.json")
|
|
106
106
|
vectors = np.load(Path(base_path) / "vectors.npy")
|
|
107
107
|
|
|
108
|
-
|
|
109
|
-
pynndescent_metric = metric_enum.value
|
|
110
|
-
|
|
111
|
-
index = NNDescent(vectors, n_neighbors=arguments.n_neighbors, metric=pynndescent_metric)
|
|
108
|
+
index = NNDescent(vectors, n_neighbors=arguments.n_neighbors, metric=arguments.metric.value)
|
|
112
109
|
|
|
113
110
|
# Load the neighbor graph if it was saved
|
|
114
111
|
neighbor_graph_path = base_path / "neighbor_graph.npy"
|
|
@@ -16,7 +16,7 @@ from vicinity.utils import Metric
|
|
|
16
16
|
@dataclass
|
|
17
17
|
class UsearchArgs(BaseArgs):
|
|
18
18
|
dim: int = 0
|
|
19
|
-
metric:
|
|
19
|
+
metric: Metric = Metric.COSINE
|
|
20
20
|
connectivity: int = 16
|
|
21
21
|
expansion_add: int = 128
|
|
22
22
|
expansion_search: int = 64
|
|
@@ -67,10 +67,10 @@ class UsearchBackend(AbstractBackend[UsearchArgs]):
|
|
|
67
67
|
expansion_add=expansion_add,
|
|
68
68
|
expansion_search=expansion_search,
|
|
69
69
|
)
|
|
70
|
-
index.add(keys=None, vectors=vectors) # type: ignore
|
|
70
|
+
index.add(keys=None, vectors=vectors) # type: ignore # None keys are allowed but not typed
|
|
71
71
|
arguments = UsearchArgs(
|
|
72
72
|
dim=dim,
|
|
73
|
-
metric=
|
|
73
|
+
metric=metric_enum,
|
|
74
74
|
connectivity=connectivity,
|
|
75
75
|
expansion_add=expansion_add,
|
|
76
76
|
expansion_search=expansion_search,
|
|
@@ -99,7 +99,7 @@ class UsearchBackend(AbstractBackend[UsearchArgs]):
|
|
|
99
99
|
|
|
100
100
|
index = UsearchIndex(
|
|
101
101
|
ndim=arguments.dim,
|
|
102
|
-
metric=arguments.metric,
|
|
102
|
+
metric=cls._map_metric_to_string(arguments.metric),
|
|
103
103
|
connectivity=arguments.connectivity,
|
|
104
104
|
expansion_add=arguments.expansion_add,
|
|
105
105
|
expansion_search=arguments.expansion_search,
|
|
@@ -122,7 +122,7 @@ class UsearchBackend(AbstractBackend[UsearchArgs]):
|
|
|
122
122
|
|
|
123
123
|
def insert(self, vectors: npt.NDArray) -> None:
|
|
124
124
|
"""Insert vectors into the backend."""
|
|
125
|
-
self.index.add(None, vectors) # type: ignore
|
|
125
|
+
self.index.add(None, vectors) # type: ignore # None keys are allowed, but not typed.
|
|
126
126
|
|
|
127
127
|
def delete(self, indices: list[int]) -> None:
|
|
128
128
|
"""Delete vectors from the index (not supported by Usearch)."""
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from typing import Any, Union
|
|
6
|
+
|
|
7
|
+
from numpy import typing as npt
|
|
8
|
+
from voyager import Index, Space
|
|
9
|
+
|
|
10
|
+
from vicinity.backends.base import AbstractBackend, BaseArgs
|
|
11
|
+
from vicinity.datatypes import Backend, QueryResult
|
|
12
|
+
from vicinity.utils import Metric
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@dataclass
|
|
16
|
+
class VoyagerArgs(BaseArgs):
|
|
17
|
+
dim: int = 0
|
|
18
|
+
metric: Metric = Metric.COSINE
|
|
19
|
+
ef_construction: int = 200
|
|
20
|
+
m: int = 16
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class VoyagerBackend(AbstractBackend[VoyagerArgs]):
|
|
24
|
+
argument_class = VoyagerArgs
|
|
25
|
+
supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
|
|
26
|
+
_metric_to_space = {
|
|
27
|
+
Metric.COSINE: Space.Cosine,
|
|
28
|
+
Metric.EUCLIDEAN: Space.Euclidean,
|
|
29
|
+
}
|
|
30
|
+
|
|
31
|
+
def __init__(
|
|
32
|
+
self,
|
|
33
|
+
index: Index,
|
|
34
|
+
arguments: VoyagerArgs,
|
|
35
|
+
) -> None:
|
|
36
|
+
"""Initialize the backend using vectors."""
|
|
37
|
+
super().__init__(arguments)
|
|
38
|
+
self.index = index
|
|
39
|
+
|
|
40
|
+
@classmethod
|
|
41
|
+
def from_vectors(
|
|
42
|
+
cls: type[VoyagerBackend],
|
|
43
|
+
vectors: npt.NDArray,
|
|
44
|
+
metric: Union[str, Metric],
|
|
45
|
+
ef_construction: int,
|
|
46
|
+
m: int,
|
|
47
|
+
**kwargs: Any,
|
|
48
|
+
) -> VoyagerBackend:
|
|
49
|
+
"""Create a new instance from vectors."""
|
|
50
|
+
metric_enum = Metric.from_string(metric)
|
|
51
|
+
|
|
52
|
+
if metric_enum not in cls.supported_metrics:
|
|
53
|
+
raise ValueError(f"Metric '{metric_enum.value}' is not supported by VoyagerBackend.")
|
|
54
|
+
|
|
55
|
+
space = cls._metric_to_space[metric_enum]
|
|
56
|
+
dim = vectors.shape[1]
|
|
57
|
+
index = Index(
|
|
58
|
+
space=space,
|
|
59
|
+
num_dimensions=dim,
|
|
60
|
+
M=m,
|
|
61
|
+
ef_construction=ef_construction,
|
|
62
|
+
)
|
|
63
|
+
index.add_items(vectors)
|
|
64
|
+
return cls(
|
|
65
|
+
index,
|
|
66
|
+
VoyagerArgs(dim=dim, metric=metric_enum, ef_construction=ef_construction, m=m),
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
def query(self, query: npt.NDArray, k: int) -> QueryResult:
|
|
70
|
+
"""Query the backend for the nearest neighbors."""
|
|
71
|
+
indices, distances = self.index.query(query, k)
|
|
72
|
+
return list(zip(indices, distances))
|
|
73
|
+
|
|
74
|
+
@classmethod
|
|
75
|
+
def load(cls: type[VoyagerBackend], base_path: Path) -> VoyagerBackend:
|
|
76
|
+
"""Load the vectors from a path."""
|
|
77
|
+
path = Path(base_path) / "index.bin"
|
|
78
|
+
arguments = VoyagerArgs.load(base_path / "arguments.json")
|
|
79
|
+
index = Index.load(str(path))
|
|
80
|
+
return cls(index, arguments=arguments)
|
|
81
|
+
|
|
82
|
+
def save(self, base_path: Path) -> None:
|
|
83
|
+
"""Save the vectors to a path."""
|
|
84
|
+
path = Path(base_path) / "index.bin"
|
|
85
|
+
self.index.save(str(path))
|
|
86
|
+
self.arguments.dump(base_path / "arguments.json")
|
|
87
|
+
|
|
88
|
+
def insert(self, vectors: npt.NDArray) -> None:
|
|
89
|
+
"""Insert vectors into the backend."""
|
|
90
|
+
self.index.add_items(vectors)
|
|
91
|
+
|
|
92
|
+
def delete(self, indices: list[int]) -> None:
|
|
93
|
+
"""Delete vectors from the backend."""
|
|
94
|
+
raise NotImplementedError("Deletion is not supported in Voyager backend.")
|
|
95
|
+
|
|
96
|
+
def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
|
|
97
|
+
"""Threshold the backend."""
|
|
98
|
+
out: list[npt.NDArray] = []
|
|
99
|
+
for x, y in self.query(vectors, len(self)):
|
|
100
|
+
out.append(x[y < threshold])
|
|
101
|
+
|
|
102
|
+
return out
|
|
103
|
+
|
|
104
|
+
@property
|
|
105
|
+
def backend_type(self) -> Backend:
|
|
106
|
+
"""The type of the backend."""
|
|
107
|
+
return Backend.VOYAGER
|
|
108
|
+
|
|
109
|
+
@property
|
|
110
|
+
def dim(self) -> int:
|
|
111
|
+
"""Get the dimension of the space."""
|
|
112
|
+
return self.index.num_dimensions
|
|
113
|
+
|
|
114
|
+
def __len__(self) -> int:
|
|
115
|
+
"""Get the number of vectors."""
|
|
116
|
+
return self.index.num_elements
|
|
@@ -2,7 +2,6 @@ from enum import Enum
|
|
|
2
2
|
from pathlib import Path
|
|
3
3
|
from typing import Iterable, List, Tuple, Union
|
|
4
4
|
|
|
5
|
-
import numpy as np
|
|
6
5
|
from numpy import typing as npt
|
|
7
6
|
|
|
8
7
|
PathLike = Union[str, Path]
|
|
@@ -22,3 +21,4 @@ class Backend(str, Enum):
|
|
|
22
21
|
PYNNDESCENT = "pynndescent"
|
|
23
22
|
FAISS = "faiss"
|
|
24
23
|
USEARCH = "usearch"
|
|
24
|
+
VOYAGER = "voyager"
|
|
@@ -1,2 +1,2 @@
|
|
|
1
|
-
__version_triple__ = (0, 3,
|
|
1
|
+
__version_triple__ = (0, 3, 2)
|
|
2
2
|
__version__ = ".".join(map(str, __version_triple__))
|
|
@@ -3,18 +3,17 @@
|
|
|
3
3
|
from __future__ import annotations
|
|
4
4
|
|
|
5
5
|
import logging
|
|
6
|
-
import time
|
|
7
6
|
from io import open
|
|
8
7
|
from pathlib import Path
|
|
9
8
|
from time import perf_counter
|
|
10
|
-
from typing import Any, Sequence, Union
|
|
9
|
+
from typing import Any, Iterable, Sequence, Union
|
|
11
10
|
|
|
12
11
|
import numpy as np
|
|
13
12
|
import orjson
|
|
14
13
|
from numpy import typing as npt
|
|
15
14
|
|
|
16
15
|
from vicinity import Metric
|
|
17
|
-
from vicinity.backends import AbstractBackend, BasicBackend, get_backend_class
|
|
16
|
+
from vicinity.backends import AbstractBackend, BasicBackend, BasicVectorStore, get_backend_class
|
|
18
17
|
from vicinity.datatypes import Backend, PathLike
|
|
19
18
|
|
|
20
19
|
logger = logging.getLogger(__name__)
|
|
@@ -33,6 +32,7 @@ class Vicinity:
|
|
|
33
32
|
items: Sequence[str],
|
|
34
33
|
backend: AbstractBackend,
|
|
35
34
|
metadata: Union[dict[str, Any], None] = None,
|
|
35
|
+
vector_store: BasicVectorStore | None = None,
|
|
36
36
|
) -> None:
|
|
37
37
|
"""
|
|
38
38
|
Initialize a Vicinity instance with an array and list of items.
|
|
@@ -42,6 +42,7 @@ class Vicinity:
|
|
|
42
42
|
aligned with the vectors.
|
|
43
43
|
:param backend: The backend to use for the vector space.
|
|
44
44
|
:param metadata: A dictionary containing metadata about the vector space.
|
|
45
|
+
:param vector_store: A simple vector store only used for storing actual vectors.
|
|
45
46
|
:raises ValueError: If the length of the items and vectors are not the same.
|
|
46
47
|
"""
|
|
47
48
|
if len(items) != len(backend):
|
|
@@ -51,6 +52,19 @@ class Vicinity:
|
|
|
51
52
|
self.items: list[str] = list(items)
|
|
52
53
|
self.backend: AbstractBackend = backend
|
|
53
54
|
self.metadata = metadata or {}
|
|
55
|
+
self.vector_store = vector_store
|
|
56
|
+
|
|
57
|
+
def get_vector_by_index(self, index: int | Iterable[int]) -> npt.NDArray:
|
|
58
|
+
"""Get a vector by index."""
|
|
59
|
+
if isinstance(index, int):
|
|
60
|
+
index = [index]
|
|
61
|
+
if not all(0 <= i < len(self.items) for i in index):
|
|
62
|
+
raise ValueError("Index out of bounds.")
|
|
63
|
+
if self.vector_store is None:
|
|
64
|
+
raise ValueError(
|
|
65
|
+
"No vector store was provided. To get items by index, create a vicinity index by passing store_vectors=True on index creation."
|
|
66
|
+
)
|
|
67
|
+
return self.vector_store.get_by_index(list(index))
|
|
54
68
|
|
|
55
69
|
def __len__(self) -> int:
|
|
56
70
|
"""The number of the items in the vector space."""
|
|
@@ -62,6 +76,7 @@ class Vicinity:
|
|
|
62
76
|
vectors: npt.NDArray,
|
|
63
77
|
items: Sequence[str],
|
|
64
78
|
backend_type: Backend | str = Backend.BASIC,
|
|
79
|
+
store_vectors: bool = False,
|
|
65
80
|
**kwargs: Any,
|
|
66
81
|
) -> Vicinity:
|
|
67
82
|
"""
|
|
@@ -70,6 +85,7 @@ class Vicinity:
|
|
|
70
85
|
:param vectors: The vectors to use.
|
|
71
86
|
:param items: The items to use.
|
|
72
87
|
:param backend_type: The type of backend to use.
|
|
88
|
+
:param store_vectors: Whether to store the raw vectors in the backend.
|
|
73
89
|
:param **kwargs: Additional arguments to pass to the backend.
|
|
74
90
|
:return: A Vicinity instance.
|
|
75
91
|
"""
|
|
@@ -77,8 +93,12 @@ class Vicinity:
|
|
|
77
93
|
backend_cls = get_backend_class(backend_type)
|
|
78
94
|
arguments = backend_cls.argument_class(**kwargs)
|
|
79
95
|
backend = backend_cls.from_vectors(vectors, **arguments.dict())
|
|
96
|
+
if store_vectors:
|
|
97
|
+
vector_store = BasicVectorStore(vectors=vectors)
|
|
98
|
+
else:
|
|
99
|
+
vector_store = None
|
|
80
100
|
|
|
81
|
-
return cls(items, backend)
|
|
101
|
+
return cls(items, backend, vector_store=vector_store)
|
|
82
102
|
|
|
83
103
|
@property
|
|
84
104
|
def dim(self) -> int:
|
|
@@ -166,6 +186,10 @@ class Vicinity:
|
|
|
166
186
|
file_handle.write(orjson.dumps(items_dict))
|
|
167
187
|
|
|
168
188
|
self.backend.save(path)
|
|
189
|
+
if self.vector_store is not None:
|
|
190
|
+
store_path = path / "store"
|
|
191
|
+
store_path.mkdir(exist_ok=overwrite)
|
|
192
|
+
self.vector_store.save(store_path)
|
|
169
193
|
|
|
170
194
|
@classmethod
|
|
171
195
|
def load(cls, filename: PathLike) -> Vicinity:
|
|
@@ -190,8 +214,12 @@ class Vicinity:
|
|
|
190
214
|
|
|
191
215
|
backend_cls: type[AbstractBackend] = get_backend_class(backend_type)
|
|
192
216
|
backend = backend_cls.load(folder_path)
|
|
217
|
+
if Path(folder_path / "store").exists():
|
|
218
|
+
vector_store = BasicVectorStore.load(folder_path / "store")
|
|
219
|
+
else:
|
|
220
|
+
vector_store = None
|
|
193
221
|
|
|
194
|
-
instance = cls(items, backend, metadata=metadata)
|
|
222
|
+
instance = cls(items, backend, metadata=metadata, vector_store=vector_store)
|
|
195
223
|
|
|
196
224
|
return instance
|
|
197
225
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: vicinity
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: Lightweight Nearest Neighbors with Flexible Backends
|
|
5
5
|
Author-email: Stéphan Tulkens <stephantul@gmail.com>, Thomas van Dongen <thomas123@live.nl>
|
|
6
6
|
License: MIT License
|
|
@@ -67,6 +67,8 @@ Provides-Extra: faiss
|
|
|
67
67
|
Requires-Dist: faiss-cpu; extra == "faiss"
|
|
68
68
|
Provides-Extra: usearch
|
|
69
69
|
Requires-Dist: usearch; extra == "usearch"
|
|
70
|
+
Provides-Extra: voyager
|
|
71
|
+
Requires-Dist: voyager; extra == "voyager"
|
|
70
72
|
Provides-Extra: all
|
|
71
73
|
Requires-Dist: hnswlib; extra == "all"
|
|
72
74
|
Requires-Dist: pynndescent>=0.5.10; extra == "all"
|
|
@@ -76,13 +78,19 @@ Requires-Dist: numpy>=1.24.0; extra == "all"
|
|
|
76
78
|
Requires-Dist: annoy; extra == "all"
|
|
77
79
|
Requires-Dist: faiss-cpu; extra == "all"
|
|
78
80
|
Requires-Dist: usearch; extra == "all"
|
|
81
|
+
Requires-Dist: voyager; extra == "all"
|
|
79
82
|
|
|
80
|
-
<div align="center">
|
|
81
|
-
|
|
82
|
-
# Vicinity: Lightweight Nearest Neighbors
|
|
83
83
|
|
|
84
|
+
<div align="center">
|
|
85
|
+
<picture>
|
|
86
|
+
<img width="40%" alt="Model2Vec logo" src="assets/images/vicinity_logo.png">
|
|
87
|
+
</picture>
|
|
88
|
+
</a>
|
|
84
89
|
</div>
|
|
85
90
|
|
|
91
|
+
<div align="center">
|
|
92
|
+
<h2>Lightweight Nearest Neighbors with Flexible Backends</h2>
|
|
93
|
+
</div>
|
|
86
94
|
|
|
87
95
|
<div align="center">
|
|
88
96
|
<h2>
|
|
@@ -118,7 +126,7 @@ Install the package with:
|
|
|
118
126
|
```bash
|
|
119
127
|
pip install vicinity
|
|
120
128
|
```
|
|
121
|
-
|
|
129
|
+
Optionally, [install any of the supported backends](#installation), or simply install all of them with:
|
|
122
130
|
```bash
|
|
123
131
|
pip install vicinity[all]
|
|
124
132
|
```
|
|
@@ -134,7 +142,12 @@ items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
|
|
|
134
142
|
vectors = np.random.rand(len(items), 128)
|
|
135
143
|
|
|
136
144
|
# Initialize the Vicinity instance (using the basic backend and cosine metric)
|
|
137
|
-
vicinity = Vicinity.from_vectors_and_items(
|
|
145
|
+
vicinity = Vicinity.from_vectors_and_items(
|
|
146
|
+
vectors=vectors,
|
|
147
|
+
items=items,
|
|
148
|
+
backend_type=Backend.BASIC,
|
|
149
|
+
metric=Metric.COSINE
|
|
150
|
+
)
|
|
138
151
|
|
|
139
152
|
# Create a query vector
|
|
140
153
|
query_vector = np.random.rand(128)
|
|
@@ -189,7 +202,7 @@ The following backends are supported:
|
|
|
189
202
|
- `ivf_scalar`: Inverted file search with scalar quantizer.
|
|
190
203
|
- `ivfpq`: Inverted file search with product quantizer.
|
|
191
204
|
- `ivfpqr`: Inverted file search with product quantizer and refinement.
|
|
192
|
-
|
|
205
|
+
- [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors.
|
|
193
206
|
|
|
194
207
|
|
|
195
208
|
|
|
@@ -200,7 +213,7 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
200
213
|
|
|
201
214
|
| Backend | Parameter | Description | Default Value |
|
|
202
215
|
|-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
|
|
203
|
-
| **
|
|
216
|
+
| **ANNOY** | `metric` | Similarity metric to use (`dot`, `euclidean`, `cosine`). | `"cosine"` |
|
|
204
217
|
| | `trees` | Number of trees to use for indexing. | `100` |
|
|
205
218
|
| | `length` | Optional length of the dataset. | `None` |
|
|
206
219
|
| **FAISS** | `metric` | Similarity metric to use (`cosine`, `l2`). | `"cosine"` |
|
|
@@ -212,13 +225,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
212
225
|
| **HNSW** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
213
226
|
| | `ef_construction` | Size of the dynamic list during index construction. | `200` |
|
|
214
227
|
| | `m` | Number of connections per layer. | `16` |
|
|
215
|
-
| **
|
|
228
|
+
| **PYNNDESCENT** | `metric` | Similarity metric to use (`cosine`, `euclidean`, `manhattan`). | `"cosine"` |
|
|
216
229
|
| | `n_neighbors` | Number of neighbors to use for search. | `15` |
|
|
217
|
-
| **
|
|
230
|
+
| **USEARCH** | `metric` | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`). | `"cos"` |
|
|
218
231
|
| | `connectivity` | Number of connections per node in the graph. | `16` |
|
|
219
232
|
| | `expansion_add` | Number of candidates considered during graph construction. | `128` |
|
|
220
233
|
| | `expansion_search` | Number of candidates considered during search. | `64` |
|
|
221
|
-
|
|
234
|
+
| **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
235
|
+
| | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` |
|
|
236
|
+
| | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` |
|
|
222
237
|
|
|
223
238
|
## Installation
|
|
224
239
|
The following installation options are available:
|
|
@@ -235,6 +250,7 @@ pip install vicinity[faiss]
|
|
|
235
250
|
pip install vicinity[hnsw]
|
|
236
251
|
pip install vicinity[pynndescent]
|
|
237
252
|
pip install vicinity[usearch]
|
|
253
|
+
pip install vicinity[voyager]
|
|
238
254
|
```
|
|
239
255
|
|
|
240
256
|
## License
|
|
@@ -6,6 +6,7 @@ README.md
|
|
|
6
6
|
pyproject.toml
|
|
7
7
|
uv.lock
|
|
8
8
|
.github/workflows/ci.yaml
|
|
9
|
+
assets/images/vicinity_logo.png
|
|
9
10
|
tests/conftest.py
|
|
10
11
|
tests/test_utils.py
|
|
11
12
|
tests/test_vicinity.py
|
|
@@ -27,4 +28,5 @@ vicinity/backends/basic.py
|
|
|
27
28
|
vicinity/backends/faiss.py
|
|
28
29
|
vicinity/backends/hnsw.py
|
|
29
30
|
vicinity/backends/pynndescent.py
|
|
30
|
-
vicinity/backends/usearch.py
|
|
31
|
+
vicinity/backends/usearch.py
|
|
32
|
+
vicinity/backends/voyager.py
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|