"""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 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/_.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[""] 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[""], tokd[""], tokd[""] 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_.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 # open_clip-free, reproduces cache/alltxt_text_biomedclip.npy (cos 1.000) 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)