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