| from pathlib import Path |
| from types import SimpleNamespace |
| import re |
| import torch |
|
|
| from models import EquivariantThreeHopGINE |
| from smiles_to_npy_discretize import smiles_to_graph_with_labels |
|
|
|
|
| class SimpleGraph: |
| """The tiny subset of DGLGraph used by EquivariantThreeHopGINE.""" |
| def __init__(self, src, dst, weights, features): |
| self._src, self._dst = src, dst |
| self.edata = {"weight": weights} |
| self.ndata = {"feat": features} |
|
|
| def edges(self): |
| return self._src, self._dst |
|
|
| def num_edges(self): |
| return int(self._src.numel()) |
|
|
| def to(self, device): |
| self._src = self._src.to(device) |
| self._dst = self._dst.to(device) |
| self.edata = {k: v.to(device) for k, v in self.edata.items()} |
| self.ndata = {k: v.to(device) for k, v in self.ndata.items()} |
| return self |
|
|
|
|
| def build_args(): |
| return SimpleNamespace( |
| hidden_dim=16, |
| codebook_size=10000, |
| edge_emb_dim=32, |
| ema_decay=0.8, |
| dynamic_threshold=True, |
| epoch_at_mode_shift=0, |
| ss_max_total_latent_count=40000, |
| train_or_infer="infer", |
| use_checkpoint=True, |
| ) |
|
|
|
|
| def ensure_codebook_keys(model, state_dict): |
| cb = model.vq._codebook |
| prefixes = ["vq._codebook.", ""] |
| for prefix in prefixes: |
| ea_pat = re.compile(rf"^{re.escape(prefix)}embed_avg_(.+)$") |
| cs_pat = re.compile(rf"^{re.escape(prefix)}cluster_size_(.+)$") |
| emb_pat = re.compile(rf"^{re.escape(prefix)}embed\.(.+)$") |
| created = False |
| for key, value in state_dict.items(): |
| match = ea_pat.match(key) |
| if match and torch.is_tensor(value) and value.ndim == 2: |
| original = match.group(1) |
| k, d = value.shape |
| if not hasattr(cb, f"cluster_size_{original}"): |
| cb.register_buffer(f"cluster_size_{original}", torch.zeros(k)) |
| if not hasattr(cb, f"embed_avg_{original}"): |
| cb.register_buffer(f"embed_avg_{original}", torch.zeros(k, d)) |
| cb._get_or_create_safe_key(original, K_e=int(k), D=int(d), device="cpu") |
| created = True |
| for key, value in state_dict.items(): |
| match = cs_pat.match(key) |
| if match and torch.is_tensor(value) and value.ndim == 1: |
| original = match.group(1) |
| if not hasattr(cb, f"cluster_size_{original}"): |
| cb.register_buffer(f"cluster_size_{original}", torch.zeros_like(value, device="cpu")) |
| created = True |
| for key, value in state_dict.items(): |
| match = emb_pat.match(key) |
| if match and torch.is_tensor(value) and value.ndim == 2: |
| safe = match.group(1) |
| if safe not in cb.embed: |
| cb.embed[safe] = torch.nn.Parameter(torch.zeros_like(value, device="cpu")) |
| created = True |
| if created: |
| return |
|
|
|
|
| def make_graph(adj, features, device): |
| x = torch.as_tensor(features, dtype=torch.float32, device=device) |
| w1 = torch.as_tensor(adj, dtype=torch.float32, device=device) |
| n = x.shape[0] |
| a1 = w1 > 0 |
| a1 = a1 | a1.T |
| a1.fill_diagonal_(True) |
| w1 = w1.clone() |
| w1.fill_diagonal_(1.0) |
| a1f = a1.float() |
| a2 = (a1f @ a1f) > 0 |
| a3 = (a2.float() @ a1f) > 0 |
| two_only = a2 & ~a1 |
| three_only = a3 & ~(a1 | a2) |
| full_w = w1 + two_only.float() * 0.5 + three_only.float() * 0.3 |
| full_w.fill_diagonal_(1.0) |
| src, dst = (full_w > 0).nonzero(as_tuple=True) |
| return SimpleGraph(src, dst, full_w[src, dst], x), n |
|
|
|
|
| def make_masks(features, device): |
| masks = {} |
| for i, row in enumerate(features): |
| key = "_".join(str(int(row[j])) for j in (0, 2, 3, 4, 5, 6)) |
| masks.setdefault(key, []).append(i) |
| return {k: torch.tensor(v, dtype=torch.long, device=device) for k, v in masks.items()} |
|
|
|
|
| class VQAtomTokenizer: |
| def __init__(self, checkpoint="data/model_epoch_3.pt", device=None): |
| self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| args = build_args() |
| self.model = EquivariantThreeHopGINE( |
| in_feats=64, hidden_feats=16, out_feats=16, args=args |
| ) |
| ckpt = torch.load(Path(checkpoint), map_location="cpu", weights_only=False) |
| state = ckpt.get("model", ckpt.get("state_dict", ckpt)) if isinstance(ckpt, dict) else ckpt |
| drop = ( |
| "vq._codebook.usage_ema_k_", "vq._codebook.split_cd_k_", |
| "vq._codebook.ever_used_k_", "vq._codebook.last_used_ep_k_", |
| ) |
| state = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop)} |
| ensure_codebook_keys(self.model, state) |
| missing, unexpected = self.model.load_state_dict(state, strict=False) |
| essential_missing = [k for k in missing if not any(x in k for x in ("usage_ema", "split_cd", "ever_used", "last_used"))] |
| if essential_missing: |
| raise RuntimeError(f"Checkpoint is missing required weights: {essential_missing[:8]}") |
| self.model.to(self.device).eval() |
|
|
| @torch.inference_mode() |
| def encode(self, smiles): |
| adj, features, _ = smiles_to_graph_with_labels(smiles, 0) |
| if features.shape[1] != 79: |
| raise RuntimeError(f"Expected 79 atom features, got {features.shape[1]}") |
| if features.shape[0] >= 100: |
| raise ValueError("This demo supports molecules with fewer than 100 heavy atoms.") |
| graph, _ = make_graph(adj, features, self.device) |
| masks = make_masks(features, self.device) |
| out = self.model( |
| graph, graph.ndata["feat"], 0, masks, None, 0, |
| None, "infer", [graph.ndata["feat"]], |
| ) |
| _, ids, _ = out |
| key_ids, cluster_ids, global_ids, id2safe = ids |
| return { |
| "tokens": global_ids.detach().cpu().long().tolist(), |
| "key_ids": key_ids.detach().cpu().long().tolist(), |
| "cluster_ids": cluster_ids.detach().cpu().long().tolist(), |
| "id2safe": id2safe, |
| } |
|
|
