Blame for experiments/43_rnafm_base.py

"""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  # 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)