| """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): |
| |
| 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"): |
| |
| 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) |
|
|
