snap-stanford/cellrater-metadata

CellRater (metadata)

The Metadata CellRater scores a cell from a sentence embedding of its metadata (cell_type, tissue_general, tissue, development_stage, disease, assay, suspension_type, sex; 256-dim text-embedding-3-small Matryoshka embeddings) plus, when use_numerical_features is set, log1p of the cell's nnz and raw_sum.

CellRater is a meta-learned per-cell data valuation model for single-cell RNA-seq: it is trained by bilevel optimisation so that its scores, used as per-cell loss weights, improve the held-out loss of a UCE-style proxy model. Scores are raw logits; higher means the cell is more useful to the proxy in the published setting, and the relative ranking within one dataset is what carries meaning. Format: cellrater-npz-v1, cellrater 0.1.0.

Inputs

• A sentence table (sentences.parquet) and row-aligned embeddings.npy covering your metadata combinations; combinations absent from the table are scored on the zero embedding and counted as unknown_combos.
• A corpus with the eight metadata columns (and nnz, raw_sum). cellrater-prepare writes one from an .h5ad; the corpus layout needs an ordered gene list even though this model never reads the counts, so pass any gene list that overlaps your data (the GEX model's genes.txt works).

Files

• model.npz: weights and constructor config (__config__)
• config.json: kind, input_embedding_dim and the checkpoint config
• README.md: this card

Usage

uv sync --extra anndata --extra hub
hf download snap-stanford/cellrater-metadata --local-dir models/cellrater-metadata
# The corpus layout needs an ordered gene list although this model never reads the counts:
# reuse the GEX model's genes.txt, or any gene list that overlaps your data.
hf download snap-stanford/cellrater-gex --local-dir models/cellrater-gex
uv run cellrater-prepare --h5ad my_cells.h5ad --genes models/cellrater-gex/genes.txt \
    --out my_cells.dataset
JAX_PLATFORMS=cpu uv run cellrater-score checkpoint=models/cellrater-metadata/model.npz \
    dataset_path=my_cells.dataset sentences_parquet_path=sentences.parquet \
    embeddings_npy_path=embeddings.npy output_dir=scores/my_cells

scores/my_cells/gold_metadata_scores.npy is a (2, N) array: row 0 the corpus row index (your adata.obs order), row 1 the score. See docs/tutorial_score_your_cells.md in the cellrater repository for the full walk-through, including how to write the scores back into adata.obs.