You can deploy this model and start earning money today!

KuanP/
counts_flash_hog_4gpu_a6_long_u2_lowilr_top75

Name: counts_flash_hog_4gpu_a6_long_u2_lowilr_top75
Availability: InStock

Views

No views yet

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


Steps	131,072
Global batch size	512
Learning rate	5e-05 (cosine, 500 warmup)
Weight decay	0.0001
Precision	bf16
Architecture	8L / d512 / 4-head, frozen 5120-d gene embeddings
Parameters	773.7M total, 28.9M trainable
Sequence length	2048 tokens
Data selection	`ExcludeDatasetSelector(dataset_id=53d208b0-2cfd-4366-9866-c3c6114081bc) + PrecomputedScoreSelector(scores_path=/scratch/m000137-pm06/kuan/data_rater_infer_output/counts_datarater_flash_hog_4gpu_a6_long_u2_lowilr/final/all_scores.npy, top_percentile=75)`
Source dataset	cellxgene_2025_top75_counts_flash_hog_4gpu_a6_long_u2_lowilr_exclude_ts_sparse_marlowe (62,634,100 cells / 61,888 genes)

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).

python

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:


`pad_length`	2048
`cls_token_idx`	1
`chrom_token_offset`	143574
`chrom_token_right_idx`	2
`pad_token_idx`	0
`vocab_size`	145469