File size: 6,220 Bytes
e8fa066 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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]) # length-sorted batches (less padding), restored below
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")
|