File size: 12,517 Bytes
40c9a4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 | """Forward EI predictors. All public prediction outputs are RAW relative intensity."""
from __future__ import annotations
import warnings
from pathlib import Path
from typing import Optional
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from rdkit import Chem, DataStructs
from rdkit.Chem import Descriptors, rdFingerprintGenerator
def choose_device(value="auto"):
if value == "auto":
value = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(value)
if device.type == "cuda" and not torch.cuda.is_available():
raise RuntimeError("CUDA requested but unavailable. Check the AutoDL PyTorch image / GPU.")
return device
def molecular_features(smiles, n_bits=2048, radius=2):
generator = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=n_bits)
count_rows, bit_rows, mass_rows = [], [], []
for value in smiles:
molecule = Chem.MolFromSmiles(str(value))
if molecule is None:
raise ValueError(f"Invalid SMILES: {value!r}")
counts = np.zeros(n_bits, dtype=np.float32)
bits = np.zeros(n_bits, dtype=np.float32)
DataStructs.ConvertToNumpyArray(generator.GetCountFingerprint(molecule), counts)
DataStructs.ConvertToNumpyArray(generator.GetFingerprint(molecule), bits)
count_rows.append(counts)
bit_rows.append(bits)
mass_rows.append(Descriptors.ExactMolWt(molecule))
if not count_rows:
return np.empty((0, n_bits), np.float32), np.empty((0, n_bits), np.float32), np.empty(0, np.float32)
return np.stack(count_rows), np.stack(bit_rows), np.asarray(mass_rows, dtype=np.float32)
class ResidualBlock(nn.Module):
def __init__(self, width, dropout):
super().__init__()
self.block = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, width), nn.GELU(),
nn.Dropout(dropout), nn.Linear(width, width))
def forward(self, inputs):
return inputs + self.block(inputs)
class BidirectionalSpectrumHead(nn.Module):
"""NEIMS-inspired mass/neutral-loss head; output represents sqrt(intensity)."""
def __init__(self, input_dim, output_dim, hidden_dim=512, blocks=2, dropout=0.1,
isotope_margin=4):
super().__init__()
self.output_dim = int(output_dim)
self.isotope_margin = int(isotope_margin)
self.body = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.GELU(),
*[ResidualBlock(hidden_dim, dropout) for _ in range(blocks)],
nn.LayerNorm(hidden_dim))
self.forward_head = nn.Linear(hidden_dim, output_dim)
self.reverse_head = nn.Linear(hidden_dim, output_dim)
self.gate = nn.Linear(hidden_dim, output_dim)
self.register_buffer("bins", torch.arange(output_dim), persistent=False)
def forward(self, embeddings, exact_mass):
hidden = self.body(embeddings)
forward = F.softplus(self.forward_head(hidden))
reverse = F.softplus(self.reverse_head(hidden))
center = torch.floor(exact_mass.float() + 0.5).long()
reverse_index = center[:, None] + self.isotope_margin - self.bins[None, :]
reverse_valid = (reverse_index >= 0) & (reverse_index < self.output_dim)
aligned = reverse.gather(1, reverse_index.clamp(0, self.output_dim - 1)) * reverse_valid
gate = torch.sigmoid(self.gate(hidden))
prediction = gate * forward + (1.0 - gate) * aligned
valid = (self.bins[None, :] > 0) & (self.bins[None, :] <= center[:, None] + self.isotope_margin)
return prediction * valid
class SafeBatchNorm(nn.BatchNorm1d):
def forward(self, values):
if self.training and values.shape[0] == 1:
return F.batch_norm(values, self.running_mean, self.running_var, self.weight, self.bias,
training=False, momentum=self.momentum, eps=self.eps)
return super().forward(values)
class NEIMSResidualBlock(nn.Module):
def __init__(self, width, bottleneck_factor, dropout):
super().__init__()
inner = max(1, int(width * bottleneck_factor))
self.block = nn.Sequential(SafeBatchNorm(width), nn.ReLU(), nn.Dropout(dropout),
nn.Linear(width, inner), SafeBatchNorm(inner), nn.ReLU(),
nn.Linear(inner, width))
def forward(self, values):
return values + self.block(values)
class NEIMSReimplementation(nn.Module):
"""Independent implementation aligned to the published NEIMS-style MLP defaults.
4096 raw Morgan count features, width 2000, seven BN/ReLU residual bottlenecks,
dropout .25, gated forward/reverse outputs. Shared split, sqrt-cosine loss and
0..750 bins differ from the original TensorFlow/NIST training protocol.
"""
def __init__(self, config):
super().__init__()
self.config = config
width = config.get("hidden_dim", 2000)
dimension = config["spectrum_dim"]
self.isotope_margin = config.get("isotope_margin", 5)
self.input = nn.Sequential(SafeBatchNorm(config.get("fingerprint_bits", 4096)),
nn.Linear(config.get("fingerprint_bits", 4096), width), nn.ReLU())
self.blocks = nn.Sequential(*[NEIMSResidualBlock(width, config.get("bottleneck_factor", .5),
config.get("dropout", .25))
for _ in range(config.get("residual_blocks", 7))],
SafeBatchNorm(width), nn.ReLU())
self.forward_head = nn.Linear(width, dimension)
self.reverse_head = nn.Linear(width, dimension)
self.gate = nn.Linear(width, dimension)
self.register_buffer("bins", torch.arange(dimension), persistent=False)
def forward(self, inputs, exact_mass):
features = self.blocks(self.input(inputs["fingerprints"]))
direct = self.forward_head(features)
reverse = self.reverse_head(features)
ceiling = torch.floor(exact_mass.float() + .5).long() + self.isotope_margin
source = ceiling[:, None] - self.bins[None, :]
valid = (source >= 0) & (source < len(self.bins)) & (self.bins[None, :] > 0)
aligned = reverse.gather(1, source.clamp(0, len(self.bins) - 1)) * valid
gate = torch.sigmoid(self.gate(features))
direct_valid = (self.bins[None, :] > 0) & (self.bins[None, :] <= ceiling[:, None])
return F.relu(gate * direct + (1 - gate) * aligned) * direct_valid
class MolformerSpectrumModel(nn.Module):
def __init__(self, config, model_dir_override=None):
super().__init__()
from transformers import AutoModel, AutoTokenizer
from peft import LoraConfig, get_peft_model
self.config = config
model_dir = Path(model_dir_override or config.get("model_dir", "models/molformer"))
if not model_dir.is_dir() or not (model_dir / "config.json").is_file():
raise FileNotFoundError(f"MoLFormer files missing in {model_dir}. Run download_resources.py --molformer first.")
self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True,
local_files_only=True)
backbone = AutoModel.from_pretrained(str(model_dir), trust_remote_code=True,
local_files_only=True, deterministic_eval=True)
expected = config.get("lora_targets", ["query", "value"])
observed = {name.rsplit(".", 1)[-1] for name, module in backbone.named_modules()
if isinstance(module, nn.Linear)}
if not set(expected).issubset(observed):
raise RuntimeError(f"LoRA modules {expected} not found. Available linear leaf names: {sorted(observed)}")
lora = LoraConfig(r=config.get("lora_r", 8), lora_alpha=config.get("lora_alpha", 16),
lora_dropout=config.get("lora_dropout", 0.05), target_modules=expected,
bias="none")
self.encoder = get_peft_model(backbone, lora)
self.head = BidirectionalSpectrumHead(
backbone.config.hidden_size, config["spectrum_dim"], config.get("hidden_dim", 512),
config.get("residual_blocks", 2), config.get("dropout", 0.1),
config.get("isotope_margin", 4))
def tokenize(self, smiles):
# Check rather than silently truncate molecules.
encoded = self.tokenizer(list(smiles), padding=True, truncation=False, return_tensors="pt")
limit = self.config.get("max_tokens", 256)
if encoded["input_ids"].shape[1] > limit:
raise ValueError(f"SMILES exceeds max_tokens={limit}; review long structures or raise max_tokens explicitly.")
return {key: value for key, value in encoded.items() if key in {"input_ids", "attention_mask"}}
def forward(self, inputs, exact_mass):
encoded = self.encoder(**inputs, return_dict=True)
mask = inputs["attention_mask"].unsqueeze(-1).to(encoded.last_hidden_state.dtype)
pooled = (encoded.last_hidden_state * mask).sum(1) / mask.sum(1).clamp_min(1)
return self.head(pooled, exact_mass)
def build_model(config, model_dir_override=None):
if config["model_type"] == "neims_reimpl":
return NEIMSReimplementation(config)
if config["model_type"] == "molformer":
return MolformerSpectrumModel(config, model_dir_override)
raise ValueError(f"Unknown model_type: {config['model_type']}")
def load_checkpoint(path, map_location="cpu"):
# Only load checkpoints produced by this package or another trusted source.
return torch.load(path, map_location=map_location, weights_only=False)
def raw_intensity(sqrt_prediction):
raw = sqrt_prediction.float().clamp_min(0).square()
return raw / raw.amax(dim=-1, keepdim=True).clamp_min(1e-12)
class SpectrumPredictor:
def __init__(self, checkpoint, device="auto", model_dir_override=None):
self.device = choose_device(device)
saved = load_checkpoint(checkpoint)
self.config = saved["config"]
self.model = build_model(self.config, model_dir_override)
self.model.load_state_dict(saved["model_state"], strict=True)
self.model.to(self.device).eval()
if "train_keys" not in saved:
raise ValueError("Checkpoint lacks training identity audit (train_keys)")
self.train_keys = set(saved["train_keys"])
self.metadata = {key: value for key, value in saved.items()
if key not in {"model_state", "optimizer_state", "rng_state", "train_keys"}}
@torch.inference_mode()
def __call__(self, smiles, exact_mass=None, batch_size=32):
smiles = list(smiles)
if batch_size < 1:
raise ValueError("batch_size must be positive")
if not smiles:
return np.empty((0, self.config["spectrum_dim"]), dtype=np.float32)
supplied_mass = None if exact_mass is None else np.asarray(exact_mass, dtype=np.float32)
if supplied_mass is not None and supplied_mass.shape != (len(smiles),):
raise ValueError("exact_mass must have one value per SMILES")
results = []
for start in range(0, len(smiles), batch_size):
batch_smiles = smiles[start:start + batch_size]
counts, _, calculated = molecular_features(batch_smiles, self.config.get("fingerprint_bits", 2048),
self.config.get("fingerprint_radius", 2))
mass = calculated if supplied_mass is None else supplied_mass[start:start + batch_size]
if not np.isfinite(mass).all() or (mass <= 0).any():
raise ValueError("Non-finite or nonpositive exact_mass")
if self.config["model_type"] == "neims_reimpl":
inputs = {"fingerprints": torch.from_numpy(counts)}
else:
inputs = self.model.tokenize(batch_smiles)
inputs = {key: value.to(self.device) for key, value in inputs.items()}
values = self.model(inputs, torch.as_tensor(mass, device=self.device))
results.append(raw_intensity(values).cpu().numpy())
return np.concatenate(results)
def load_predictor(checkpoint, device="auto", model_dir_override=None):
return SpectrumPredictor(checkpoint, device, model_dir_override)
|