| 44 | |
| 45 | |
| 46 | class SelectableBasicBackend(CosineBasicBackend): |
| 47 | def _selector_dist(self, x: npt.NDArray, selector: npt.NDArray[np.int_]) -> npt.NDArray: |
| 48 | """Compute cosine distance.""" |
| 49 | x_norm = normalize(x) |
| 50 | sim = x_norm.dot(self._vectors[selector].T) |
| 51 | return 1 - sim |
| 52 | |
| 53 | def query(self, vectors: npt.NDArray, k: int, selector: npt.NDArray[np.int_] | None = None) -> QueryResult: |
| 54 | """Batched distance query. |
| 55 | |
| 56 | :param vectors: The vectors to query. |
| 57 | :param k: The number of nearest neighbors to return. |
| 58 | :param selector: Optional array of chunk indices to filter results by. |
| 59 | :return: A list of tuples with the indices and distances. |
| 60 | :raises ValueError: If k is less than 1. |
| 61 | """ |
| 62 | if k < 1: |
| 63 | raise ValueError(f"k should be >= 1, is now {k}") |
| 64 | |
| 65 | out: QueryResult = [] |
| 66 | num_vectors = len(self.vectors) |
| 67 | effective_k = min(k, num_vectors) |
| 68 | if selector is not None: |
| 69 | effective_k = min(effective_k, len(selector)) |
| 70 | |
| 71 | # Batch the queries |
| 72 | for index in range(0, len(vectors), 1024): |
| 73 | batch = vectors[index : index + 1024] |
| 74 | if selector is not None: |
| 75 | distances = self._selector_dist(batch, selector) |
| 76 | else: |
| 77 | distances = self._dist(batch) |
| 78 | |
| 79 | # Efficiently get the k smallest distances |
| 80 | indices = np.argpartition(distances, kth=effective_k - 1, axis=1)[:, :effective_k] |
| 81 | sorted_indices = np.take_along_axis( |
| 82 | indices, np.argsort(np.take_along_axis(distances, indices, axis=1)), axis=1 |
| 83 | ) |
| 84 | sorted_distances = np.take_along_axis(distances, sorted_indices, axis=1) |
| 85 | |
| 86 | # Extend the output with tuples of (indices, distances) |
| 87 | if selector is not None: |
| 88 | sorted_indices = selector[sorted_indices] |
| 89 | out.extend(zip(sorted_indices, sorted_distances)) |
| 90 | |
| 91 | return out |
| 92 | |
| 93 | def save(self, path: Path) -> None: |
| 94 | """Save the selectable basic backend.""" |
| 95 | path.mkdir(parents=True, exist_ok=True) |
| 96 | super().save(path) |
| 97 | |
| 98 | @classmethod |
| 99 | def load(cls, path: Path) -> "SelectableBasicBackend": |
| 100 | """Load a selectable basic backend.""" |
| 101 | loaded = super().load(path) |
| 102 | return SelectableBasicBackend(loaded.vectors, loaded.arguments) |
no outgoing calls