"""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_.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()