"""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_.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) # with_context=False, mean incl. pad positions (official) 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() # spot x model genes (aligned like concat with model) 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": # sensitivity: modality token = dissociated 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()