Blame for 43_rnafm_aidocell.py

"""Frozen AIDO.Cell-100M (a.k.a. GB.Cell-100M) cell embeddings of the RNA sets (each spot = one pseudo-cell).

    python 43_rnafm_aidocell.py [set ...]        (default: all of rna_sets.SETS)

Model   : HF genbio-ai/AIDO.Cell-100M @ a983b388e1025860f2c413ecef60e76078253c24 (model.safetensors, config.json;
          identical sha256 to genbio-ai/GB.Cell-100M, the renamed repo). Licence "other" = GenBio AI Community
          License: non-commercial use only; outputs (explicitly incl. embeddings) may not be used to train models
          "similar to the GenBio Materials" (other cell foundation models) or for distillation.
          Code reference: github genbio-ai/ModelGenerator @ bda6f2da (cloned to ext/modelgenerator, read, NOT
          installed / imported), huggingface/gb.cell/gb_cell/{models/modeling_cellfoundation.py,
          utils/gene_alignment.py, utils/preprocessing.py}, quickstart.ipynb, embed.py.
          The forward pass of CellFoundationModel is re-implemented below; the bert.* tensors load with strict=True
          (the only other tensors, cls.* = MLM head, are unused):
            gene_embedding : auto-discretisation  v -> Linear(1,100) -> LeakyReLU(0.1) -> x + Linear(100,100)(x)
                             (bin_alpha=1) -> softmax over 100 bins -> @ lookup_table (100 x 640)
            + absolute position embedding (19266 x 640; position = gene index, 19264 genes + 2 depth tokens)
            18 x layer     : pre-LN; MHA 20 heads x 32, no linear biases, scale 1/sqrt(32), FULL attention over all
                             19,266 tokens (the official flash-attn path ignores the attention mask: "assuming NO
                             PADDING"; its SDPA fallback likewise uses attn_mask=None) -> torch SDPA here, no
                             flash-attn package needed; SwiGLU MLP 640->1664->640; final LayerNorm.
          Weights and inputs in bf16 as in embed.py (model.to(torch.bfloat16)).
Input   : official gene_alignment.align_adata: data genes matched by HGNC symbol to the fixed 19,264-gene list
          (OS_scRNA_gene_index.19264.tsv, = scFoundation's list), other genes dropped, missing genes = 0 (still
          attended to, as in the official code); duplicate data symbols keep the first column.
          official preprocessing.preprocess_counts on the aligned raw counts: total = sum over the 19,264 aligned genes;
          two depth tokens max(log10(total), 5) appended; X = log1p(count / total * 1e4); clip at 20; bf16.
          (Deviation: an all-zero spot -- 1 each in hest_v / tcga_v after rounding the virtual counts -- gets an
          all-zero gene input instead of the official 0/0 = NaN.)
Readout : quickstart.ipynb cell-level embedding: last_hidden_state (after final LN) with the 2 depth tokens dropped,
          mean over the gene positions whose gene is present in the data (align_adata attention_mask = 1; i.e. the
          genes the dataset measured, whether or not expressed in that spot).  D = 640.  Not L2-normalised.
Output  : rna_sets.save("aidocell", set, E) -> cache/rnafm/aidocell_<set>.npy  (n_spot, 640) float16.
Log     : logs/rnafm/aidocell.log (stdout of this script, tee'd by the caller).
"""
import json
import os
import sys
import time

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from common import CACHE, ROOT
from rna_sets import SETS, load, save

MODEL = "aidocell"
REPO, REV = "genbio-ai/AIDO.Cell-100M", "a983b388e1025860f2c413ecef60e76078253c24"
GENE_LIST = os.path.join(ROOT, "ext", "modelgenerator", "huggingface", "gb.cell", "gb_cell", "models", "gene_lists",
                         "OS_scRNA_gene_index.19264.tsv")
DEV = "cuda"


# ----------------------------------------------------------------------------- model (re-implementation)
class Layer(nn.Module):
    def __init__(self, d, h, ff, eps):
        super().__init__()
        self.h = h
        self.attention = nn.Module()
        self.attention.ln = nn.LayerNorm(d, eps=eps)
        self.attention.self = nn.Module()
        for n in ["query", "key", "value"]:
            setattr(self.attention.self, n, nn.Linear(d, d, bias=False))
        self.attention.output = nn.Module()
        self.attention.output.dense = nn.Linear(d, d, bias=False)
        self.ln = nn.LayerNorm(d, eps=eps)
        self.mlp = nn.Module()
        self.mlp.up_proj, self.mlp.gate_proj = nn.Linear(d, ff, bias=False), nn.Linear(d, ff, bias=False)
        self.mlp.down_proj = nn.Linear(ff, d, bias=False)

    def forward(self, x):
        B, S, D = x.shape
        a, y = self.attention, self.attention.ln(x)
        q, k, v = (getattr(a.self, n)(y).view(B, S, self.h, D // self.h).transpose(1, 2)
                   for n in ["query", "key", "value"])
        o = F.scaled_dot_product_attention(q, k, v, scale=(D // self.h) ** -0.5)
        x = x + a.output.dense(o.transpose(1, 2).reshape(B, S, D))
        y = self.ln(x)
        return x + self.mlp.down_proj(F.silu(self.mlp.gate_proj(y)) * self.mlp.up_proj(y))


class AIDOCell(nn.Module):
    def __init__(self, c):
        super().__init__()
        assert c["hidden_act"] == "swiglu" and c["normalization_type"] == "LayerNorm" and not c["add_linear_bias"]
        assert c["position_embedding_type"] == "absolute"
        d, nb = c["hidden_size"], c["bin_num"]
        self.alpha = c["bin_alpha"]
        ge = nn.Module()
        ge.linear1, ge.linear2 = nn.Linear(1, nb), nn.Linear(nb, nb)
        ge.lookup_table, ge.mask_embedding, ge.pad_embedding = nn.Embedding(nb, d), nn.Embedding(1, d), nn.Embedding(1, d)
        self.gene_embedding = ge
        self.position_embedding = nn.Embedding(c["max_position_embeddings"] + 2, d)
        self.encoder = nn.Module()
        self.encoder.layer = nn.ModuleList([Layer(d, c["num_attention_heads"], c["intermediate_size"],
                                                  c["layer_norm_eps"]) for _ in range(c["num_hidden_layers"])])
        self.encoder.ln = nn.LayerNorm(d, eps=c["layer_norm_eps"])

    def forward(self, X):                                 # X (B, S) preprocessed values >= 0
        ge = self.gene_embedding
        x = F.leaky_relu(ge.linear1(X.unsqueeze(-1)), 0.1)
        x = self.alpha * x + ge.linear2(x)
        x = torch.softmax(x, dim=-1) @ ge.lookup_table.weight
        x = x + self.position_embedding.weight[:X.shape[1]]
        for lyr in self.encoder.layer:
            x = lyr(x)
        return self.encoder.ln(x)


def load_model():
    cfg = json.load(open(hf_hub_download(REPO, "config.json", revision=REV)))
    sd = load_file(hf_hub_download(REPO, "model.safetensors", revision=REV))
    bert = {k[len("bert."):]: v for k, v in sd.items() if k.startswith("bert.")}
    rest = sorted(k for k in sd if not k.startswith("bert."))
    assert all(k.startswith("cls.") for k in rest), rest
    model = AIDOCell(cfg)
    model.load_state_dict(bert, strict=True)
    print(f"loaded {REPO}@{REV[:8]}/model.safetensors: {len(bert)} bert.* tensors strictly "
          f"({sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params); unused MLM head: {rest}", flush=True)
    return model.to(torch.bfloat16).eval().to(DEV), cfg


# ----------------------------------------------------------------------------- input (official alignment / preprocessing)
MODEL_GENES = pd.read_csv(GENE_LIST, sep="\t")["gene_name"].to_numpy()
assert len(MODEL_GENES) == 19264 and len(set(MODEL_GENES)) == 19264


def align(X, genes, name):
    pos = {}
    for j, g in enumerate(genes):
        pos.setdefault(g, j)
    mi = np.array([pos.get(g, -1) for g in MODEL_GENES])
    present = mi >= 0
    rows, cols = mi[present], np.nonzero(present)[0]
    P = sp.csr_matrix((np.ones(len(rows), np.float32), (rows, cols)), shape=(len(genes), len(MODEL_GENES)))
    Xa = (X @ P).tocsr()
    n_dup = len(genes) - len(pos)
    print(f"[{name}] data genes {len(genes)} ({n_dup} duplicate symbols); used {present.sum()}; "
          f"coverage of the 19,264 AIDO genes: {present.mean():.1%}; data genes not in AIDO list: "
          f"{len(pos) - present.sum()}", flush=True)
    return Xa, present


def preprocess(counts):
    X = torch.from_numpy(counts).to(device=DEV, dtype=torch.float32)
    tot = X.sum(1, keepdim=True)
    depth = torch.clamp(torch.log10(tot), min=5.0)          # log10(0) = -inf -> 5 for an empty spot
    X = torch.log1p(X / tot.clamp(min=1) * 10000)           # empty spot: all-zero input instead of 0/0 = NaN
    X = torch.cat([X, depth, depth], dim=1)
    return torch.clamp(X, max=20).to(torch.bfloat16)


@torch.no_grad()
def embed(model, Xa, present, bs=4):
    G = Xa.shape[1]
    m = torch.from_numpy(present.astype(np.float32)).to(DEV)[None, :, None]
    E = np.zeros((Xa.shape[0], model.encoder.ln.normalized_shape[0]), np.float32)
    t0, k = time.time(), 0
    while k < Xa.shape[0]:
        try:
            h = model(preprocess(Xa[k:k + bs].toarray()))[:, :G].float()      # drop the 2 depth tokens
            E[k:k + bs] = ((h * m).sum(1) / m.sum(1)).cpu().numpy()
            k += bs
            if (k // bs) % 250 == 0:
                print(f"  {k}/{Xa.shape[0]}  {time.time() - t0:.0f}s", flush=True)
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            bs = max(1, bs // 2)
            print(f"  OOM -> batch {bs}", flush=True)
    return E


# ----------------------------------------------------------------------------- sanity
def sanity(E, name, rng):
    out = {"finite": bool(np.isfinite(E).all())}
    s = rng.choice(len(E), min(200, len(E)), replace=False)
    Z = E[s] / np.linalg.norm(E[s], axis=1, keepdims=True)
    C = Z @ Z.T
    out["mean_pairwise_cos200"] = float(C[np.triu_indices(len(s), 1)].mean())
    if name == "hest":
        organ = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")).organ.values
        assert len(organ) == len(E)
        Zt = torch.from_numpy(E / np.linalg.norm(E, axis=1, keepdims=True)).to(DEV)
        nn_idx = []
        for i in range(0, len(Zt), 2048):
            sim = Zt[i:i + 2048] @ Zt.T
            sim[torch.arange(sim.shape[0]), torch.arange(i, i + sim.shape[0])] = -2
            nn_idx.append(sim.topk(10, dim=1).indices.cpu().numpy())
        nn_idx = np.concatenate(nn_idx)
        out["knn10_same_organ"] = float((organ[nn_idx] == organ[:, None]).mean())
        cnt = pd.Series(organ).value_counts()
        out["organ_base_rate"] = float(((cnt * (cnt - 1)).sum()) / (len(organ) * (len(organ) - 1)))
    return out


def main():
    sets = sys.argv[1:] or SETS
    t0 = time.time()
    model, cfg = load_model()
    report = {}
    for s in sets:
        t1 = time.time()
        X, genes = load(s)
        Xa, present = align(X, genes, s)
        tot = np.asarray(Xa.sum(1)).ravel()
        nexp = np.diff(Xa.indptr)
        print(f"[{s}] aligned UMI/spot median {np.median(tot):.0f} (min {tot.min():.0f}); expressed genes/spot median "
              f"{np.median(nexp):.0f}; spots with total>=1e5 (depth token > 5): {(tot >= 1e5).sum()}", flush=True)
        if (tot == 0).any():
            print(f"[{s}] {(tot == 0).sum()} empty spot(s) (all-zero counts): input = all zeros + depth tokens 5 "
                  f"(the official code would give NaN from 0/0)", flush=True)
        E = embed(model, Xa, present)
        save(MODEL, s, E)
        info = dict(n_genes=len(genes), n_aido_genes_present=int(present.sum()), coverage=float(present.mean()),
                    median_expr=float(np.median(nexp)))
        info.update(sanity(E, s, np.random.default_rng(0)))
        info["D"], info["seconds"] = int(E.shape[1]), round(time.time() - t1, 1)
        report[s] = info
        print(f"[{s}] {json.dumps(info)}", flush=True)
    print(f"done in {time.time() - t0:.0f}s; peak GPU mem {torch.cuda.max_memory_allocated() / 2**30:.2f} GB")
    print(json.dumps(report, indent=1))


if __name__ == "__main__":
    main()