| """Frozen Nicheformer spot embeddings of the 5 RNA sets (each Visium spot / virtual spot = one pseudo-cell). |
| |
| python 43_rnafm_nicheformer.py [set ...] (default: all of rna_sets.SETS) |
| |
| Model : theislab/Nicheformer (HF, MIT), revision 0ba4ba162f9ad3d6c8b713761954c31be2fa69ed, file model.safetensors |
| (12-layer post-norm nn.TransformerEncoder, d=512, 16 heads, FFN 1024, learnable PE, context 1500). |
| Downloaded to ext/nicheformer/hf/; the HF remote code (modeling_/tokenization_nicheformer.py) was read but is |
| NOT executed: the forward pass is re-implemented below (same nn.TransformerEncoderLayer stack) and loaded |
| with a strict state-dict load (the unused template `encoder_layer.*`, the MLM head and the `pos` buffer |
| are the only checkpoint keys skipped, and that set is asserted). |
| Technology means: ext/nicheformer/repo/data/model_means/dissociated_mean_script.npy from the official |
| GitHub repo theislab/nicheformer (commit 485cadbc5caa15119adfd54228f8a8af835fcabc); gene order = |
| model.h5ad var (20,310 Ensembl ids, token = column index + 30, == vocab.json). |
| Inputs : rna_sets.load(set) raw(-like) counts + HGNC symbols. Symbols -> Ensembl via cache/external/ |
| dlpfc_genes_full.csv (10x GRCh38 reference), fallback Geneformer gene_name_id_dict_gc104M.pkl; genes not in |
| the Nicheformer vocabulary are dropped (several symbols hitting one Ensembl id are summed), i.e. the spot is |
| aligned to model.h5ad like the official `ad.concat([model, adata])` step. |
| Tokenising: official NicheformerDataset / tokenization notebooks: counts -> size-factor normalise to 1e4 (over the |
| aligned genes) -> divide by the technology mean (nan->0, rounded half-up, 0 -> 1, as in the notebooks) -> |
| rank the non-zero genes by descending value (stable sort; official uses numba quicksort, so only ties may |
| be ordered differently) -> token = gene index + 30, zero-padded to 1500. Context tokens prepended in the |
| official order [specie, assay, modality] and the sequence truncated to context_length 1500 (=> at most |
| 1497 gene tokens). Pad 0 -> 1 and attention padding mask = (token == 1), as in models/_utils.complete_masking |
| with p = 0. |
| Context : Visium was not in the pretraining corpus (SpatialCorpus-110M = dissociated 10x/Smart-seq + MERFISH/CosMx/ |
| Xenium/ISS) and has no technology token (the "visium": 9 entry in the get_embeddings notebook is a typo |
| for Xenium = 9 in every tokenisation notebook and the HF vocab). Chosen: specie = human (5), |
| modality = spatial (4), assay = 10x 3' v3 (11) -- Visium v1/v2 is a whole-transcriptome poly(A) 3' |
| sequencing assay, so the closest documented technology is 10x 3' v3, and its matching technology mean is |
| the dissociated (10x) one (the image-based means only cover their 300-1000-gene panels). A sensitivity |
| run with modality = dissociated (3) is logged for hest only (not saved). |
| Readout : official Nicheformer.get_embeddings(layer=-1, with_context=False): last-layer output, drop the 3 context |
| positions, mean over the remaining 1497 positions (this INCLUDES padding positions, exactly as the |
| official code; only spots with < 1497 detected vocabulary genes are affected and the fraction is logged). |
| D = 512, not L2-normalised. bf16 autocast inference. |
| Output : rna_sets.save("nicheformer", set, E) -> cache/rnafm/nicheformer_<set>.npy (n_spot, 512) float16. |
| Sanity : finite, mean pairwise cosine of 200 random spots, and (hest) organ agreement of cosine 10-NN vs base rate. |
| """ |
| import hashlib |
| import json |
| import os |
| import pickle |
| import sys |
| import time |
|
|
| import anndata as ad |
| import numpy as np |
| import pandas as pd |
| import scipy.sparse as sp |
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
|
|
| from common import CACHE, ROOT |
| from rna_sets import SETS, load, save |
|
|
| MODEL = "nicheformer" |
| REPO, REV = "theislab/Nicheformer", "0ba4ba162f9ad3d6c8b713761954c31be2fa69ed" |
| GH_COMMIT = "485cadbc5caa15119adfd54228f8a8af835fcabc" |
| EXT = os.path.join(ROOT, "ext", "nicheformer") |
| HF_DIR, GH_DIR = os.path.join(EXT, "hf"), os.path.join(EXT, "repo") |
| PAD, AUX, CTX_LEN = 1, 30, 1500 |
| SPECIE_HUMAN, ASSAY_10X3V3, MOD_SPATIAL, MOD_DISSOC = 5, 11, 4, 3 |
| dev = "cuda" |
| torch.manual_seed(0) |
|
|
|
|
| def sha(p): |
| return hashlib.sha256(open(p, "rb").read()).hexdigest()[:16] |
|
|
|
|
| def ensure_download(): |
| if not os.path.exists(os.path.join(HF_DIR, "model.safetensors")): |
| from huggingface_hub import snapshot_download |
| snapshot_download(REPO, revision=REV, local_dir=HF_DIR) |
| assert os.path.exists(os.path.join(GH_DIR, "data", "model_means", "dissociated_mean_script.npy")), \ |
| f"git clone https://github.com/theislab/nicheformer {GH_DIR} (commit {GH_COMMIT})" |
|
|
|
|
| class Nicheformer(nn.Module): |
| """Re-implementation of NicheformerModel (HF modeling_nicheformer.py / official models/_nicheformer.py).""" |
|
|
| def __init__(self, cfg): |
| super().__init__() |
| self.embeddings = nn.Embedding(cfg["n_tokens"] + 5, cfg["dim_model"], padding_idx=1) |
| self.positional_embedding = nn.Embedding(cfg["context_length"], cfg["dim_model"]) |
| layer = nn.TransformerEncoderLayer(cfg["dim_model"], cfg["nheads"], cfg["dim_feedforward"], dropout=0.0, |
| batch_first=True, layer_norm_eps=1e-12) |
| self.encoder = nn.TransformerEncoder(layer, cfg["nlayers"], enable_nested_tensor=False) |
|
|
| def get_embeddings(self, ids): |
| pad = ids.eq(PAD) |
| x = self.embeddings(ids) + self.positional_embedding(torch.arange(ids.shape[1], device=ids.device)) |
| for blk in self.encoder.layers: |
| x = blk(x, src_key_padding_mask=pad, is_causal=False) |
| return x[:, 3:].float().mean(1) |
|
|
|
|
| def load_model(): |
| cfg = json.load(open(os.path.join(HF_DIR, "config.json"))) |
| assert cfg["learnable_pe"] and cfg["specie"] and cfg["assay"] and cfg["modality"] |
| sd = load_file(os.path.join(HF_DIR, "model.safetensors")) |
| keep = {k[len("nicheformer."):]: v for k, v in sd.items() |
| if k.startswith("nicheformer.") and not k.startswith("nicheformer.encoder_layer.") and k != "nicheformer.pos"} |
| skipped = sorted(set(sd) - {"nicheformer." + k for k in keep}) |
| assert all(k.startswith(("nicheformer.encoder_layer.", "classifier_head.")) or k == "nicheformer.pos" |
| for k in skipped), skipped |
| m = Nicheformer(cfg) |
| m.load_state_dict(keep, strict=True) |
| print(f"weights {sha(os.path.join(HF_DIR, 'model.safetensors'))}: loaded {len(keep)} tensors strictly; " |
| f"skipped {len(skipped)} (template encoder_layer / MLM head / pos buffer)", flush=True) |
| return m.eval().to(dev), cfg |
|
|
|
|
| def gene_reference(): |
| ref = ad.read_h5ad(os.path.join(HF_DIR, "model.h5ad"), backed=None) |
| ens = np.asarray(ref.var_names, dtype=str) |
| vocab = json.load(open(os.path.join(HF_DIR, "vocab.json"))) |
| assert all(vocab[g] == i + AUX for i, g in enumerate(ens)) |
| mean = np.load(os.path.join(GH_DIR, "data", "model_means", "dissociated_mean_script.npy")) |
| assert mean.shape == ens.shape |
| mean = np.nan_to_num(mean) |
| r = np.where((mean % 1) >= 0.5, np.ceil(mean), np.floor(mean)) |
| mean = np.where(r == 0, 1, r) |
| sym2ens = dict(pd.read_csv(os.path.join(CACHE, "external", "dlpfc_genes_full.csv"))[["gene", "ensembl_id"]].values) |
| from huggingface_hub import hf_hub_download |
| gf = pickle.load(open(hf_hub_download("ctheodoris/Geneformer", "geneformer/gene_name_id_dict_gc104M.pkl"), "rb")) |
| for k, v in gf.items(): |
| sym2ens.setdefault(k, v) |
| return ens, mean, sym2ens |
|
|
|
|
| def tokenize(X, genes, ens, mean, sym2ens, modality): |
| e2i = {e: i for i, e in enumerate(ens)} |
| src = [(j, e2i[sym2ens[g]]) for j, g in enumerate(genes) if g in sym2ens and sym2ens[g] in e2i] |
| rows, cols = zip(*src) |
| M = sp.csr_matrix((np.ones(len(src), np.float32), (rows, cols)), shape=(len(genes), len(ens))) |
| A = (X @ M).tocsr() |
| lib = np.asarray(A.sum(1)).ravel() |
| cov = dict(n_in=len(genes), n_mapped=len(src), n_model_genes=int((np.asarray(M.sum(0)).ravel() > 0).sum()), |
| umi_frac=float(np.median(lib / np.maximum(np.asarray(X.sum(1)).ravel(), 1)))) |
| lib = lib + (lib == 0) |
| A = (sp.diags(1e4 / lib) @ A).tocsr() |
| ids = np.full((A.shape[0], CTX_LEN), PAD, np.int64) |
| ids[:, :3] = [SPECIE_HUMAN, ASSAY_10X3V3, modality] |
| ngen = np.zeros(A.shape[0], int) |
| for r in range(A.shape[0]): |
| s, e = A.indptr[r], A.indptr[r + 1] |
| nz, v = A.indices[s:e], A.data[s:e] / mean[A.indices[s:e]] |
| keep = v > 0 |
| nz, v = nz[keep], v[keep] |
| order = nz[np.argsort(-v, kind="stable")][:CTX_LEN - 3] |
| ids[r, 3:3 + len(order)] = order + AUX |
| ngen[r] = len(order) |
| cov["expr_genes_median"] = float(np.median(np.diff(A.indptr))) |
| cov["tokens_median"] = float(np.median(ngen)) |
| cov["frac_spots_padded"] = float((ngen < CTX_LEN - 3).mean()) |
| return ids, cov |
|
|
|
|
| @torch.no_grad() |
| def embed(model, ids, bs=32): |
| out = np.zeros((len(ids), 512), np.float32) |
| i = 0 |
| while i < len(ids): |
| try: |
| b = torch.from_numpy(ids[i:i + bs]).to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out[i:i + bs] = model.get_embeddings(b).cpu().numpy() |
| i += bs |
| except torch.OutOfMemoryError: |
| torch.cuda.empty_cache() |
| bs = max(1, bs // 2) |
| print(f" OOM -> batch {bs}", flush=True) |
| return out |
|
|
|
|
| def sanity(E, name): |
| rng = np.random.default_rng(0) |
| S = E[rng.choice(len(E), min(200, len(E)), replace=False)] |
| S = S / np.linalg.norm(S, axis=1, keepdims=True) |
| C = S @ S.T |
| mc = float(C[np.triu_indices(len(S), 1)].mean()) |
| msg = f"[{MODEL}] {name}: finite={np.isfinite(E).all()} std={E.std(0).mean():.4f} mean-pair-cos(200)={mc:.4f}" |
| if name == "hest": |
| org = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")).organ.values |
| assert len(org) == len(E) |
| Z = torch.from_numpy(E / np.linalg.norm(E, axis=1, keepdims=True)).to(dev) |
| nn_idx = [] |
| for k in range(0, len(Z), 2048): |
| s = Z[k:k + 2048] @ Z.T |
| s[torch.arange(s.shape[0]), torch.arange(k, k + s.shape[0])] = -2 |
| nn_idx.append(s.topk(10, dim=1).indices.cpu().numpy()) |
| nn_idx = np.concatenate(nn_idx) |
| knn = float((org[nn_idx] == org[:, None]).mean()) |
| p = pd.Series(org).value_counts(normalize=True).values |
| msg += f" 10-NN same-organ={knn:.3f} (base rate {float((p ** 2).sum()):.3f})" |
| print(msg, flush=True) |
| return mc |
|
|
|
|
| def main(): |
| sets = sys.argv[1:] or SETS |
| t0 = time.time() |
| ensure_download() |
| model, cfg = load_model() |
| ens, mean, sym2ens = gene_reference() |
| print(f"context tokens: specie=human({SPECIE_HUMAN}) assay=10x 3' v3({ASSAY_10X3V3}) modality=spatial({MOD_SPATIAL}); " |
| f"technology mean = dissociated_mean_script.npy (GitHub {GH_COMMIT[:7]})", flush=True) |
| for name in sets: |
| t1 = time.time() |
| X, genes = load(name) |
| ids, cov = tokenize(X, genes, ens, mean, sym2ens, MOD_SPATIAL) |
| print(f"[{MODEL}] {name}: {X.shape[0]} spots; genes in {cov['n_in']}, mapped to vocab {cov['n_mapped']} " |
| f"({cov['n_mapped'] / cov['n_in']:.1%}) -> {cov['n_model_genes']} model genes; median UMI fraction kept " |
| f"{cov['umi_frac']:.3f}; expressed vocab genes/spot median {cov['expr_genes_median']:.0f}; gene tokens " |
| f"median {cov['tokens_median']:.0f}; spots with padding {cov['frac_spots_padded']:.1%}", flush=True) |
| E = embed(model, ids) |
| sanity(E, name) |
| save(MODEL, name, E) |
| if name == "hest": |
| ids[:, 2] = MOD_DISSOC |
| E2 = embed(model, ids) |
| c = (E * E2).sum(1) / np.linalg.norm(E, axis=1) / np.linalg.norm(E2, axis=1) |
| print(f" sensitivity modality=dissociated(3): cosine to saved per spot median {np.median(c):.4f}", flush=True) |
| sanity(E2, "hest") |
| print(f"[{MODEL}] {name}: {time.time() - t1:.0f}s", flush=True) |
| print(f"[{MODEL}] done {time.time() - t0:.0f}s; peak GPU {torch.cuda.max_memory_allocated() / 2**30:.2f} GB", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
