vicinity 0.2.1__tar.gz → 0.3.0__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.
Files changed (34) hide show
  1. {vicinity-0.2.1 → vicinity-0.3.0}/PKG-INFO +75 -97
  2. {vicinity-0.2.1 → vicinity-0.3.0}/README.md +65 -96
  3. {vicinity-0.2.1 → vicinity-0.3.0}/pyproject.toml +10 -0
  4. {vicinity-0.2.1 → vicinity-0.3.0}/tests/test_vicinity.py +20 -0
  5. vicinity-0.3.0/vicinity/__init__.py +8 -0
  6. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/annoy.py +28 -24
  7. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/base.py +7 -0
  8. vicinity-0.3.0/vicinity/backends/basic.py +218 -0
  9. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/faiss.py +33 -51
  10. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/hnsw.py +19 -6
  11. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/pynndescent.py +21 -12
  12. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/usearch.py +26 -22
  13. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/utils.py +39 -0
  14. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/version.py +1 -1
  15. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/vicinity.py +78 -2
  16. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity.egg-info/PKG-INFO +75 -97
  17. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity.egg-info/requires.txt +10 -0
  18. vicinity-0.2.1/vicinity/__init__.py +0 -7
  19. vicinity-0.2.1/vicinity/backends/basic.py +0 -149
  20. {vicinity-0.2.1 → vicinity-0.3.0}/.github/workflows/ci.yaml +0 -0
  21. {vicinity-0.2.1 → vicinity-0.3.0}/.gitignore +0 -0
  22. {vicinity-0.2.1 → vicinity-0.3.0}/.pre-commit-config.yaml +0 -0
  23. {vicinity-0.2.1 → vicinity-0.3.0}/LICENSE +0 -0
  24. {vicinity-0.2.1 → vicinity-0.3.0}/Makefile +0 -0
  25. {vicinity-0.2.1 → vicinity-0.3.0}/setup.cfg +0 -0
  26. {vicinity-0.2.1 → vicinity-0.3.0}/tests/conftest.py +0 -0
  27. {vicinity-0.2.1 → vicinity-0.3.0}/tests/test_utils.py +0 -0
  28. {vicinity-0.2.1 → vicinity-0.3.0}/uv.lock +0 -0
  29. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/backends/__init__.py +0 -0
  30. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/datatypes.py +0 -0
  31. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity/py.typed +0 -0
  32. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity.egg-info/SOURCES.txt +0 -0
  33. {vicinity-0.2.1 → vicinity-0.3.0}/vicinity.egg-info/dependency_links.txt +0 -0
  34. {vicinity-0.2.1 → vicinity-0.3.0}/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.2.1
3
+ Version: 0.3.0
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,10 +67,19 @@ 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: all
71
+ Requires-Dist: hnswlib; extra == "all"
72
+ Requires-Dist: pynndescent>=0.5.10; extra == "all"
73
+ Requires-Dist: numba>=0.59.0; extra == "all"
74
+ Requires-Dist: llvmlite>=0.42.0; extra == "all"
75
+ Requires-Dist: numpy>=1.24.0; extra == "all"
76
+ Requires-Dist: annoy; extra == "all"
77
+ Requires-Dist: faiss-cpu; extra == "all"
78
+ Requires-Dist: usearch; extra == "all"
70
79
 
71
80
  <div align="center">
72
81
 
73
- # Vicinity: The Lightweight Vector Store
82
+ # Vicinity: Lightweight Nearest Neighbors
74
83
 
75
84
  </div>
76
85
 
@@ -87,18 +96,21 @@ Requires-Dist: usearch; extra == "usearch"
87
96
  </a>
88
97
  <a href="https://github.com/MinishLab/vicinity/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-MIT-green" alt="License - MIT"></a>
89
98
  </h2>
99
+
100
+ [Quickstart](#quickstart) •
101
+ [Main Features](#main-features) •
102
+ [Supported Backends](#supported-backends) •
103
+ [Installation](#installation)
104
+
90
105
  </div>
91
106
 
92
107
 
93
- ## Table of contents
108
+ Vicinity is a light-weight, low-dependency vector store. It provides a simple and intuitive interface for nearest neighbor search, with support for different backends and evaluation.
109
+
110
+ There are many nearest neighbors packages and methods out there. However, we found it difficult to compare them. Every package has its own interface, quirks, and limitations, and learning a new package can be time-consuming. In addition to that, how do you effectively evaluate different packages? How do you know which one is the best for your use case?
94
111
 
95
- - [Quickstart](#quickstart)
96
- - [Main Features](#main-features)
97
- - [Supported Backends](#supported-backends)
98
- - [Backend Parameters](#backend-parameters)
99
- - [Usage](#usage)
100
112
 
101
- Vicinity is the lightest-weight vector store. Just put in some vectors, calculate query vectors, and off you go. It provides a simple and intuitive API for nearest neighbor search, with support for different backends.
113
+ This is where Vicinity comes in. Instead of learning a new interface for each new package or backend, Vicinity provides a unified interface for all backends. This allows you to easily experiment with different indexing methods and distance metrics and choose the best one for your use case. Vicinity also provides a simple way to evaluate the performance of different backends, allowing you to measure the queries per second and recall.
102
114
 
103
115
  ## Quickstart
104
116
 
@@ -106,50 +118,81 @@ Install the package with:
106
118
  ```bash
107
119
  pip install vicinity
108
120
  ```
121
+ Optinally, [install any of the supported backends](#installation), or simply install all of them with:
122
+ ```bash
123
+ pip install vicinity[all]
124
+ ```
109
125
 
110
126
 
111
127
  The following code snippet demonstrates how to use Vicinity for nearest neighbor search:
112
128
  ```python
113
129
  import numpy as np
114
130
  from vicinity import Vicinity
115
- from vicinity.datatypes import Backend
131
+ from vicinity.datatypes import Backend, Metric
116
132
 
117
133
  # Create some dummy data
118
134
  items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
119
135
  vectors = np.random.rand(len(items), 128)
120
136
 
121
- # Initialize the Vicinity instance (using the basic backend)
122
- vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC)
137
+ # Initialize the Vicinity instance (using the basic backend and cosine metric)
138
+ vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC, metric=Metric.COSINE)
123
139
 
124
- # Query for nearest neighbors with a top-k search
140
+ # Create a query vector
125
141
  query_vector = np.random.rand(128)
142
+
143
+ # Query for nearest neighbors with a top-k search
126
144
  results = vicinity.query([query_vector], k=3)
127
145
 
128
146
  # Query for nearest neighbors with a threshold search
129
147
  results = vicinity.query_threshold([query_vector], threshold=0.9)
148
+ ```
130
149
 
131
- # Save the vector store
150
+ Saving and loading a vector store:
151
+ ```python
132
152
  vicinity.save('my_vector_store')
133
-
134
- # Load the vector store
135
153
  vicinity = Vicinity.load('my_vector_store')
136
154
  ```
137
155
 
156
+ Evaluating a backend:
157
+ ```python
158
+ # Use the first 1000 vectors as query vectors
159
+ query_vectors = vectors[:1000]
160
+
161
+ # Evaluate the Vicinity instance by measuring the queries per second and recall
162
+ qps, recall = vicinity.evaluate(
163
+ full_vectors=vectors,
164
+ query_vectors=query_vectors,
165
+ )
166
+ ```
167
+
138
168
  ## Main Features
139
169
  Vicinity provides the following features:
140
170
  - Lightweight: Minimal dependencies and fast performance.
141
171
  - Flexible Backend Support: Use different backends for vector storage and search.
142
172
  - Serialization: Save and load vector stores for persistence.
173
+ - Evaluation: Easily evaluate the performance of different backends.
143
174
  - Easy to Use: Simple and intuitive API.
144
175
 
145
176
  ## Supported Backends
146
177
  The following backends are supported:
147
- - `BASIC`: A simple flat index for vector storage and search.
178
+ - `BASIC`: A simple (exact matching) flat index for vector storage and search.
148
179
  - [HNSW](https://github.com/nmslib/hnswlib): Hierarchical Navigable Small World Graph (HNSW) for ANN search using hnswlib.
149
- - [FAISS](https://github.com/facebookresearch/faiss): ANN search using FAISS. All FAISS indexes are supported.
180
+ - [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
150
181
  - [ANNOY](https://github.com/spotify/annoy): "Approximate Nearest Neighbors Oh Yeah" for approximate nearest neighbor search.
151
182
  - [PYNNDescent](https://github.com/lmcinnes/pynndescent): ANN search using PyNNDescent.
152
- - [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
183
+ - [FAISS](https://github.com/facebookresearch/faiss): All FAISS indexes are supported:
184
+ - `flat`: Exact search.
185
+ - `ivf`: Inverted file search.
186
+ - `hnsw`: Hierarchical Navigable Small World Graph.
187
+ - `lsh`: Locality Sensitive Hashing.
188
+ - `scalar`: Scalar quantizer.
189
+ - `pq`: Product Quantizer.
190
+ - `ivf_scalar`: Inverted file search with scalar quantizer.
191
+ - `ivfpq`: Inverted file search with product quantizer.
192
+ - `ivfpqr`: Inverted file search with product quantizer and refinement.
193
+
194
+
195
+
153
196
 
154
197
  NOTE: the ANN backends do not support dynamic deletion. To delete items, you need to recreate the index. Insertion is supported in the following backends: `FAISS`, `HNSW`, and `Usearch`. The `BASIC` backend supports both insertion and deletion.
155
198
 
@@ -178,87 +221,22 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
178
221
  | | `expansion_search` | Number of candidates considered during search. | `64` |
179
222
 
180
223
 
224
+ ## Installation
225
+ The following installation options are available:
226
+ ```bash
227
+ # Install the base package
228
+ pip install vicinity
181
229
 
182
- ## Usage
183
-
184
- <details>
185
- <summary> Creating a Vector Store
186
- </summary>
187
- <br>
188
-
189
- You can create a Vicinity instance by providing items and their corresponding vectors:
190
-
191
-
192
- ```python
193
- from vicinity import Vicinity
194
- import numpy as np
195
-
196
- items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
197
- vectors = np.random.rand(len(items), 128)
198
-
199
- vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items)
200
- ```
201
-
202
- </details>
203
-
204
- <details>
205
- <summary> Querying
206
- </summary>
207
- <br>
208
-
209
- Find the k nearest neighbors for a given vector:
210
-
211
- ```python
212
- query_vector = np.random.rand(128)
213
- results = vicinity.query([query_vector], k=3)
214
- ```
215
-
216
- Find all neighbors within a given threshold:
217
-
218
- ```python
219
- query_vector = np.random.rand(128)
220
- results = vicinity.query_threshold([query_vector], threshold=0.9)
221
- ```
222
- </details>
223
-
224
- <details>
225
-
226
- <summary> Inserting and Deleting Items
227
- </summary>
228
- <br>
229
-
230
- Insert new items:
231
-
232
- ```python
233
- new_items = ["ocarina", "bow"]
234
- new_vectors = np.random.rand(2, 128)
235
- vicinity.insert(new_items, new_vectors)
236
- ```
237
-
238
- Delete items:
239
-
240
- ```python
241
- vicinity.delete(["hookshot"])
242
- ```
243
- </details>
244
-
245
- <details>
246
- <summary> Saving and Loading
247
- </summary>
248
- <br>
249
-
250
- Save the vector store:
251
-
252
- ```python
253
- vicinity.save('my_vector_store')
254
- ```
255
-
256
- Load the vector store:
230
+ # Install all backends
231
+ pip install vicinity[all]
257
232
 
258
- ```python
259
- vicinity = Vicinity.load('my_vector_store')
233
+ # Install specific backends
234
+ pip install vicinity[annoy]
235
+ pip install vicinity[faiss]
236
+ pip install vicinity[hnsw]
237
+ pip install vicinity[pynndescent]
238
+ pip install vicinity[usearch]
260
239
  ```
261
- </details>
262
240
 
263
241
  ## License
264
242
 
@@ -1,6 +1,6 @@
1
1
  <div align="center">
2
2
 
3
- # Vicinity: The Lightweight Vector Store
3
+ # Vicinity: Lightweight Nearest Neighbors
4
4
 
5
5
  </div>
6
6
 
@@ -17,18 +17,21 @@
17
17
  </a>
18
18
  <a href="https://github.com/MinishLab/vicinity/blob/main/LICENSE"><img src="https://img.shields.io/badge/license-MIT-green" alt="License - MIT"></a>
19
19
  </h2>
20
+
21
+ [Quickstart](#quickstart) •
22
+ [Main Features](#main-features) •
23
+ [Supported Backends](#supported-backends) •
24
+ [Installation](#installation)
25
+
20
26
  </div>
21
27
 
22
28
 
23
- ## Table of contents
29
+ Vicinity is a light-weight, low-dependency vector store. It provides a simple and intuitive interface for nearest neighbor search, with support for different backends and evaluation.
30
+
31
+ There are many nearest neighbors packages and methods out there. However, we found it difficult to compare them. Every package has its own interface, quirks, and limitations, and learning a new package can be time-consuming. In addition to that, how do you effectively evaluate different packages? How do you know which one is the best for your use case?
24
32
 
25
- - [Quickstart](#quickstart)
26
- - [Main Features](#main-features)
27
- - [Supported Backends](#supported-backends)
28
- - [Backend Parameters](#backend-parameters)
29
- - [Usage](#usage)
30
33
 
31
- Vicinity is the lightest-weight vector store. Just put in some vectors, calculate query vectors, and off you go. It provides a simple and intuitive API for nearest neighbor search, with support for different backends.
34
+ This is where Vicinity comes in. Instead of learning a new interface for each new package or backend, Vicinity provides a unified interface for all backends. This allows you to easily experiment with different indexing methods and distance metrics and choose the best one for your use case. Vicinity also provides a simple way to evaluate the performance of different backends, allowing you to measure the queries per second and recall.
32
35
 
33
36
  ## Quickstart
34
37
 
@@ -36,50 +39,81 @@ Install the package with:
36
39
  ```bash
37
40
  pip install vicinity
38
41
  ```
42
+ Optinally, [install any of the supported backends](#installation), or simply install all of them with:
43
+ ```bash
44
+ pip install vicinity[all]
45
+ ```
39
46
 
40
47
 
41
48
  The following code snippet demonstrates how to use Vicinity for nearest neighbor search:
42
49
  ```python
43
50
  import numpy as np
44
51
  from vicinity import Vicinity
45
- from vicinity.datatypes import Backend
52
+ from vicinity.datatypes import Backend, Metric
46
53
 
47
54
  # Create some dummy data
48
55
  items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
49
56
  vectors = np.random.rand(len(items), 128)
50
57
 
51
- # Initialize the Vicinity instance (using the basic backend)
52
- vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC)
58
+ # Initialize the Vicinity instance (using the basic backend and cosine metric)
59
+ vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC, metric=Metric.COSINE)
53
60
 
54
- # Query for nearest neighbors with a top-k search
61
+ # Create a query vector
55
62
  query_vector = np.random.rand(128)
63
+
64
+ # Query for nearest neighbors with a top-k search
56
65
  results = vicinity.query([query_vector], k=3)
57
66
 
58
67
  # Query for nearest neighbors with a threshold search
59
68
  results = vicinity.query_threshold([query_vector], threshold=0.9)
69
+ ```
60
70
 
61
- # Save the vector store
71
+ Saving and loading a vector store:
72
+ ```python
62
73
  vicinity.save('my_vector_store')
63
-
64
- # Load the vector store
65
74
  vicinity = Vicinity.load('my_vector_store')
66
75
  ```
67
76
 
77
+ Evaluating a backend:
78
+ ```python
79
+ # Use the first 1000 vectors as query vectors
80
+ query_vectors = vectors[:1000]
81
+
82
+ # Evaluate the Vicinity instance by measuring the queries per second and recall
83
+ qps, recall = vicinity.evaluate(
84
+ full_vectors=vectors,
85
+ query_vectors=query_vectors,
86
+ )
87
+ ```
88
+
68
89
  ## Main Features
69
90
  Vicinity provides the following features:
70
91
  - Lightweight: Minimal dependencies and fast performance.
71
92
  - Flexible Backend Support: Use different backends for vector storage and search.
72
93
  - Serialization: Save and load vector stores for persistence.
94
+ - Evaluation: Easily evaluate the performance of different backends.
73
95
  - Easy to Use: Simple and intuitive API.
74
96
 
75
97
  ## Supported Backends
76
98
  The following backends are supported:
77
- - `BASIC`: A simple flat index for vector storage and search.
99
+ - `BASIC`: A simple (exact matching) flat index for vector storage and search.
78
100
  - [HNSW](https://github.com/nmslib/hnswlib): Hierarchical Navigable Small World Graph (HNSW) for ANN search using hnswlib.
79
- - [FAISS](https://github.com/facebookresearch/faiss): ANN search using FAISS. All FAISS indexes are supported.
101
+ - [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
80
102
  - [ANNOY](https://github.com/spotify/annoy): "Approximate Nearest Neighbors Oh Yeah" for approximate nearest neighbor search.
81
103
  - [PYNNDescent](https://github.com/lmcinnes/pynndescent): ANN search using PyNNDescent.
82
- - [USEARCH](https://github.com/unum-cloud/usearch): ANN search using Usearch. This uses a highly optimized version of the HNSW algorithm.
104
+ - [FAISS](https://github.com/facebookresearch/faiss): All FAISS indexes are supported:
105
+ - `flat`: Exact search.
106
+ - `ivf`: Inverted file search.
107
+ - `hnsw`: Hierarchical Navigable Small World Graph.
108
+ - `lsh`: Locality Sensitive Hashing.
109
+ - `scalar`: Scalar quantizer.
110
+ - `pq`: Product Quantizer.
111
+ - `ivf_scalar`: Inverted file search with scalar quantizer.
112
+ - `ivfpq`: Inverted file search with product quantizer.
113
+ - `ivfpqr`: Inverted file search with product quantizer and refinement.
114
+
115
+
116
+
83
117
 
84
118
  NOTE: the ANN backends do not support dynamic deletion. To delete items, you need to recreate the index. Insertion is supported in the following backends: `FAISS`, `HNSW`, and `Usearch`. The `BASIC` backend supports both insertion and deletion.
85
119
 
@@ -108,87 +142,22 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
108
142
  | | `expansion_search` | Number of candidates considered during search. | `64` |
109
143
 
110
144
 
145
+ ## Installation
146
+ The following installation options are available:
147
+ ```bash
148
+ # Install the base package
149
+ pip install vicinity
111
150
 
112
- ## Usage
113
-
114
- <details>
115
- <summary> Creating a Vector Store
116
- </summary>
117
- <br>
118
-
119
- You can create a Vicinity instance by providing items and their corresponding vectors:
120
-
121
-
122
- ```python
123
- from vicinity import Vicinity
124
- import numpy as np
125
-
126
- items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
127
- vectors = np.random.rand(len(items), 128)
128
-
129
- vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items)
130
- ```
131
-
132
- </details>
133
-
134
- <details>
135
- <summary> Querying
136
- </summary>
137
- <br>
138
-
139
- Find the k nearest neighbors for a given vector:
140
-
141
- ```python
142
- query_vector = np.random.rand(128)
143
- results = vicinity.query([query_vector], k=3)
144
- ```
145
-
146
- Find all neighbors within a given threshold:
147
-
148
- ```python
149
- query_vector = np.random.rand(128)
150
- results = vicinity.query_threshold([query_vector], threshold=0.9)
151
- ```
152
- </details>
153
-
154
- <details>
155
-
156
- <summary> Inserting and Deleting Items
157
- </summary>
158
- <br>
159
-
160
- Insert new items:
161
-
162
- ```python
163
- new_items = ["ocarina", "bow"]
164
- new_vectors = np.random.rand(2, 128)
165
- vicinity.insert(new_items, new_vectors)
166
- ```
167
-
168
- Delete items:
169
-
170
- ```python
171
- vicinity.delete(["hookshot"])
172
- ```
173
- </details>
174
-
175
- <details>
176
- <summary> Saving and Loading
177
- </summary>
178
- <br>
179
-
180
- Save the vector store:
181
-
182
- ```python
183
- vicinity.save('my_vector_store')
184
- ```
185
-
186
- Load the vector store:
151
+ # Install all backends
152
+ pip install vicinity[all]
187
153
 
188
- ```python
189
- vicinity = Vicinity.load('my_vector_store')
154
+ # Install specific backends
155
+ pip install vicinity[annoy]
156
+ pip install vicinity[faiss]
157
+ pip install vicinity[hnsw]
158
+ pip install vicinity[pynndescent]
159
+ pip install vicinity[usearch]
190
160
  ```
191
- </details>
192
161
 
193
162
  ## License
194
163
 
@@ -52,6 +52,16 @@ pynndescent = [
52
52
  annoy = ["annoy"]
53
53
  faiss = ["faiss-cpu"]
54
54
  usearch = ["usearch"]
55
+ all = [
56
+ "hnswlib",
57
+ "pynndescent>=0.5.10",
58
+ "numba>=0.59.0",
59
+ "llvmlite>=0.42.0",
60
+ "numpy>=1.24.0",
61
+ "annoy",
62
+ "faiss-cpu",
63
+ "usearch"
64
+ ]
55
65
 
56
66
  [project.urls]
57
67
  "Homepage" = "https://github.com/MinishLab"
@@ -220,3 +220,23 @@ def test_vicinity_delete_and_query(vicinity_instance: Vicinity, items: list[str]
220
220
 
221
221
  # Check that the queried item is in the results
222
222
  assert "item3" in returned_items
223
+
224
+
225
+ def test_vicinity_evaluate(vicinity_instance: Vicinity, vectors: np.ndarray) -> None:
226
+ """
227
+ Test the evaluate method of the Vicinity instance.
228
+
229
+ :param vicinity_instance: A Vicinity instance.
230
+ :param vectors: The full dataset vectors used to build the index.
231
+ """
232
+ query_vectors = vectors[:10]
233
+ qps, recall = vicinity_instance.evaluate(vectors, query_vectors)
234
+
235
+ # Ensure the QPS and recall values are within valid ranges
236
+ assert qps > 0
237
+ assert 0 <= recall <= 1
238
+
239
+ # Test with an unsupported metric
240
+ vicinity_instance.backend.arguments.metric = "manhattan"
241
+ with pytest.raises(ValueError):
242
+ vicinity_instance.evaluate(vectors, query_vectors)
@@ -0,0 +1,8 @@
1
+ """Small vector store."""
2
+
3
+ from vicinity.datatypes import Backend
4
+ from vicinity.utils import Metric, normalize
5
+ from vicinity.version import __version__
6
+ from vicinity.vicinity import Vicinity
7
+
8
+ __all__ = ["Backend", "Metric", "Vicinity", "normalize", "__version__"]
@@ -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, Literal
5
+ from typing import Any, Union
6
6
 
7
7
  import numpy as np
8
8
  from annoy import AnnoyIndex
@@ -10,19 +10,25 @@ from numpy import typing as npt
10
10
 
11
11
  from vicinity.backends.base import AbstractBackend, BaseArgs
12
12
  from vicinity.datatypes import Backend, QueryResult
13
- from vicinity.utils import normalize
13
+ from vicinity.utils import Metric, normalize
14
14
 
15
15
 
16
16
  @dataclass
17
17
  class AnnoyArgs(BaseArgs):
18
18
  dim: int = 0
19
- metric: Literal["dot", "euclidean", "cosine"] = "cosine"
19
+ metric: str = "cosine"
20
20
  trees: int = 100
21
21
  length: int | None = None
22
22
 
23
23
 
24
24
  class AnnoyBackend(AbstractBackend[AnnoyArgs]):
25
25
  argument_class = AnnoyArgs
26
+ supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN, Metric.INNER_PRODUCT}
27
+ inverse_metric_mapping = {
28
+ Metric.COSINE: "dot",
29
+ Metric.EUCLIDEAN: "euclidean",
30
+ Metric.INNER_PRODUCT: "dot",
31
+ }
26
32
 
27
33
  def __init__(
28
34
  self,
@@ -40,25 +46,28 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
40
46
  def from_vectors(
41
47
  cls: type[AnnoyBackend],
42
48
  vectors: npt.NDArray,
43
- metric: Literal["dot", "euclidean", "cosine"],
49
+ metric: Union[str, Metric],
44
50
  trees: int,
45
51
  **kwargs: Any,
46
52
  ) -> AnnoyBackend:
47
53
  """Create a new instance from vectors."""
48
- dim = vectors.shape[1]
49
- actual_metric: Literal["dot", "euclidean"]
50
- if metric == "cosine":
51
- actual_metric = "dot"
54
+ metric_enum = Metric.from_string(metric)
55
+
56
+ if metric_enum not in cls.supported_metrics:
57
+ raise ValueError(f"Metric '{metric_enum.value}' is not supported by AnnoyBackend.")
58
+
59
+ metric = cls._map_metric_to_string(metric_enum)
60
+
61
+ if metric == "dot":
52
62
  vectors = normalize(vectors)
53
- else:
54
- actual_metric = metric
55
63
 
56
- index = AnnoyIndex(f=dim, metric=actual_metric)
64
+ dim = vectors.shape[1]
65
+ index = AnnoyIndex(f=dim, metric=metric) # type: ignore
57
66
  for i, vector in enumerate(vectors):
58
67
  index.add_item(i, vector)
59
68
  index.build(trees)
60
69
 
61
- arguments = AnnoyArgs(dim=dim, trees=trees, metric=metric, length=len(vectors))
70
+ arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors)) # type: ignore
62
71
  return AnnoyBackend(index, arguments=arguments)
63
72
 
64
73
  @property
@@ -80,11 +89,7 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
80
89
  """Load the vectors from a path."""
81
90
  path = Path(base_path) / "index.bin"
82
91
  arguments = AnnoyArgs.load(base_path / "arguments.json")
83
-
84
- metric = arguments.metric
85
- actual_metric = "dot" if metric == "cosine" else metric
86
-
87
- index = AnnoyIndex(arguments.dim, actual_metric)
92
+ index = AnnoyIndex(arguments.dim, arguments.metric) # type: ignore
88
93
  index.load(str(path))
89
94
 
90
95
  return cls(index, arguments=arguments)
@@ -93,7 +98,7 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
93
98
  """Save the vectors to a path."""
94
99
  path = Path(base_path) / "index.bin"
95
100
  self.index.save(str(path))
96
- # NOTE: set the length before saving.
101
+ # Ensure the length is set before saving
97
102
  self.arguments.length = len(self)
98
103
  self.arguments.dump(base_path / "arguments.json")
99
104
 
@@ -101,28 +106,27 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
101
106
  """Query the backend."""
102
107
  out = []
103
108
  for vec in vectors:
104
- if self.arguments.metric == "cosine":
109
+ if self.arguments.metric == "dot":
105
110
  vec = normalize(vec)
106
111
  indices, scores = self.index.get_nns_by_vector(vec, k, include_distances=True)
107
112
  scores_array = np.asarray(scores)
108
- if self.arguments.metric == "cosine":
109
- # Turn cosine similarity into cosine distance.
113
+ if self.arguments.metric == "dot":
114
+ # Convert cosine similarity to cosine distance
110
115
  scores_array = 1 - scores_array
111
116
  out.append((np.asarray(indices), scores_array))
112
117
  return out
113
118
 
114
119
  def insert(self, vectors: npt.NDArray) -> None:
115
120
  """Insert vectors into the backend."""
116
- raise NotImplementedError("Insertion is not supported in ANNOY backend.")
121
+ raise NotImplementedError("Insertion is not supported in Annoy backend.")
117
122
 
118
123
  def delete(self, indices: list[int]) -> None:
119
124
  """Delete vectors from the backend."""
120
- raise NotImplementedError("Deletion is not supported in ANNOY backend.")
125
+ raise NotImplementedError("Deletion is not supported in Annoy backend.")
121
126
 
122
127
  def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
123
128
  """Threshold the backend."""
124
129
  out: list[npt.NDArray] = []
125
130
  for x, y in self.query(vectors, 100):
126
131
  out.append(x[y < threshold])
127
-
128
132
  return out