counts_flash_hog_4gpu_a6_long_u2_lowilr_top75

counts_flash_hog_4gpu_a6_long_u2_lowilr_top75

Counts DataRater (a6_long_u2_lowilr), top 75% of cells by score.

UCE (Universal Cell Embedding) model trained on CELLxGENE 2025-01-30. This is the final model at 131,072 steps. One of a set of runs that share an identical training recipe and differ only in which training cells were selected, so their representations can be compared against each other.

Training recipe

The gene embedding table is a frozen nn.Parameter, so it lives in the state dict -- hence the ~3 GB model.safetensors. You do not need all_tokens.torch at inference.

Usage

Requires uce-training-suite (uce_suite).

from huggingface_hub import snapshot_download
from uce_suite.inference import (
    load_uce_checkpoint, load_cell_sentence_params, load_gene_artifacts, embed_dataset,
)

local = snapshot_download("KuanP/counts_flash_hog_4gpu_a6_long_u2_lowilr_top75")           # private repo: needs auth

model  = load_uce_checkpoint(local, device="cuda")
params = load_cell_sentence_params(local)             # reads config.yaml in this repo
gene_names, gene_mapping = load_gene_artifacts(
    f"{local}/gene_names.txt", f"{local}/all_species_gene_dict.json", species="human",
)

emb = embed_dataset(
    model, dataset_path="/path/to/cells.dataset",
    gene_names=gene_names, gene_mapping=gene_mapping, **params.as_kwargs(),
)

Use load_uce_checkpoint rather than a bare from_pretrained: on transformers >=5 from_pretrained re-runs _init_weights after loading and overwrites the trained nn.Linear weights. The loader force-reloads the state dict to undo that. The symptom if it is skipped is a per-cell loss pinned at log(2) = 0.693 and noise-like embeddings.

Tokenisation

config.yaml ships in this repo and is what load_cell_sentence_params reads. These values must match at inference or the embeddings are silently wrong: