| """Frozen Arc Institute STATE embedding (SE-600M) spot embeddings of the 5 RNA sets (each spot = one pseudo-cell). |
| |
| python 43_rnafm_se600m.py [set ...] (default: all of rna_sets.SETS) |
| |
| Model : arcinstitute/SE-600M (HF), revision 5a9a80f44f7ce32ce57059933ef0d735d7c10ce5, files model.safetensors |
| (metadata: converted from se600m_epoch16.ckpt, i.e. the final 16-epoch checkpoint), config.yaml and |
| protein_embeddings.pt (19,790 gene symbols -> ESM2-15B mean protein embedding, 5120-d; identical, in |
| the same order, to the checkpoint's pe_embedding.weight -- asserted). Licence: Arc Research Institute |
| State Model Non-Commercial License + Acceptable Use Policy (weights; non-commercial research use only), |
| code CC BY-NC-SA 4.0. Official code: github.com/ArcInstitute/state (commit 9bbfe78a, cloned to |
| ext/state/state, read but NOT imported / installed): the forward pass of src/state/emb/nn/model.py |
| (StateEmbeddingModel) + nn/flash_transformer.py is re-implemented below and loaded with a strict |
| state-dict load of all 238 tensors (the encoder / gene_embedding_layer tie is asserted); the unused |
| heads (binary_decoder, dataset_encoder, gene_embedding_layer) are dropped after loading. |
| 16-layer post-LN transformer, d=2048, 16 heads, FFN 2048, GELU, no positional encoding, no attention mask. |
| Inputs : rna_sets.load(set) raw(-like) counts + HGNC symbols. Genes = the set's genes whose symbol is a key of |
| protein_embeddings.pt (official FilteredGenesCounts valid mask, set gene order kept); other genes dropped. |
| Preprocess: official `state emb transform` path = Inference.encode_adata -> create_dataloader -> |
| VCIDatasetSentenceCollator.sample_cell_sentences (is_train=False): |
| * counts are raw -> log1p (the official per-cell raw-vs-log1p heuristic, max>35 or sum(expm1)>5e6, is |
| bypassed because all inputs are known raw counts; how many spots it would misclassify is logged); |
| * expression weight w = log1p / sum(log1p) over the valid genes (uniform if the spot is all-zero); |
| * cell sentence (pad_length 2048) = [CLS] + 2047 genes ranked by descending expression with random |
| tie-breaking (every set has > 2047 valid genes, so: all detected genes in rank order, top 2047 if more, |
| then randomly drawn undetected genes without replacement); seeded numpy RNG (seed 0) instead of |
| torch.randperm; |
| * token count value = 100 * w[gene]; position 0 carries 100 * w[local gene index 3] exactly as the |
| official collator (cls_token_idx = 3 is written as a local gene index before the count lookup); |
| * token embedding = L2-normalised ESM2 embedding, position 0 replaced by the learned cls_token, the |
| learned dataset_token appended (count embedding 0) -> 2049 tokens. |
| Forward : encoder(Linear 5120->2048, LayerNorm, SiLU) * sqrt(2048) + soft-binned count embedding |
| (softmax(count_encoder(c)) @ bin_encoder) -> 16 layers -> decoder (SkipBlock + Linear). |
| Whole model in bf16 (official get_precision_config('cuda') casts the model to bf16 + bf16 autocast). |
| Readout : official embedding = L2-normalised decoder output at the CLS position (StateEmbeddingModel.forward |
| `embedding`), D = 2048. `state emb transform` additionally concatenates a 10-d dataset-token vector |
| (dataset_embedder(decoder output at the dataset token)) -> 2058-d X_state; that batch-identity part is NOT |
| saved here (only the 2048-d cell CLS embedding). |
| Output : rna_sets.save("se600m", set, E) -> cache/rnafm/se600m_<set>.npy (n_spot, 2048) float16. |
| Sanity : finite, mean pairwise cosine of 200 random spots, (hest) organ agreement of cosine 10-NN vs base rate, |
| and (hest, 300 spots) cosine between two runs with different tie-breaking seeds (stochasticity of the |
| undetected-gene filler tokens). |
| """ |
| import math |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import yaml |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
|
|
| from common import CACHE |
| from rna_sets import SETS, load, save |
|
|
| MODEL = "se600m" |
| REPO, REV = "arcinstitute/SE-600M", "5a9a80f44f7ce32ce57059933ef0d735d7c10ce5" |
| dev = "cuda" |
| DT = torch.bfloat16 |
|
|
|
|
| |
| class SkipBlock(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| self.intermediate_dense = nn.Linear(d, 2 * d) |
| self.dense = nn.Linear(2 * d, d) |
| self.layer_norm = nn.LayerNorm(d) |
|
|
| def forward(self, x): |
| return self.layer_norm(x + self.dense(F.relu(self.intermediate_dense(x)))) |
|
|
|
|
| class Layer(nn.Module): |
| def __init__(self, d, nhead, d_hid): |
| super().__init__() |
| self.nhead = nhead |
| self.qkv_proj = nn.Linear(d, 3 * d) |
| self.out_proj = nn.Linear(d, d) |
| self.norm1, self.norm2 = nn.LayerNorm(d), nn.LayerNorm(d) |
| self.linear1, self.linear2 = nn.Linear(d, d_hid), nn.Linear(d_hid, d) |
|
|
| def forward(self, x): |
| B, T, D = x.shape |
| q, k, v = self.qkv_proj(x).chunk(3, dim=-1) |
| q, k, v = (t.view(B, T, self.nhead, D // self.nhead).transpose(1, 2) for t in (q, k, v)) |
| a = F.scaled_dot_product_attention(q, k, v, is_causal=False).transpose(1, 2).reshape(B, T, D) |
| x = self.norm1(x + self.out_proj(a)) |
| return self.norm2(x + self.linear2(F.gelu(self.linear1(x)))) |
|
|
|
|
| class TransformerEncoder(nn.Module): |
| def __init__(self, layers): |
| super().__init__() |
| self.layers = nn.ModuleList(layers) |
|
|
|
|
| class SE(nn.Module): |
| def __init__(self, token_dim, d, nhead, d_hid, nlayers, out_dim, n_pe, z_dim, n_ds): |
| super().__init__() |
| self.d = d |
| self.cls_token = nn.Parameter(torch.zeros(1, token_dim)) |
| self.dataset_token = nn.Parameter(torch.zeros(1, token_dim)) |
| self.encoder = nn.Sequential(nn.Linear(token_dim, d), nn.LayerNorm(d), nn.SiLU()) |
| self.gene_embedding_layer = self.encoder |
| self.transformer_encoder = TransformerEncoder([Layer(d, nhead, d_hid) for _ in range(nlayers)]) |
| self.decoder = nn.Sequential(SkipBlock(d), nn.Linear(d, out_dim)) |
| self.binary_decoder = nn.Sequential(SkipBlock(out_dim + d + z_dim), SkipBlock(out_dim + d + z_dim), |
| nn.Linear(out_dim + d + z_dim, 1)) |
| self.bin_encoder = nn.Embedding(10, d) |
| self.count_encoder = nn.Sequential(nn.Linear(1, 512), nn.LeakyReLU(), nn.Linear(512, 10)) |
| self.dataset_embedder = nn.Linear(out_dim, 10) |
| self.dataset_encoder = nn.Sequential(nn.Linear(out_dim, d), nn.SiLU(), nn.LayerNorm(d), nn.Dropout(0.1), |
| nn.Linear(d, n_ds)) |
| self.pe_embedding = nn.Embedding(n_pe, token_dim) |
|
|
| def forward(self, tok, cnt): |
| """tok: (B, L) global gene ids (position 0 = placeholder for CLS); cnt: (B, L) count values.""" |
| x = F.normalize(self.pe_embedding(tok), dim=2) |
| x[:, 0, :] = self.cls_token.expand(x.size(0), -1) |
| x = torch.cat([x, self.dataset_token.expand(x.size(0), -1).unsqueeze(1)], 1) |
| x = self.encoder(x) * math.sqrt(self.d) |
| w = F.softmax(self.count_encoder(cnt.unsqueeze(-1)), dim=-1) |
| ce = w @ self.bin_encoder.weight |
| x = x + torch.cat([ce, torch.zeros_like(ce[:, :1])], 1) |
| for lyr in self.transformer_encoder.layers: |
| x = lyr(x) |
| return F.normalize(self.decoder(x)[:, 0, :], dim=1) |
|
|
|
|
| def build(): |
| cfg = yaml.safe_load(open(hf_hub_download(REPO, "config.yaml", revision=REV))) |
| m, ds = cfg["model"], cfg["dataset"] |
| assert m["counts"] and m["rda"] and m["dataset_correction"] and ds["pad_length"] == 2048 |
| assert ds["cls_token_idx"] == 3 and cfg["loss"]["name"] == "tabular" |
| pe = torch.load(hf_hub_download(REPO, "protein_embeddings.pt", revision=REV), map_location="cpu", |
| weights_only=True) |
| sd = load_file(hf_hub_download(REPO, "model.safetensors", revision=REV)) |
| n_ds = cfg["dataset"][cfg["dataset"]["current"]]["num_datasets"] |
| model = SE(cfg["tokenizer"]["token_dim"], m["emsize"], m["nhead"], m["d_hid"], m["nlayers"], m["output_dim"], |
| len(pe), 1 + 10, n_ds) |
| for k in ["0.weight", "0.bias", "1.weight", "1.bias"]: |
| assert torch.equal(sd[f"encoder.{k}"], sd[f"gene_embedding_layer.{k}"]) |
| model.load_state_dict(sd, strict=True) |
| assert torch.equal(model.pe_embedding.weight, torch.vstack(list(pe.values()))) |
| del model.binary_decoder, model.dataset_encoder, model.gene_embedding_layer |
| model = model.eval().to(dev, DT) |
| print(f"[{MODEL}] loaded {len(sd)} tensors strictly; config pad_length={ds['pad_length']} " |
| f"{m['nlayers']}x{m['emsize']} heads={m['nhead']}", flush=True) |
| return model, cfg, {g: i for i, g in enumerate(pe)} |
|
|
|
|
| |
| def sentences(X, genes, g2i, pad_len, cls_idx, seed): |
| valid = np.array([g in g2i for g in genes]) |
| Xv = X[:, valid].tocsr() |
| gid = np.array([g2i[g] for g in genes[valid]], dtype=np.int64) |
| n = Xv.shape[1] |
| assert n >= pad_len - 1 |
| rng = np.random.default_rng(seed) |
| T = np.zeros((Xv.shape[0], pad_len), np.int64) |
| C = np.zeros((Xv.shape[0], pad_len), np.float32) |
| with np.errstate(over="ignore"): |
| n_heur = int(sum((r.data.max(initial=0) <= 35) and (np.expm1(r.data.astype(np.float64)).sum() <= 5e6) |
| for r in (X.getrow(i) for i in range(X.shape[0])))) |
| n_trunc = 0 |
| for r in range(Xv.shape[0]): |
| a, b = Xv.indptr[r], Xv.indptr[r + 1] |
| nz, raw = Xv.indices[a:b], Xv.data[a:b].astype(np.float64) |
| lv = np.log1p(raw) |
| s = lv.sum() |
| w = np.zeros(n) |
| if s > 0: |
| w[nz] = lv / s |
| else: |
| w[:] = 1.0 / n |
| p = rng.permutation(len(nz)) |
| order = nz[p][np.argsort(-lv[p], kind="stable")][:pad_len - 1] |
| n_trunc += len(nz) > pad_len - 1 |
| k = pad_len - 1 - len(order) |
| if k > 0: |
| zero = np.setdiff1d(np.arange(n), nz, assume_unique=True) |
| order = np.r_[order, rng.choice(zero, k, replace=False)] |
| loc = np.r_[cls_idx, order] |
| T[r] = gid[loc] |
| C[r] = 100.0 * w[loc] |
| stats = dict(valid=int(valid.sum()), n_gene=len(genes), heur_log=n_heur, trunc=n_trunc, |
| det_med=float(np.median(np.diff(Xv.indptr)))) |
| return T, C, stats |
|
|
|
|
| def embed(model, T, C, bs=16): |
| E = np.zeros((len(T), model.decoder[1].out_features), np.float32) |
| i = 0 |
| with torch.no_grad(), torch.autocast("cuda", dtype=DT): |
| while i < len(T): |
| try: |
| t = torch.from_numpy(T[i:i + bs]).to(dev) |
| c = torch.from_numpy(C[i:i + bs]).to(dev, DT) |
| E[i:i + bs] = model(t, c).float().cpu().numpy() |
| i += bs |
| except torch.OutOfMemoryError: |
| torch.cuda.empty_cache() |
| assert bs > 1, "OOM at batch size 1" |
| bs //= 2 |
| print(f" OOM -> batch size {bs}", flush=True) |
| return E |
|
|
|
|
| def sanity(E, name): |
| rng = np.random.default_rng(0) |
| S = E[rng.choice(len(E), min(200, len(E)), replace=False)] |
| S = S / np.linalg.norm(S, axis=1, keepdims=True) |
| mc = float((S @ S.T)[np.triu_indices(len(S), 1)].mean()) |
| msg = f"[{MODEL}] {name}: finite={np.isfinite(E).all()} std={E.std(0).mean():.4f} mean-pair-cos(200)={mc:.4f}" |
| if name == "hest": |
| org = pd.read_csv(os.path.join(CACHE, "hest", "meta.csv")).organ.values |
| assert len(org) == len(E) |
| Z = torch.from_numpy(E / np.linalg.norm(E, axis=1, keepdims=True)).to(dev) |
| nn_idx = [] |
| for k in range(0, len(Z), 2048): |
| s = Z[k:k + 2048] @ Z.T |
| s[torch.arange(s.shape[0]), torch.arange(k, k + s.shape[0])] = -2 |
| nn_idx.append(s.topk(10, dim=1).indices.cpu().numpy()) |
| nn_idx = np.concatenate(nn_idx) |
| knn = float((org[nn_idx] == org[:, None]).mean()) |
| p = pd.Series(org).value_counts(normalize=True).values |
| msg += f" 10-NN same-organ={knn:.3f} (base rate {float((p ** 2).sum()):.3f})" |
| print(msg, flush=True) |
| return mc |
|
|
|
|
| if __name__ == "__main__": |
| torch.cuda.set_per_process_memory_fraction(min(1.0, 8 * 2 ** 30 / torch.cuda.get_device_properties(0).total_memory)) |
| t0 = time.time() |
| model, cfg, g2i = build() |
| PAD, CLS = cfg["dataset"]["pad_length"], cfg["dataset"]["cls_token_idx"] |
| for name in (sys.argv[1:] or SETS): |
| t1 = time.time() |
| X, genes = load(name) |
| T, C, st = sentences(X, genes, g2i, PAD, CLS, seed=0) |
| print(f"[{MODEL}] {name}: {X.shape[0]} spots; gene coverage {st['valid']}/{st['n_gene']} " |
| f"({st['valid'] / st['n_gene']:.1%}) in protein_embeddings; detected valid genes/spot median " |
| f"{st['det_med']:.0f}; spots with >2047 detected (truncated) {st['trunc']}; spots the official " |
| f"raw-vs-log1p heuristic would treat as log1p {st['heur_log']}", flush=True) |
| E = embed(model, T, C) |
| if name == "hest": |
| idx = np.random.default_rng(1).choice(len(T), 300, replace=False) |
| T2, C2, _ = sentences(X[idx], genes, g2i, PAD, CLS, seed=1) |
| c = (embed(model, T2, C2) * E[idx]).sum(1) |
| print(f" seed-0 vs seed-1 tie-breaking: per-spot cosine median {np.median(c):.4f} min {c.min():.4f}", |
| flush=True) |
| sanity(E, name) |
| save(MODEL, name, E) |
| print(f" {name}: {time.time() - t1:.0f}s peak GPU {torch.cuda.max_memory_allocated() / 2 ** 30:.2f} GB", |
| flush=True) |
| print(f"[{MODEL}] done in {time.time() - t0:.0f}s", flush=True) |
|
|
