README.md

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