Query vectors within a distance threshold, using range_search if supported.
(self, vectors: npt.NDArray, threshold: float, max_k: int)
| 165 | raise NotImplementedError("Deletion is not supported in FAISS backends.") |
| 166 | |
| 167 | def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult: |
| 168 | """Query vectors within a distance threshold, using range_search if supported.""" |
| 169 | out: QueryResult = [] |
| 170 | if self.arguments.metric == "cosine": |
| 171 | vectors = normalize(vectors) |
| 172 | |
| 173 | if isinstance(self.index, RANGE_SEARCH_INDEXES): |
| 174 | radius = threshold |
| 175 | lims, D, I = self.index.range_search(vectors, radius) |
| 176 | for i in range(vectors.shape[0]): |
| 177 | start, end = lims[i], lims[i + 1] |
| 178 | idx = I[start:end] |
| 179 | dist = D[start:end] |
| 180 | if self.arguments.metric == "cosine": |
| 181 | dist = 1 - dist |
| 182 | mask = dist < threshold |
| 183 | out.append((idx[mask], dist[mask])) |
| 184 | else: |
| 185 | distances, indices = self.index.search(vectors, max_k) |
| 186 | for dist, idx in zip(distances, indices): |
| 187 | if self.arguments.metric == "cosine": |
| 188 | dist = 1 - dist |
| 189 | mask = dist < threshold |
| 190 | out.append((idx[mask], dist[mask])) |
| 191 | |
| 192 | return out |
| 193 | |
| 194 | def save(self, path: Path) -> None: |
| 195 | """Save the FAISS index and arguments.""" |
nothing calls this directly
no test coverage detected