Embeddings search experimental API #1164
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1164 +/- ##
==========================================
+ Coverage 91.11% 91.29% +0.17%
==========================================
Files 77 79 +2
Lines 5923 6044 +121
==========================================
+ Hits 5397 5518 +121
Misses 526 526
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
mlin
force-pushed
the
mlin/similarity-search-api
branch
from
6bf1181
to
639e64c
Compare
api/python/cellxgene_census/src/cellxgene_census/experimental/_embedding_search.py
Outdated
Show resolved
Hide resolved
| """ | ||
|
|
||
|
|
||
| def find_nearest_obs( |
There was a problem hiding this comment.
On the API side, it would be nice if this could produce output that can be directly with sklearn style classes. For example, if this returned a KNNTransformer subclass, that could be used directly with the KNeighborsClassifier and KNeighborsRegressor classes.
There was a problem hiding this comment.
@ivirshup I like this idea very much, but I'm not quite sure it's workable (albeit I'm not as familiar with those APIs)...
Those scikit-learn classes seem oriented around the scenario where you're providing either all the points (in the "universe") or the complete distance matrix for them. Here we're working with a more limited view of the query points and their neighbor distances; we don't have or want the complete distance matrix, and actually we don't even have the coordinates of the neighbors immediately handy.
Do you think the shoe fits? I see there's some stuff about the "K neighbors graph" that might be relevant, but I'm not personally familiar enough to use them in an unconventional/advanced way like this.
|
|
||
|
|
||
| def find_nearest_obs( | ||
| embedding_metadata: Dict[str, Any], |
There was a problem hiding this comment.
Why does this use a different way to specify an embedding than get_embedding does?
There was a problem hiding this comment.
get_embedding is a little low level in that it wants the full URI to the embeddings TileDB array, which isn't actually needed to find the index. embedding_metadata is the information returned from get_embedding_metadata() which seems like the appropriate level (especially in view of #1181 wherein we will actually put the relative URIs to the index arrays in there), although of course it'd be nice if it were more typesafe.
There was a problem hiding this comment.
@mlin can we change it to?
embedding_name: str,
organism: str,
census_version: str
I think that provides an easier entry point to users and it aligns to get_embedding_metadata_by_name
api/python/cellxgene_census/src/cellxgene_census/experimental/_embedding_search.py
Show resolved
Hide resolved
|
|
||
|
|
||
| def find_nearest_obs( | ||
| embedding_metadata: Dict[str, Any], |
There was a problem hiding this comment.
@mlin can we change it to?
embedding_name: str,
organism: str,
census_version: str
I think that provides an easier entry point to users and it aligns to get_embedding_metadata_by_name
| def _resolve_embedding_index( | ||
| embedding_metadata: Dict[str, Any], | ||
| mirror: Optional[str] = None, | ||
| ) -> Optional[Tuple[str, str]]: |
There was a problem hiding this comment.
@ebezzi new index resolution method here
|
@ebezzi @pablo-gar @ivirshup Updated this to resolve indexes through mirrors/contributions json and remove the need for caller to use |
Adds two new functions to
cellxgene_census.experimental:find_nearest_obsuses TileDB-Vector-Search indexes of Census embeddings to find nearest neighbors of given embedding vectors (in an AnnData obsm layer). Census cell similarity search: experimental Python API for searching given AnnData #1114predict_obs_metadatauses the nearest neighbors to predict metadata attributes like cell_type and tissue_general for the query cells. Naive initial implementation is just a starting point to start experimenting with. Census cell similarity search: experimental Python API for metadata prediction #1115The TileDB-Vector-Search query speed seems to be very S3-latency-sensitive, even moreso than typical Census queries. It's many times faster to run from within AWS us-west-2 than externally.