File size: 2,742 Bytes
ad5806c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | ---
license: other
library_name: transformers
tags: [single-cell, cell-embedding, uce, transcriptomics, scrna-seq]
---
# cxg-random75
Uninformed random 75% of samples (`SampleKeyFractionSelector(fraction=0.75, seed=42)`).
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) + SampleKeyFractionSelector(fraction=0.75, seed=42)` |
| Source dataset | cellxgene_2025_exclude_ts_sparse_random75pct (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](https://github.com/) (`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/cxg-random75") # 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 |
|