| """Frozen single-cell foundation-model embeddings of HEST Visium spot expression (each spot = one pseudo-cell). |
| |
| python 22_extract_gene_fm.py scgpt | geneformer |
| |
| scgpt : wanglab/scGPT-human (same re-implementation / strict load as 15_extract_marker_fm.py); input = expressed |
| genes among the 2000 HVGs that are in the scGPT vocab, raw counts -> log1p CP10k -> per-spot quantile |
| binning (51 bins); at most 1200 genes (random subset, as scgpt.tasks.embed_data max_length=1200); |
| <cls> output, L2-normalised. |
| geneformer : ctheodoris/Geneformer V2-104M; input = rank-value encoding of all common genes (count / total * 1e4 / |
| gc104M gene median, descending, nonzero), truncated to 4096 tokens incl. <cls>/<eos>; |
| penultimate layer, mean over gene tokens. |
| Output cache/hest/gene_<name>.npy (n_spot, 1, D) float16 |
| """ |
| import importlib |
| import os |
| import pickle |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import scipy.sparse as sp |
| import torch |
|
|
| from common import CACHE |
|
|
| M15 = importlib.import_module("15_extract_marker_fm") |
| OUT = os.path.join(CACHE, "hest") |
| name = sys.argv[1] |
| dev = "cuda" |
| X = sp.load_npz(os.path.join(OUT, "counts.npz")).tocsr() |
| genes = pd.read_csv(os.path.join(OUT, "genes.csv")).gene.values |
| tot = np.asarray(X.sum(1)).ravel() |
| t0 = time.time() |
|
|
|
|
| def batches(seqs, bs): |
| for i in range(0, len(seqs), bs): |
| yield seqs[i:i + bs] |
|
|
|
|
| if name == "scgpt": |
| import json |
| from huggingface_hub import hf_hub_download |
| args = json.load(open(hf_hub_download(M15.SCGPT_REPO, "args.json"))) |
| vocab = json.load(open(hf_hub_download(M15.SCGPT_REPO, "vocab.json"))) |
| sd = torch.load(hf_hub_download(M15.SCGPT_REPO, "best_model.pt"), map_location="cpu", weights_only=True) |
| sd = {k.replace("self_attn.Wqkv.", "self_attn.in_proj_"): v for k, v in sd.items()} |
| PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"] |
| model = M15.ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], PAD) |
| model.load_state_dict(sd, strict=True) |
| model.eval().to(dev) |
| hvg = pd.read_csv(os.path.join(OUT, "hvg_genes.csv")).gene.values |
| gi = np.array([i for i, g in enumerate(genes) if g in set(hvg) and g in vocab]) |
| gid = np.array([vocab[g] for g in genes[gi]]) |
| L = sp.diags(1e4 / tot) @ X[:, gi] |
| L.data = np.log1p(L.data) |
| L = L.tocsr() |
| rng = np.random.default_rng(0) |
| seqs = [] |
| for r in range(L.shape[0]): |
| row = L.getrow(r) |
| nz, v = row.indices, row.data |
| if len(nz) > 1199: |
| keep = np.sort(rng.choice(len(nz), 1199, replace=False)) |
| nz, v = nz[keep], v[keep] |
| vals = M15.scgpt_binning(v, args["n_bins"], rng).astype(np.float32) |
| seqs.append((np.r_[CLS, gid[nz]], np.r_[np.float32(args["pad_value"]), vals])) |
| print(f"scGPT: {len(gi)} HVGs in vocab; genes/spot median {np.median([len(s[0]) - 1 for s in seqs])}", flush=True) |
| embs = [] |
| with torch.no_grad(): |
| for b in batches(seqs, 32): |
| Lm = max(len(g) for g, _ in b) |
| g = torch.full((len(b), Lm), PAD, dtype=torch.long) |
| v = torch.full((len(b), Lm), float(args["pad_value"])) |
| for j, (gg, vv) in enumerate(b): |
| g[j, :len(gg)], v[j, :len(vv)] = torch.from_numpy(gg), torch.from_numpy(vv) |
| g, v = g.to(dev), v.to(dev) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| h = model.encode(g, v, g.eq(PAD)) |
| embs.append(h[:, 0].float().cpu()) |
| E = torch.cat(embs).numpy() |
| E /= np.linalg.norm(E, axis=1, keepdims=True) |
| else: |
| from transformers import BertModel |
| tokd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/token_dictionary_gc104M.pkl"), "rb")) |
| medd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_median_dictionary_gc104M.pkl"), "rb")) |
| n2id = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_name_id_dict_gc104M.pkl"), "rb")) |
| mdir = os.path.dirname(M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/model.safetensors")) |
| M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/config.json") |
| gi = np.array([i for i, g in enumerate(genes) if g in n2id and n2id[g] in tokd and n2id[g] in medd]) |
| tok = np.array([tokd[n2id[g]] for g in genes[gi]]) |
| med = np.array([medd[n2id[g]] for g in genes[gi]]) |
| CLS, EOS, PAD = tokd["<cls>"], tokd["<eos>"], tokd["<pad>"] |
| model = BertModel.from_pretrained(mdir, add_pooling_layer=False, attn_implementation="sdpa").eval().to(dev) |
| Xg = (sp.diags(1e4 / tot) @ X[:, gi]).tocsr() |
| seqs = [] |
| for r in range(Xg.shape[0]): |
| row = Xg.getrow(r) |
| nz, v = row.indices, row.data / med[row.indices] |
| order = nz[np.argsort(-v, kind="stable")][:4094] |
| seqs.append([CLS] + tok[order].tolist() + [EOS]) |
| print(f"Geneformer: {len(gi)} genes in dictionary; tokens/spot median {np.median([len(s) for s in seqs])}", flush=True) |
| order_len = np.argsort([len(s) for s in seqs]) |
| embs = np.zeros((len(seqs), model.config.hidden_size), np.float32) |
| with torch.no_grad(): |
| for k in range(0, len(seqs), 8): |
| idx = order_len[k:k + 8] |
| b = [seqs[i] for i in idx] |
| Lm = max(map(len, b)) |
| ids = torch.full((len(b), Lm), PAD, dtype=torch.long) |
| for j, s in enumerate(b): |
| ids[j, :len(s)] = torch.tensor(s) |
| ids = ids.to(dev) |
| att = (ids != PAD).long() |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| h = model(input_ids=ids, attention_mask=att, output_hidden_states=True).hidden_states[-2].float() |
| gm = att.clone().float() |
| gm[:, 0] = 0 |
| gm[torch.arange(len(b)), att.sum(1) - 1] = 0 |
| embs[idx] = ((h * gm[..., None]).sum(1) / gm.sum(1, keepdim=True).clamp(min=1)).cpu().numpy() |
| if k % 800 == 0: |
| print(f" {k}/{len(seqs)} {time.time() - t0:.0f}s", flush=True) |
| E = embs |
| np.save(os.path.join(OUT, f"gene_{name}.npy"), E.astype(np.float16)[:, None]) |
| print(f"[{name}] saved {E.shape} finite={np.isfinite(E).all()} {time.time() - t0:.0f}s") |
|
|
