Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,75 @@
+---
+license: other
+library_name: transformers
+tags: [single-cell, cell-embedding, uce, transcriptomics, scrna-seq]
+---
+
+# counts_flash_hog_4gpu_a6_long_u2_lowilr_bot75
+
+Counts DataRater (a6_long_u2_lowilr), **bottom** 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, bottom_percentile=75)` |
+| Source dataset | cellxgene_2025_bot75_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](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/counts_flash_hog_4gpu_a6_long_u2_lowilr_bot75")           # 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 |