vicinity 0.3.1__tar.gz → 0.3.3__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.3}/Makefile +1 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/PKG-INFO +35 -14
- {vicinity-0.3.1 → vicinity-0.3.3}/README.md +31 -13
- vicinity-0.3.3/assets/images/vicinity_logo.png +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/pyproject.toml +4 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/tests/conftest.py +51 -3
- {vicinity-0.3.1 → vicinity-0.3.3}/tests/test_vicinity.py +70 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/uv.lock +41 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/__init__.py +7 -2
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/annoy.py +13 -11
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/base.py +9 -3
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/basic.py +81 -60
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/faiss.py +2 -3
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/hnsw.py +5 -4
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/pynndescent.py +3 -6
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/backends/usearch.py +5 -5
- vicinity-0.3.3/vicinity/backends/voyager.py +116 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/datatypes.py +1 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/version.py +1 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/vicinity.py +37 -5
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity.egg-info/PKG-INFO +35 -14
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity.egg-info/SOURCES.txt +3 -1
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity.egg-info/requires.txt +4 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/.github/workflows/ci.yaml +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/.gitignore +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/.pre-commit-config.yaml +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/LICENSE +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/setup.cfg +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/tests/test_utils.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/__init__.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/py.typed +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity/utils.py +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/vicinity.egg-info/dependency_links.txt +0 -0
- {vicinity-0.3.1 → vicinity-0.3.3}/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.3
|
|
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,16 +142,25 @@ 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)
|
|
141
154
|
|
|
142
155
|
# Query for nearest neighbors with a top-k search
|
|
143
|
-
results = vicinity.query(
|
|
156
|
+
results = vicinity.query(query_vector, k=3)
|
|
144
157
|
|
|
145
158
|
# Query for nearest neighbors with a threshold search
|
|
146
|
-
results = vicinity.query_threshold(
|
|
159
|
+
results = vicinity.query_threshold(query_vector, threshold=0.9)
|
|
160
|
+
|
|
161
|
+
# Query with a list of query vectors
|
|
162
|
+
query_vectors = np.random.rand(5, 128)
|
|
163
|
+
results = vicinity.query(query_vectors, k=3)
|
|
147
164
|
```
|
|
148
165
|
|
|
149
166
|
Saving and loading a vector store:
|
|
@@ -178,7 +195,7 @@ The following backends are supported:
|
|
|
178
195
|
- [HNSW](https://github.com/nmslib/hnswlib): Hierarchical Navigable Small World Graph (HNSW) for ANN search using hnswlib.
|
|
179
196
|
- [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
|
|
180
197
|
- [ANNOY](https://github.com/spotify/annoy): "Approximate Nearest Neighbors Oh Yeah" for approximate nearest neighbor search.
|
|
181
|
-
- [
|
|
198
|
+
- [PYNNDESCENT](https://github.com/lmcinnes/pynndescent): ANN search using PyNNDescent.
|
|
182
199
|
- [FAISS](https://github.com/facebookresearch/faiss): All FAISS indexes are supported:
|
|
183
200
|
- `flat`: Exact search.
|
|
184
201
|
- `ivf`: Inverted file search.
|
|
@@ -189,7 +206,7 @@ The following backends are supported:
|
|
|
189
206
|
- `ivf_scalar`: Inverted file search with scalar quantizer.
|
|
190
207
|
- `ivfpq`: Inverted file search with product quantizer.
|
|
191
208
|
- `ivfpqr`: Inverted file search with product quantizer and refinement.
|
|
192
|
-
|
|
209
|
+
- [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
210
|
|
|
194
211
|
|
|
195
212
|
|
|
@@ -200,7 +217,8 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
200
217
|
|
|
201
218
|
| Backend | Parameter | Description | Default Value |
|
|
202
219
|
|-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
|
|
203
|
-
| **
|
|
220
|
+
| **BASIC** | `metric` | Similarity metric to use (`cosine`, `euclidean`). | `"cosine"` |
|
|
221
|
+
| **ANNOY** | `metric` | Similarity metric to use (`dot`, `euclidean`, `cosine`). | `"cosine"` |
|
|
204
222
|
| | `trees` | Number of trees to use for indexing. | `100` |
|
|
205
223
|
| | `length` | Optional length of the dataset. | `None` |
|
|
206
224
|
| **FAISS** | `metric` | Similarity metric to use (`cosine`, `l2`). | `"cosine"` |
|
|
@@ -212,13 +230,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
212
230
|
| **HNSW** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
213
231
|
| | `ef_construction` | Size of the dynamic list during index construction. | `200` |
|
|
214
232
|
| | `m` | Number of connections per layer. | `16` |
|
|
215
|
-
| **
|
|
233
|
+
| **PYNNDESCENT** | `metric` | Similarity metric to use (`cosine`, `euclidean`, `manhattan`). | `"cosine"` |
|
|
216
234
|
| | `n_neighbors` | Number of neighbors to use for search. | `15` |
|
|
217
|
-
| **
|
|
235
|
+
| **USEARCH** | `metric` | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`). | `"cos"` |
|
|
218
236
|
| | `connectivity` | Number of connections per node in the graph. | `16` |
|
|
219
237
|
| | `expansion_add` | Number of candidates considered during graph construction. | `128` |
|
|
220
238
|
| | `expansion_search` | Number of candidates considered during search. | `64` |
|
|
221
|
-
|
|
239
|
+
| **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
240
|
+
| | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` |
|
|
241
|
+
| | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` |
|
|
222
242
|
|
|
223
243
|
## Installation
|
|
224
244
|
The following installation options are available:
|
|
@@ -235,6 +255,7 @@ pip install vicinity[faiss]
|
|
|
235
255
|
pip install vicinity[hnsw]
|
|
236
256
|
pip install vicinity[pynndescent]
|
|
237
257
|
pip install vicinity[usearch]
|
|
258
|
+
pip install vicinity[voyager]
|
|
238
259
|
```
|
|
239
260
|
|
|
240
261
|
## 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,16 +60,25 @@ 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)
|
|
62
72
|
|
|
63
73
|
# Query for nearest neighbors with a top-k search
|
|
64
|
-
results = vicinity.query(
|
|
74
|
+
results = vicinity.query(query_vector, k=3)
|
|
65
75
|
|
|
66
76
|
# Query for nearest neighbors with a threshold search
|
|
67
|
-
results = vicinity.query_threshold(
|
|
77
|
+
results = vicinity.query_threshold(query_vector, threshold=0.9)
|
|
78
|
+
|
|
79
|
+
# Query with a list of query vectors
|
|
80
|
+
query_vectors = np.random.rand(5, 128)
|
|
81
|
+
results = vicinity.query(query_vectors, k=3)
|
|
68
82
|
```
|
|
69
83
|
|
|
70
84
|
Saving and loading a vector store:
|
|
@@ -99,7 +113,7 @@ The following backends are supported:
|
|
|
99
113
|
- [HNSW](https://github.com/nmslib/hnswlib): Hierarchical Navigable Small World Graph (HNSW) for ANN search using hnswlib.
|
|
100
114
|
- [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
|
|
101
115
|
- [ANNOY](https://github.com/spotify/annoy): "Approximate Nearest Neighbors Oh Yeah" for approximate nearest neighbor search.
|
|
102
|
-
- [
|
|
116
|
+
- [PYNNDESCENT](https://github.com/lmcinnes/pynndescent): ANN search using PyNNDescent.
|
|
103
117
|
- [FAISS](https://github.com/facebookresearch/faiss): All FAISS indexes are supported:
|
|
104
118
|
- `flat`: Exact search.
|
|
105
119
|
- `ivf`: Inverted file search.
|
|
@@ -110,7 +124,7 @@ The following backends are supported:
|
|
|
110
124
|
- `ivf_scalar`: Inverted file search with scalar quantizer.
|
|
111
125
|
- `ivfpq`: Inverted file search with product quantizer.
|
|
112
126
|
- `ivfpqr`: Inverted file search with product quantizer and refinement.
|
|
113
|
-
|
|
127
|
+
- [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
128
|
|
|
115
129
|
|
|
116
130
|
|
|
@@ -121,7 +135,8 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
121
135
|
|
|
122
136
|
| Backend | Parameter | Description | Default Value |
|
|
123
137
|
|-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
|
|
124
|
-
| **
|
|
138
|
+
| **BASIC** | `metric` | Similarity metric to use (`cosine`, `euclidean`). | `"cosine"` |
|
|
139
|
+
| **ANNOY** | `metric` | Similarity metric to use (`dot`, `euclidean`, `cosine`). | `"cosine"` |
|
|
125
140
|
| | `trees` | Number of trees to use for indexing. | `100` |
|
|
126
141
|
| | `length` | Optional length of the dataset. | `None` |
|
|
127
142
|
| **FAISS** | `metric` | Similarity metric to use (`cosine`, `l2`). | `"cosine"` |
|
|
@@ -133,13 +148,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
|
|
|
133
148
|
| **HNSW** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
134
149
|
| | `ef_construction` | Size of the dynamic list during index construction. | `200` |
|
|
135
150
|
| | `m` | Number of connections per layer. | `16` |
|
|
136
|
-
| **
|
|
151
|
+
| **PYNNDESCENT** | `metric` | Similarity metric to use (`cosine`, `euclidean`, `manhattan`). | `"cosine"` |
|
|
137
152
|
| | `n_neighbors` | Number of neighbors to use for search. | `15` |
|
|
138
|
-
| **
|
|
153
|
+
| **USEARCH** | `metric` | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`). | `"cos"` |
|
|
139
154
|
| | `connectivity` | Number of connections per node in the graph. | `16` |
|
|
140
155
|
| | `expansion_add` | Number of candidates considered during graph construction. | `128` |
|
|
141
156
|
| | `expansion_search` | Number of candidates considered during search. | `64` |
|
|
142
|
-
|
|
157
|
+
| **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` |
|
|
158
|
+
| | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` |
|
|
159
|
+
| | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` |
|
|
143
160
|
|
|
144
161
|
## Installation
|
|
145
162
|
The following installation options are available:
|
|
@@ -156,6 +173,7 @@ pip install vicinity[faiss]
|
|
|
156
173
|
pip install vicinity[hnsw]
|
|
157
174
|
pip install vicinity[pynndescent]
|
|
158
175
|
pip install vicinity[usearch]
|
|
176
|
+
pip install vicinity[voyager]
|
|
159
177
|
```
|
|
160
178
|
|
|
161
179
|
## 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_and_store(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_and_store: 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_and_store.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_and_store.get_vector_by_index(idx)
|
|
169
|
+
assert np.allclose(v, vectors[idx])
|
|
170
|
+
|
|
171
|
+
with pytest.raises(ValueError):
|
|
172
|
+
vicinity_with_basic_backend_and_store.get_vector_by_index([10_000])
|
|
173
|
+
|
|
174
|
+
with pytest.raises(ValueError):
|
|
175
|
+
vicinity_with_basic_backend_and_store.get_vector_by_index([-1])
|
|
136
176
|
|
|
137
177
|
|
|
138
178
|
def test_vicinity_insert_duplicate(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:
|
|
@@ -163,6 +203,35 @@ def test_vicinity_delete_nonexistent(vicinity_instance: Vicinity) -> None:
|
|
|
163
203
|
vicinity_instance.delete(["item10002"])
|
|
164
204
|
|
|
165
205
|
|
|
206
|
+
def test_vicinity_insert_with_store(vicinity_with_basic_backend_and_store: Vicinity) -> None:
|
|
207
|
+
"""
|
|
208
|
+
Test that Vicinity.insert raises ValueError when trying to insert vectors into a Vicinity instance with stored vectors.
|
|
209
|
+
|
|
210
|
+
:param vicinity_with_basic_backend_and_store: A Vicinity instance with stored vectors.
|
|
211
|
+
"""
|
|
212
|
+
new_item = ["item10002"]
|
|
213
|
+
new_vector = np.full((1, vicinity_with_basic_backend_and_store.dim), 0.5)
|
|
214
|
+
|
|
215
|
+
vicinity_with_basic_backend_and_store.insert(new_item, new_vector)
|
|
216
|
+
assert vicinity_with_basic_backend_and_store.vector_store is not None
|
|
217
|
+
assert len(vicinity_with_basic_backend_and_store) == len(vicinity_with_basic_backend_and_store.vector_store)
|
|
218
|
+
|
|
219
|
+
|
|
220
|
+
def test_vicinity_delete_with_store(vicinity_with_basic_backend_and_store: Vicinity) -> None:
|
|
221
|
+
"""
|
|
222
|
+
Test Vicinity.delete method by verifying that the vector for a deleted item is not returned in subsequent queries.
|
|
223
|
+
|
|
224
|
+
:param vicinity_with_basic_backend_and_store: A Vicinity instance.
|
|
225
|
+
"""
|
|
226
|
+
assert vicinity_with_basic_backend_and_store.vector_store is not None
|
|
227
|
+
# Delete "item2" from the Vicinity instance
|
|
228
|
+
vicinity_with_basic_backend_and_store.delete(["item2"])
|
|
229
|
+
|
|
230
|
+
# Ensure "item2" is no longer in the items list
|
|
231
|
+
assert "item2" not in vicinity_with_basic_backend_and_store.items
|
|
232
|
+
assert len(vicinity_with_basic_backend_and_store) == len(vicinity_with_basic_backend_and_store.vector_store)
|
|
233
|
+
|
|
234
|
+
|
|
166
235
|
def test_vicinity_insert_mismatched_lengths(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:
|
|
167
236
|
"""
|
|
168
237
|
Test that Vicinity.insert raises ValueError when tokens and vectors lengths do not match.
|
|
@@ -1087,7 +1087,7 @@ wheels = [
|
|
|
1087
1087
|
|
|
1088
1088
|
[[package]]
|
|
1089
1089
|
name = "vicinity"
|
|
1090
|
-
version = "0.3.
|
|
1090
|
+
version = "0.3.1"
|
|
1091
1091
|
source = { editable = "." }
|
|
1092
1092
|
dependencies = [
|
|
1093
1093
|
{ name = "numpy" },
|
|
@@ -1105,6 +1105,7 @@ all = [
|
|
|
1105
1105
|
{ name = "numpy" },
|
|
1106
1106
|
{ name = "pynndescent" },
|
|
1107
1107
|
{ name = "usearch" },
|
|
1108
|
+
{ name = "voyager" },
|
|
1108
1109
|
]
|
|
1109
1110
|
annoy = [
|
|
1110
1111
|
{ name = "annoy" },
|
|
@@ -1134,6 +1135,9 @@ pynndescent = [
|
|
|
1134
1135
|
usearch = [
|
|
1135
1136
|
{ name = "usearch" },
|
|
1136
1137
|
]
|
|
1138
|
+
voyager = [
|
|
1139
|
+
{ name = "voyager" },
|
|
1140
|
+
]
|
|
1137
1141
|
|
|
1138
1142
|
[package.metadata]
|
|
1139
1143
|
requires-dist = [
|
|
@@ -1164,6 +1168,8 @@ requires-dist = [
|
|
|
1164
1168
|
{ name = "tqdm" },
|
|
1165
1169
|
{ name = "usearch", marker = "extra == 'all'" },
|
|
1166
1170
|
{ name = "usearch", marker = "extra == 'usearch'" },
|
|
1171
|
+
{ name = "voyager", marker = "extra == 'all'" },
|
|
1172
|
+
{ name = "voyager", marker = "extra == 'voyager'" },
|
|
1167
1173
|
]
|
|
1168
1174
|
|
|
1169
1175
|
[[package]]
|
|
@@ -1180,6 +1186,40 @@ wheels = [
|
|
|
1180
1186
|
{ url = "https://files.pythonhosted.org/packages/ae/92/78324ff89391e00c8f4cf6b8526c41c6ef36b4ea2d2c132250b1a6fc2b8d/virtualenv-20.27.1-py3-none-any.whl", hash = "sha256:f11f1b8a29525562925f745563bfd48b189450f61fb34c4f9cc79dd5aa32a1f4", size = 3117838 },
|
|
1181
1187
|
]
|
|
1182
1188
|
|
|
1189
|
+
[[package]]
|
|
1190
|
+
name = "voyager"
|
|
1191
|
+
version = "2.0.9"
|
|
1192
|
+
source = { registry = "https://pypi.org/simple" }
|
|
1193
|
+
dependencies = [
|
|
1194
|
+
{ name = "numpy" },
|
|
1195
|
+
]
|
|
1196
|
+
wheels = [
|
|
1197
|
+
{ url = "https://files.pythonhosted.org/packages/79/d9/62399eb8a907da7f9eba0e2facf78925d54ba644a9be671cd98e7b4f547f/voyager-2.0.9-cp310-cp310-macosx_10_13_universal2.whl", hash = "sha256:f3cfe84666314daaf51836f77dabaafeb064bcbe5d9435d9c97886af19ae2876", size = 731200 },
|
|
1198
|
+
{ url = "https://files.pythonhosted.org/packages/51/70/0d90ebf90ca6fa415513b63925f4b21aeebd9f6400ecce5dffdb3743039f/voyager-2.0.9-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:956a60f864a586991cdde09373caf5630ab3319c47b6cdad4bf05ea8ba96a578", size = 377180 },
|
|
1199
|
+
{ url = "https://files.pythonhosted.org/packages/6e/af/76bb88b390bbf9c1dcbab90acd669168294c6d20b6a230b7333b432eee1b/voyager-2.0.9-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:56e6e7d6f1597207141d224d5420dc690ae27a2feb6c4164180cecc453cd5816", size = 356644 },
|
|
1200
|
+
{ url = "https://files.pythonhosted.org/packages/2b/b9/a514e2f88b7cdf35cb182064d5859258488cfe7e7b82bfff012d9cffd746/voyager-2.0.9-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7be032a8ae5447ce289f1c95f349e825e9d4edc915db7ea1b2c273940de2acbe", size = 4211154 },
|
|
1201
|
+
{ url = "https://files.pythonhosted.org/packages/db/95/1032e54e27eeac13f9affaf1fdeb834cc389f728316c9e7244c0e57c3e98/voyager-2.0.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ed85d23006f98091a137f637894560d1020ba25217b5332729417ac0a200b0bd", size = 4280843 },
|
|
1202
|
+
{ url = "https://files.pythonhosted.org/packages/2c/dc/395551a312110ad53cdac4f0126708d02c2d5c11c9f63ecf6a64985fab56/voyager-2.0.9-cp310-cp310-win_amd64.whl", hash = "sha256:053f22b23bd090f6376cf026c31653d3c040f9eaf0c6d1267e766a442a5b8ff8", size = 199722 },
|
|
1203
|
+
{ url = "https://files.pythonhosted.org/packages/c2/17/ace9aaf1e65553dae2ce3ba947e922e6bd16aa89abe558543471b5fe065c/voyager-2.0.9-cp311-cp311-macosx_10_13_universal2.whl", hash = "sha256:d270b29bafb5a133abd611a03b1b9ec7b03a44361966609abef72dde2518eb5e", size = 734474 },
|
|
1204
|
+
{ url = "https://files.pythonhosted.org/packages/3b/73/83d047ad7f6cb4fad426ed6f8b6c249c3ffecffa85ebb017e6d4e725aa0c/voyager-2.0.9-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:216c517dd66e0095f5f14277627a233c081797119028737f71168f094bacb9d0", size = 378460 },
|
|
1205
|
+
{ url = "https://files.pythonhosted.org/packages/92/3e/7648441cf72258fc2e5cb9be6f39890684f15a06f66db42f12a8041e4e30/voyager-2.0.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a2ea7f5e61bb7de06d0196b1fbd43826908c6f2811576c94d97e353e7b32cb78", size = 358415 },
|
|
1206
|
+
{ url = "https://files.pythonhosted.org/packages/42/02/0f7c5e4d9eaf83d932ab0c89e93117166575163f670a6200f55a56dae268/voyager-2.0.9-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4c451b25d449ac6923f7424a3255995f97263ed694fa873119703da01a09f6d7", size = 4237185 },
|
|
1207
|
+
{ url = "https://files.pythonhosted.org/packages/38/7c/47598f5630b97cf7cf8f29ac390b8beab24adeb3f322d004244660b96e8f/voyager-2.0.9-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ecd1e690a69a044b988fd102676bb8584ed423af4b7aa9e69349203690263aad", size = 4304425 },
|
|
1208
|
+
{ url = "https://files.pythonhosted.org/packages/1a/95/fa3f090d8e1bdbd8401e9944080f3c12d1e8f1385980da4151f6ab62a783/voyager-2.0.9-cp311-cp311-win_amd64.whl", hash = "sha256:5e1c264f4202b551ca8cb5a33d456f451c598ea7b9680134177dc5443819f328", size = 200738 },
|
|
1209
|
+
{ url = "https://files.pythonhosted.org/packages/22/fa/455fbd3a709403dd98f910b1ba5bfd9bf38db3e718edd03c8758ecb3b750/voyager-2.0.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2d344617a1c97d6b317f13e801b3377a55f1e5b91786f9354dcde12b7945b3e0", size = 736030 },
|
|
1210
|
+
{ url = "https://files.pythonhosted.org/packages/d4/9f/ee29e306901be0ea521b464c3533bd95d49dbe08c54df64cd49bbb13918e/voyager-2.0.9-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d58682f2f8d9cfccb95d2278eeec01ed76f90636cf6123db7850eb9a4b7e70e7", size = 380145 },
|
|
1211
|
+
{ url = "https://files.pythonhosted.org/packages/e5/83/d9b40742c17f674cbafec55b751804e68ee9b492cfb02fddc72f060beb17/voyager-2.0.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:510485dafa539ae130f2ad8bf7c95b21a2fa454f0b25581b37be3ca36bab29a1", size = 358094 },
|
|
1212
|
+
{ url = "https://files.pythonhosted.org/packages/33/61/e705ee3154450802912d35e881261182c01ba0e65cc83c6ef16b84c56e1b/voyager-2.0.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a4015db48c0c5d596e8bc536f4b05b601a2138e6439e00b06410e18286896273", size = 4234321 },
|
|
1213
|
+
{ url = "https://files.pythonhosted.org/packages/13/7a/4ac271fb5a8068f97dd6446622d679d26237a772ab12d1e6241cc0848dca/voyager-2.0.9-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:805f1736e5fe309ae8884e572b78d9816ff4fd6f1ef364298a2f62566ea86bc9", size = 4312789 },
|
|
1214
|
+
{ url = "https://files.pythonhosted.org/packages/b4/f6/906b21d4b42df264fc231c74753ca164f14cb09d40cd6b3a32d293aa0418/voyager-2.0.9-cp312-cp312-win_amd64.whl", hash = "sha256:abc7041ba8047d25942afe457db597be6f33161a669aac00dfffc00b7d833bfa", size = 201868 },
|
|
1215
|
+
{ url = "https://files.pythonhosted.org/packages/9f/9a/f5467d52356995c9f64cb27a99211142abead24290bfb8ccfb38168ef91d/voyager-2.0.9-cp39-cp39-macosx_10_13_universal2.whl", hash = "sha256:a38b7a232367095ef28350889b0c4ec8ce9f53c0a7839f2b06f517b652938e2c", size = 731403 },
|
|
1216
|
+
{ url = "https://files.pythonhosted.org/packages/ab/be/ecef9dfce07c535af6d11a4011a4a8693c4174ca8f87495e9ba30c8c7884/voyager-2.0.9-cp39-cp39-macosx_10_13_x86_64.whl", hash = "sha256:e5968f516a94832a599c9c1b64246b7f81344ed3fb49247fbdaba7e673c1e8a9", size = 377365 },
|
|
1217
|
+
{ url = "https://files.pythonhosted.org/packages/55/99/d31f596e887d998c81671b8e345925f07ccc36dee7144dbf80f1c86fb187/voyager-2.0.9-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f4095d4012479fa53b21f8a65c7790e844c3634a8ead1fca04f69a4a34ebc362", size = 356976 },
|
|
1218
|
+
{ url = "https://files.pythonhosted.org/packages/6f/43/096fcfa3443d37d8b5e4163b49b0d91f0456022df59017a8ae25ae4482d9/voyager-2.0.9-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d55f85d36111e950d42e7d78accca84c582d5606ec1bae2d4b2da41ccc3e7f2e", size = 4209912 },
|
|
1219
|
+
{ url = "https://files.pythonhosted.org/packages/ff/51/c35a343471746e8964cb5d002a3367f5e28f5787c2c0689e0566a7406f66/voyager-2.0.9-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83e1dd648040df093ddda74a6f361ec912caf981d20dfdd4b0d3bd7fb230147b", size = 4280035 },
|
|
1220
|
+
{ url = "https://files.pythonhosted.org/packages/ba/6c/5643effb7e495f8ee5c346e7b41961cb8d4dd6ef6952b24ec929a868f41c/voyager-2.0.9-cp39-cp39-win_amd64.whl", hash = "sha256:a2b17036c9b825f7062eb09ac613ca1e3fdb11712b1fc6d58bbbc201698e09c9", size = 196377 },
|
|
1221
|
+
]
|
|
1222
|
+
|
|
1183
1223
|
[[package]]
|
|
1184
1224
|
name = "wcwidth"
|
|
1185
1225
|
version = "0.2.13"
|
|
@@ -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
|
|