| """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" |
|
|
|
|
| |
| 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): |
| 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 |
|
|
|
|
| |
| 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) |
| X = torch.log1p(X / tot.clamp(min=1) * 10000) |
| 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() |
| 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 |
|
|
|
|
| |
| 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() |
|
|
