| """Frozen Cell2Sentence (C2S-Scale-Gemma-2-2B) spot embeddings for the RNA sets of rna_sets.py (spot = pseudo-cell). |
| |
| python 43_rnafm_c2s.py [set ...] (default: all of rna_sets.SETS) |
| |
| Model : HF vandijklab/C2S-Scale-Gemma-2-2B @ 5ddf28b8f1c81b7ab7a9be192924da82b6c5d512 (cc-by-4.0), a Gemma-2 2B |
| decoder (26 layers, d=2304) instruction-tuned on cell sentences from >800 CellxGene/HCA datasets. Loaded as |
| a plain transformers Gemma2Model (no remote code, strict: no missing / unexpected keys), bf16, eager |
| attention (keeps Gemma-2's attention-logit soft-capping exactly as in training). |
| Tokenizer : the repo ships only the sentencepiece tokenizer.model; sentencepiece is not in the shared env, so it was |
| converted once to ext/cell2sentence/tokenizer/fast/tokenizer.json with transformers' own GemmaConverter |
| inside an isolated venv (ext/cell2sentence/venv) and checked token-identical to the sentencepiece |
| GemmaTokenizer on sampled prompts of every set. |
| Cell sentence (exactly as cell2sentence src/cell2sentence/utils.py generate_sentences + prompt_formatter.py, repo |
| vandijklab/cell2sentence @ 6599c9d, read-only clone in ext/cell2sentence/src): |
| raw counts -> sc.pp.normalize_total (target = median total counts of the set) -> log1p base 10 |
| (tutorial 0; rank-preserving within a spot) -> gene symbols upper-cased (generate_vocabulary) -> |
| nonzero genes of the spot, shuffled with np.random.seed(42) / sklearn.utils.shuffle and stable-sorted by |
| descending expression (= their random tie-breaking) -> top K=200 symbols (default n_genes=200 of |
| tasks.embed_cells / predict_cell_types_of_data and the tutorials), space-separated. |
| Not applied: tutorial-0 QC filters (filter_cells min_genes=200, filter_genes min_cells=3), because they |
| would drop spots / make a spot's sentence depend on the other spots of the set. |
| Prompt : the cell-type-prediction template used by tasks.embed_cells (C2SPromptFormatter task |
| "cell_type_prediction"); that formatter picks one of 15 paraphrases at random, we fix the first one, which |
| is also the template on the C2S-Scale model card: |
| "The following is a list of {K} gene names ordered by descending expression level in a Homo sapiens cell. |
| Your task is to give the cell type which this cell belongs to based on its gene expression.\\n |
| Cell sentence: {sentence}.\\nThe cell type corresponding to these genes is:" (+ <bos>, no chat template) |
| {K} is always 200, as in the formatter (spots with < 200 expressed genes keep "200"). |
| Readout : last hidden layer (post final RMSNorm = hidden_states[-1], as CSModel.embed_cells_batched). |
| "c2s" : mean over the cell-sentence tokens only (tokens whose char span overlaps the gene list; |
| official code averages over the whole padded prompt, which mixes in the constant template |
| and padding). |
| "c2s-last" : hidden state of the last prompt token (":" after "is"), i.e. the state that would generate |
| the cell type. |
| Output : cache/rnafm/c2s_<set>.npy and c2s-last_<set>.npy, (n_spot, 2304) float16, set order. Log logs/rnafm/c2s.log. |
| """ |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import scipy.sparse as sp |
| import torch |
| from sklearn.utils import shuffle |
|
|
| from common import CACHE, ROOT |
| from rna_sets import SETS, load, save |
|
|
| REPO = "vandijklab/C2S-Scale-Gemma-2-2B" |
| REV = "5ddf28b8f1c81b7ab7a9be192924da82b6c5d512" |
| TOK_DIR = os.path.join(ROOT, "ext", "cell2sentence", "tokenizer", "fast") |
| K = 200 |
| ORGANISM = "Homo sapiens" |
| PREFIX = (f"The following is a list of {K} gene names ordered by descending expression level in a {ORGANISM} cell. " |
| "Your task is to give the cell type which this cell belongs to based on its gene expression.\nCell sentence: ") |
| SUFFIX = ".\nThe cell type corresponding to these genes is:" |
| GPU_GB = 8.0 |
| TOK_BUDGET = 8 * 800 |
| dev = "cuda" |
|
|
|
|
| def log(*a): |
| print(*a, flush=True) |
|
|
|
|
| def cell_sentences(X, genes): |
| """C2S normalisation + rank -> list of top-K space-separated (upper-case) gene symbols, plus #expressed genes.""" |
| X = X.tocsr().astype(np.float64) |
| tot = np.asarray(X.sum(1)).ravel() |
| target = np.median(tot[tot > 0]) |
| Xn = sp.diags(target / np.maximum(tot, 1e-12)) @ X |
| Xn = Xn.tocsr() |
| Xn.data = np.log1p(Xn.data) / np.log(10) |
| enc = np.char.upper(np.asarray(genes, dtype=str)) |
| np.random.seed(42) |
| sents, n_expr = [], np.zeros(Xn.shape[0], int) |
| for i in range(Xn.shape[0]): |
| cols = Xn.indices[Xn.indptr[i]:Xn.indptr[i + 1]] |
| vals = Xn.data[Xn.indptr[i]:Xn.indptr[i + 1]] |
| cols, vals = shuffle(cols, vals) |
| order = cols[np.argsort(-vals, kind="stable")] |
| n_expr[i] = len(order) |
| sents.append(" ".join(enc[order[:K]])) |
| return sents, n_expr |
|
|
|
|
| def tokenize(tok, sents): |
| """Prompt ids + mask of cell-sentence tokens (by char offsets).""" |
| prompts = [PREFIX + s + SUFFIX for s in sents] |
| enc = tok(prompts, return_offsets_mapping=True, add_special_tokens=True) |
| out = [] |
| for p, s, ids, off in zip(prompts, sents, enc["input_ids"], enc["offset_mapping"]): |
| a, b = len(PREFIX), len(PREFIX) + len(s) |
| off = np.asarray(off) |
| m = (off[:, 0] < b) & (off[:, 1] > a) & (off[:, 1] > off[:, 0]) |
| assert ids[0] == tok.bos_token_id and m.sum() > 0 |
| out.append((np.asarray(ids), m)) |
| return prompts, out |
|
|
|
|
| @torch.no_grad() |
| def embed(model, seqs, pad_id): |
| n, D = len(seqs), model.config.hidden_size |
| E_mean, E_last = np.zeros((n, D), np.float32), np.zeros((n, D), np.float32) |
| order = np.argsort([len(s[0]) for s in seqs])[::-1] |
| budget, k, t0, done = TOK_BUDGET, 0, time.time(), 0 |
| while k < n: |
| Lm = len(seqs[order[k]][0]) |
| bs = max(1, budget // Lm) |
| idx = order[k:k + bs] |
| ids = torch.full((len(idx), Lm), pad_id, dtype=torch.long) |
| msk = torch.zeros((len(idx), Lm), dtype=torch.bool) |
| att = torch.zeros((len(idx), Lm), dtype=torch.long) |
| for j, i in enumerate(idx): |
| s, m = seqs[i] |
| ids[j, :len(s)], msk[j, :len(s)], att[j, :len(s)] = torch.from_numpy(s), torch.from_numpy(m), 1 |
| try: |
| h = model(input_ids=ids.to(dev), attention_mask=att.to(dev)).last_hidden_state.float() |
| except torch.OutOfMemoryError: |
| torch.cuda.empty_cache() |
| if budget <= Lm: |
| raise |
| budget //= 2 |
| log(f" OOM -> token budget {budget}") |
| continue |
| mk = msk.to(dev).float() |
| E_mean[idx] = ((h * mk[..., None]).sum(1) / mk.sum(1, keepdim=True)).cpu().numpy() |
| last = att.sum(1).to(dev) - 1 |
| E_last[idx] = h[torch.arange(len(idx), device=dev), last].cpu().numpy() |
| k += len(idx) |
| if k // 1000 > done: |
| done = k // 1000 |
| log(f" {k}/{n} {time.time() - t0:.0f}s bs={len(idx)} L={Lm} peak {torch.cuda.max_memory_allocated() / 2**30:.2f} GB") |
| return E_mean, E_last |
|
|
|
|
| def sanity(E, tag, organ=None): |
| rng = np.random.default_rng(0) |
| Z = E / np.linalg.norm(E, axis=1, keepdims=True) |
| s = Z[rng.choice(len(Z), min(200, len(Z)), replace=False)] |
| C = s @ s.T |
| mc = C[~np.eye(len(s), dtype=bool)].mean() |
| msg = f"[{tag}] finite={np.isfinite(E).all()} mean pairwise cos(200)={mc:.4f} std/dim={E.std(0).mean():.4f}" |
| if organ is not None: |
| Zc = Z - Z.mean(0) |
| Zc /= np.linalg.norm(Zc, axis=1, keepdims=True) |
| o = np.asarray(organ) |
| same, raw = [], [] |
| for ZZ, acc in ((Zc, same), (Z, raw)): |
| Zt = torch.from_numpy(ZZ).to(dev) |
| for i in range(0, len(Zt), 2048): |
| S = Zt[i:i + 2048] @ Zt.T |
| S[torch.arange(S.shape[0]), torch.arange(i, i + S.shape[0])] = -2 |
| nn = S.topk(10, dim=1).indices.cpu().numpy() |
| acc.append((o[nn] == o[i:i + S.shape[0], None]).mean(1)) |
| cnt = pd.Series(o).value_counts() |
| base = ((cnt * (cnt - 1)).sum()) / (len(o) * (len(o) - 1)) |
| msg += (f" 10-NN same-organ: {np.concatenate(raw).mean():.3f} (cos), {np.concatenate(same).mean():.3f}" |
| f" (centred cos) base rate {base:.3f}") |
| log(msg) |
|
|
|
|
| def main(): |
| from transformers import Gemma2Model, PreTrainedTokenizerFast |
| from huggingface_hub import snapshot_download |
| sets = sys.argv[1:] or SETS |
| torch.cuda.set_per_process_memory_fraction(min(1.0, GPU_GB * 2**30 / torch.cuda.get_device_properties(0).total_memory)) |
| tok = PreTrainedTokenizerFast.from_pretrained(TOK_DIR) |
| tok.padding_side = "right" |
| mdir = snapshot_download(REPO, revision=REV) |
| model, info = Gemma2Model.from_pretrained(mdir, dtype=torch.bfloat16, attn_implementation="eager", |
| output_loading_info=True) |
| log(f"loaded {REPO}@{REV}: missing={info['missing_keys']} unexpected={info['unexpected_keys']} " |
| f"mismatched={info['mismatched_keys']}") |
| assert not info["missing_keys"] and not info["mismatched_keys"] and not info["unexpected_keys"] |
| model.eval().to(dev) |
| log(f"weights {torch.cuda.memory_allocated() / 2**30:.2f} GB; prompt prefix tokens {len(tok(PREFIX)['input_ids'])}") |
| meta = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")) |
| tot_seq, tot_t = 0, 0.0 |
| for name in sets: |
| t0 = time.time() |
| X, genes = load(name) |
| sents, n_expr = cell_sentences(X, genes) |
| prompts, seqs = tokenize(tok, sents) |
| ntok = np.array([len(s[0]) for s in seqs]) |
| nsent = np.array([m.sum() for _, m in seqs]) |
| upg = np.char.upper(genes.astype(str)) |
| nonhgnc = np.mean([("." in g) or g.startswith("ENSG") for g in upg]) |
| log(f"== {name}: {X.shape[0]} spots x {X.shape[1]} genes (all symbols usable as text; {nonhgnc:.1%} non-HGNC-" |
| f"style ids e.g. ENSG / RP11-..), {len(set(upg))} unique upper-case symbols; expressed genes/spot median " |
| f"{np.median(n_expr):.0f}, spots with < {K} expressed genes {np.mean(n_expr < K):.1%}; sentence genes/spot " |
| f"median {np.median(np.minimum(n_expr, K)):.0f}; tokens/prompt median {np.median(ntok):.0f} " |
| f"(sentence tokens {np.median(nsent):.0f}, {np.median(nsent) / K:.2f} tok/gene)") |
| log(" example: " + prompts[0][:400].replace("\n", " | ") + " ...") |
| with open(os.path.join(ROOT, "ext", "cell2sentence", f"prompts_sample_{name}.txt"), "w") as f: |
| for i in np.random.default_rng(1).choice(len(prompts), min(200, len(prompts)), replace=False): |
| f.write(repr(prompts[i]) + "\n") |
| t1 = time.time() |
| E_mean, E_last = embed(model, seqs, tok.pad_token_id) |
| dt = time.time() - t1 |
| tot_seq, tot_t = tot_seq + len(seqs), tot_t + dt |
| log(f" forward {dt:.0f}s = {len(seqs) / dt:.2f} seq/s, {ntok.sum() / dt:.0f} tok/s; total {time.time() - t0:.0f}s;" |
| f" peak GPU {torch.cuda.max_memory_allocated() / 2**30:.2f} GB") |
| save("c2s", name, E_mean) |
| save("c2s-last", name, E_last) |
| organ = meta.organ.values if name in ("hest", "hest_v") else None |
| if organ is not None: |
| assert len(organ) == len(E_mean) |
| sanity(E_mean, f"c2s {name}", organ) |
| sanity(E_last, f"c2s-last {name}", organ) |
| log(f"overall throughput {tot_seq / max(tot_t, 1e-9):.2f} seq/s over {tot_seq} sequences ({tot_t / 60:.1f} min)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
