scores.KNNScore
Abstract base class for KNN distance-based uncertainty scores.
Usage
scores.KNNScore(
k=1,
stat="max",
pca=None,
save_index=False,
)Computes distance-based uncertainty scores where low scores indicate samples similar to the training distribution (low uncertainty) and high scores indicate samples deviating from the training distribution (high uncertainty).
Parameters
k: int = 1-
Number of nearest neighbors used to compute the distance score.
stat: (max, mean, median, min) = "max"-
Statistic applied to aggregate distances across the k neighbors.
pca: TensorPCA or None = None-
Optional TensorPCA object for dimensionality reduction prior to scoring.
save_index: bool or Path = False-
If
True, the HNSW index is saved to a default file. If aPathis provided (must end in.bin), the index is saved there.
See Also
Methods
| Name | Description |
|---|---|
| knn_search() | Compute the K-nearest-neighbour distances and indices for a set of query embeddings. |
knn_search()
Compute the K-nearest-neighbour distances and indices for a set of query embeddings.
Usage
knn_search(query, offset=0)Parameters
query: torch.Tensor-
A 2-D tensor of shape
(N, D)containing the embeddings for which distances are to be computed. offset: int = 0-
Number of nearest neighbours to discard from the result. This is typically used to skip self-matching when the query points are drawn from the same set that built the index (e.g.
offset=1).
Returns
distances: torch.Tensor-
A tensor of shape
(N, k)containing the KNN distances for each query point after discarding the firstoffsetnearest neighbours. indices: torch.Tensor-
A tensor of shape
(N, k)with the index positions of the nearest neighbours in the reference embedding set after discarding the firstoffsetmatches.
Notes
offsetis useful when the query set is identical to the reference set, because the nearest neighbour would be the point itself (distance zero). Skipping it yields a meaningful distance to the second nearest neighbour.