| """Frozen RNA embeddings for every RNA set of rna_sets.py (hest, hest_v, orion_v, tcga_v, dlpfc). |
| |
| python 43_rnafm_base.py scgpt | geneformer | genept-bge | genept-bmc |
| |
| scgpt / geneformer : same preprocessing and readout as 22_extract_gene_fm.py (scGPT-human <cls> on the HEST 2000 HVGs |
| in vocab, quantile-binned log1p CP10k, <=1200 genes; Geneformer-V2-104M rank-value encoding, |
| penultimate layer mean over gene tokens). 22_ only embedded HEST; this adds the virtual / external sets. |
| genept-* : GenePT-w (Chen & Zou 2024) with open text encoders: each gene = text embedding of its NCBI + UniProt |
| summary (GenePT v2 release, Zenodo 10833191, CC BY 4.0, ext/genept/); spot = average of the gene |
| embeddings weighted by log1p CP10k over the genes that have a summary, L2-normalised. |
| bge = BAAI/bge-base-en-v1.5 (CLS, 512 tokens), bmc = BiomedCLIP PubMedBERT text tower (256 tokens, |
| the text tower of the text modality -> RNA lands in the same text space). |
| Output cache/rnafm/<model>_<set>.npy (rna_sets.save) |
| """ |
| import importlib |
| import json |
| 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, ROOT |
| from rna_sets import SETS, load, save |
|
|
| M15 = importlib.import_module("15_extract_marker_fm") |
| name = sys.argv[1] |
| sets = sys.argv[2:] or SETS |
| dev = "cuda" |
| HVG = pd.read_csv(os.path.join(CACHE, "hest", "hvg_genes.csv")).gene.values |
| t0 = time.time() |
|
|
|
|
| def cp10k_log(X, cols): |
| tot = np.asarray(X.sum(1)).ravel() |
| L = (sp.diags(1e4 / np.maximum(tot, 1)) @ X[:, cols]).tocsr() |
| L.data = np.log1p(L.data) |
| return L |
|
|
|
|
| def scgpt_model(): |
| 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()} |
| model = M15.ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], vocab[args["pad_token"]]) |
| model.load_state_dict(sd, strict=True) |
| return model.eval().to(dev), args, vocab |
|
|
|
|
| def scgpt(X, genes, model, args, vocab): |
| PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"] |
| hv = set(HVG) |
| gi = np.array([i for i, g in enumerate(genes) if g in hv and g in vocab]) |
| gid = np.array([vocab[g] for g in genes[gi]]) |
| L = cp10k_log(X, gi) |
| rng = np.random.default_rng(0) |
| seqs = [] |
| for r in range(L.shape[0]): |
| nz, v = L.indices[L.indptr[r]:L.indptr[r + 1]], L.data[L.indptr[r]:L.indptr[r + 1]] |
| 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) if len(v) else np.zeros(0, 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 i in range(0, len(seqs), 32): |
| b = seqs[i:i + 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() |
| return E / np.linalg.norm(E, axis=1, keepdims=True) |
|
|
|
|
| def gf_model(): |
| 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") |
| model = BertModel.from_pretrained(mdir, add_pooling_layer=False, attn_implementation="sdpa").eval().to(dev) |
| return model, tokd, medd, n2id |
|
|
|
|
| def geneformer(X, genes, model, tokd, medd, n2id): |
| 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>"] |
| tot = np.asarray(X.sum(1)).ravel() |
| Xg = (sp.diags(1e4 / np.maximum(tot, 1)) @ X[:, gi]).tocsr() |
| seqs = [] |
| for r in range(Xg.shape[0]): |
| nz, v = Xg.indices[Xg.indptr[r]:Xg.indptr[r + 1]], Xg.data[Xg.indptr[r]:Xg.indptr[r + 1]] |
| v = v / med[nz] |
| 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] |
| ids = torch.full((len(b), max(map(len, b))), 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() |
| return embs |
|
|
|
|
| def gene_text_table(enc): |
| """Text embedding of each gene's NCBI + UniProt summary; cached in cache/rnafm/genetext_<enc>.npz.""" |
| f = os.path.join(CACHE, "rnafm", f"genetext_{enc}.npz") |
| if os.path.exists(f): |
| z = np.load(f, allow_pickle=False) |
| return z["genes"], z["E"] |
| summ = json.load(open(os.path.join(ROOT, "ext", "genept", "unz", "GenePT_emebdding_v2", "NCBI_UniProt_summary_of_genes.json"))) |
| g = np.array(sorted(summ)) |
| txt = [summ[x] for x in g] |
| out = [] |
| if enc == "bge": |
| from transformers import AutoModel, AutoTokenizer |
| tk = AutoTokenizer.from_pretrained("BAAI/bge-base-en-v1.5") |
| m = AutoModel.from_pretrained("BAAI/bge-base-en-v1.5").to(dev).eval() |
| f_ = lambda b: m(**{k: v.to(dev) for k, v in tk(b, padding=True, truncation=True, max_length=512, return_tensors="pt").items()}).last_hidden_state[:, 0] |
| else: |
| from biomedclip_text import BiomedCLIPText |
| m = BiomedCLIPText().to(dev).eval() |
| f_ = m |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| for i in range(0, len(txt), 128): |
| out.append(torch.nn.functional.normalize(f_(txt[i:i + 128]).float(), dim=-1).cpu()) |
| E = torch.cat(out).numpy() |
| os.makedirs(os.path.dirname(f), exist_ok=True) |
| np.savez(f, genes=g, E=E) |
| print(f" gene texts: {len(g)} genes embedded with {enc} {E.shape}", flush=True) |
| return g, E |
|
|
|
|
| def genept(X, genes, gt, GE): |
| idx = pd.Series(range(len(gt)), index=gt) |
| gi = np.array([i for i, g in enumerate(genes) if g in idx.index]) |
| W = cp10k_log(X, gi) |
| print(f" GenePT: {len(gi)}/{len(genes)} genes with a summary", flush=True) |
| E = W @ GE[idx.loc[genes[gi]].values] |
| E = E / np.maximum(np.asarray(W.sum(1)), 1e-8) |
| return E / np.linalg.norm(E, axis=1, keepdims=True).clip(1e-8) |
|
|
|
|
| if name == "scgpt": |
| ctx = scgpt_model(); fn = scgpt |
| elif name == "geneformer": |
| ctx = gf_model(); fn = geneformer |
| elif name.startswith("genept-"): |
| ctx = gene_text_table(name.split("-")[1]); fn = genept |
| else: |
| raise ValueError(name) |
| for s in sets: |
| X, genes = load(s) |
| print(f"[{name}] {s} {X.shape}", flush=True) |
| save(name, s, fn(X, genes, *ctx)) |
| print(f" {time.time() - t0:.0f}s", flush=True) |
|
|
